[PECOFF] Improve /export compatibility.

Looks like the rule of /export is more complicated than
I was thinking. If /export:foo, for example, is given, and
if the actual symbol name in an object file is _foo@<number>,
we need to export that symbol as foo, not as the mangled name.

If only /export:_foo@<number> is given, the symbol is exported
as _foo@<number>.

If both /export:foo and /export:_foo@<number> are given,
they are considered as duplicates, and the linker needs to
choose the unmangled name.

The basic idea seems that the linker needs to export a symbol
with the same name as given as /export.

We exported mangled symbols. This patch fixes that issue.

llvm-svn: 223341

4 files changed, 63 insertions, 42 deletions

diff --git a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
index 6d306e209aa..a59e5ae2424 100644
--- a/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
+++ b/lld/include/lld/ReaderWriter/PECOFFLinkingContext.h
@@ -65,14 +65,17 @@ public:
       return getExternalName().compare(other.getExternalName()) < 0;
     }
 
+    StringRef getRealName() const {
+      return mangledName.empty() ? name : mangledName;
+    }
+
     StringRef getExternalName() const {
-      if (!externalName.empty())
-        return externalName;
-      return name;
+      return externalName.empty() ? name : externalName;
     }
 
     std::string name;
     std::string externalName;
+    std::string mangledName;
     int ordinal;
     bool noname;
     bool isData;
diff --git a/lld/lib/ReaderWriter/PECOFF/EdataPass.cpp b/lld/lib/ReaderWriter/PECOFF/EdataPass.cpp
index 4edf1124532..3944cf7c8c0 100644
--- a/lld/lib/ReaderWriter/PECOFF/EdataPass.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/EdataPass.cpp
@@ -26,20 +26,54 @@ using llvm::object::export_directory_table_entry;
 namespace lld {
 namespace pecoff {
 
+typedef PECOFFLinkingContext::ExportDesc ExportDesc;
+
+// dedupExports removes duplicate export entries. If two exports are
+// referring the same symbol, they are considered duplicates.
+// This could happen if the same symbol name is specified as an argument
+// to /export more than once, or an unmangled and mangled name of the
+// same symbol are given to /export. In the latter case, we choose
+// unmangled (shorter) name.
+static void dedupExports(PECOFFLinkingContext &ctx) {
+  std::vector<ExportDesc> &exports = ctx.getDllExports();
+  // Pass 1: find duplicate entries
+  std::set<const ExportDesc *> dup;
+  std::map<StringRef, ExportDesc *> map;
+  for (ExportDesc &exp : exports) {
+    if (!exp.externalName.empty())
+      continue;;
+    StringRef symbol = exp.getRealName();
+    auto it = map.find(symbol);
+    if (it == map.end()) {
+      map[symbol] = &exp;
+    } else if (symbol.size() < it->second->getRealName().size()) {
+      map[symbol] = &exp;
+      dup.insert(it->second);
+    } else {
+      dup.insert(&exp);
+    }
+  }
+  // Pass 2: remove duplicate entries
+  auto pred = [&](const ExportDesc &exp) {
+    return dup.count(&exp) == 1;
+  };
+  exports.erase(std::remove_if(exports.begin(), exports.end(), pred),
+                exports.end());
+}
+
 static void assignOrdinals(PECOFFLinkingContext &ctx) {
-  std::vector<PECOFFLinkingContext::ExportDesc> &exports = ctx.getDllExports();
+  std::vector<ExportDesc> &exports = ctx.getDllExports();
   int maxOrdinal = -1;
-  for (PECOFFLinkingContext::ExportDesc &desc : exports)
+  for (ExportDesc &desc : exports)
     maxOrdinal = std::max(maxOrdinal, desc.ordinal);
 
   std::sort(exports.begin(), exports.end(),
-            [](const PECOFFLinkingContext::ExportDesc &a,
-               const PECOFFLinkingContext::ExportDesc &b) {
+            [](const ExportDesc &a, const ExportDesc &b) {
     return a.getExternalName().compare(b.getExternalName()) < 0;
   });
 
   int nextOrdinal = (maxOrdinal == -1) ? 1 : (maxOrdinal + 1);
-  for (PECOFFLinkingContext::ExportDesc &desc : exports)
+  for (ExportDesc &desc : exports)
     if (desc.ordinal == -1)
       desc.ordinal = nextOrdinal++;
 }
@@ -51,7 +85,7 @@ static bool getExportedAtoms(PECOFFLinkingContext &ctx, MutableFile *file,
     definedAtoms[atom->name()] = atom;
 
   for (PECOFFLinkingContext::ExportDesc &desc : ctx.getDllExports()) {
-    auto it = definedAtoms.find(desc.name);
+    auto it = definedAtoms.find(desc.getRealName());
     if (it == definedAtoms.end()) {
       llvm::errs() << "Symbol <" << desc.name
                    << "> is exported but not defined.\n";
@@ -142,6 +176,7 @@ EdataPass::createOrdinalTable(const std::vector<TableEntry> &entries,
 }
 
 void EdataPass::perform(std::unique_ptr<MutableFile> &file) {
+  dedupExports(_ctx);
   assignOrdinals(_ctx);
 
   std::vector<TableEntry> entries;
diff --git a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
index 4feaec6c753..be1d8434dd7 100644
--- a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
+++ b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
@@ -261,29 +261,11 @@ public:
     if (!findDecoratedSymbol(_ctx, _syms.get(), sym.str(), replace))
       return nullptr;
 
-    // We found a decorated symbol. There may be another symbol that
-    // has the same decorated name. If that's the case, we remove the
-    // duplicate item.
-    std::vector<ExportDesc> &exp = _ctx->getDllExports();
-    auto isFound = std::find_if(
-        exp.begin(), exp.end(),
-        [&](ExportDesc &e) { return e.getExternalName().equals(replace); });
-    if (isFound != exp.end()) {
-      exp.erase(
-          std::remove_if(exp.begin(), exp.end(),
-                         [&](ExportDesc &e) { return e.name == sym; }),
-          exp.end());
-    } else {
-      for (ExportDesc &e : exp) {
-        if (e.name == sym) {
-          e.name = replace;
-          break;
-        }
-      }
-      if (_ctx->deadStrip())
-        _ctx->addDeadStripRoot(_ctx->allocate(replace));
-    }
-
+    for (ExportDesc &exp : _ctx->getDllExports())
+      if (exp.name == sym)
+        exp.mangledName = replace;
+    if (_ctx->deadStrip())
+      _ctx->addDeadStripRoot(_ctx->allocate(replace));
     return new (_alloc) impl::SymbolRenameFile(sym, replace);
   }
 
diff --git a/lld/test/pecoff/export.test b/lld/test/pecoff/export.test
index 38ad8d56040..63b8677cd4b 100644
--- a/lld/test/pecoff/export.test
+++ b/lld/test/pecoff/export.test
@@ -59,30 +59,31 @@ CHECK5-NEXT:        2   0x2010  exportfn7
 CHECK6:      Export Table:
 CHECK6:      DLL name: export.test.tmp6.dll
 CHECK6:       Ordinal      RVA  Name
-CHECK6-NEXT:        1   0x2010  ?exportfn8@@YAXXZ
+CHECK6-NEXT:        1   0x2010  exportfn3@256
+CHECK6-NEXT:        2   0x2010  exportfn8
 
-# RUN: lld -flavor link /out:%t6.dll /dll /entry:init \
+# RUN: lld -flavor link /out:%t7.dll /dll /entry:init \
 # RUN:   /export:exportfn7 /export:exportfn7@8 \
 # RUN:   /export:exportfn8 /export:exportfn8 /export:exportfn3 -- %t.obj
-# RUN: llvm-objdump -p %t6.dll | FileCheck -check-prefix=DUP %s
+# RUN: llvm-objdump -p %t7.dll | FileCheck -check-prefix=DUP %s
 
 DUP:      Export Table:
-DUP:      DLL name: export.test.tmp6.dll
+DUP:      DLL name: export.test.tmp7.dll
 DUP:       Ordinal      RVA  Name
-DUP-NEXT:        1   0x2010  ?exportfn8@@YAXXZ
-DUP-NEXT:        2   0x2010  exportfn3@256
-DUP-NEXT:        3   0x2010  exportfn7@8
+DUP-NEXT:        1   0x2010  exportfn3
+DUP-NEXT:        2   0x2010  exportfn7
+DUP-NEXT:        3   0x2010  exportfn8
 DUP-NOT:  ?exportfn8@@YAXXZ
 DUP-NOT:  exportfn3@256
 
 # RUN: yaml2obj %p/Inputs/export.obj.yaml > %t.obj
 #
-# RUN: lld -flavor link /out:%t1.dll /dll /entry:init \
+# RUN: lld -flavor link /out:%t8.dll /dll /entry:init \
 # RUN:   /export:f1=exportfn1 /export:f2@4=exportfn2,private -- %t.obj
-# RUN: llvm-objdump -p %t1.dll | FileCheck -check-prefix=EQUAL %s
+# RUN: llvm-objdump -p %t8.dll | FileCheck -check-prefix=EQUAL %s
 
 EQUAL:      Export Table:
-EQUAL:      DLL name: export.test.tmp1.dll
+EQUAL:      DLL name: export.test.tmp8.dll
 EQUAL:       Ordinal      RVA  Name
 EQUAL-NEXT:       1   0x2010  exportfn3@256
 EQUAL-NEXT:       2   0x2008  f1