COFF: Allow mangled symbols as arguments for /export.

Usually dllexported symbols are defined with 'extern "C"',
so identifying them is easy. We can just do hash table lookup
to look up exported symbols.

However, C++ non-member functions are also allowed to be exported,
and they can be specified with unmangled name. So, if /export:foo
is given, we need to look up not only "foo" but also its all
mangled names. In MSVC mangling scheme, that means that we need to
look up any symbol which starts with "?foo@@Y".

In this patch, we scan the entire symbol table to search for
a mangled symbol. The symbol table is a DenseMap, and that doesn't
support table lookup by string prefix. This is of course very
inefficient. But that should be probably OK because the user
should always add 'extern "C"' to dllexported symbols.

llvm-svn: 240919

5 files changed, 49 insertions, 8 deletions

diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index 67f6034b0fa4..c8554e1435aa 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -22,13 +22,13 @@ namespace coff {
 
 using llvm::COFF::WindowsSubsystem;
 using llvm::StringRef;
-class Defined;
+struct Symbol;
 
 // Represents an /export option.
 struct Export {
   StringRef Name;
   StringRef ExtName;
-  Defined *Sym = nullptr;
+  Symbol *Sym = nullptr;
   uint16_t Ordinal = 0;
   bool Noname = false;
   bool Data = false;
diff --git a/lld/COFF/DLL.cpp b/lld/COFF/DLL.cpp
index d44dd4b4a0b6..5dda97288dc7 100644
--- a/lld/COFF/DLL.cpp
+++ b/lld/COFF/DLL.cpp
@@ -419,8 +419,10 @@ public:
   size_t getSize() const override { return Size * 4; }
 
   void writeTo(uint8_t *Buf) override {
-    for (Export &E : Config->Exports)
-      write32le(Buf + FileOff + E.Ordinal * 4, E.Sym->getRVA());
+    for (Export &E : Config->Exports) {
+      auto *D = cast<Defined>(E.Sym->Body);
+      write32le(Buf + FileOff + E.Ordinal * 4, D->getRVA());
+    }
   }
 
 private:
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 649ddfb39460..12cff8a4f9af 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -566,6 +566,20 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
       break;
   }
 
+  // Windows specific -- resolve dllexported symbols.
+  for (Export &E : Config->Exports) {
+    StringRef Name;
+    Symbol *Sym;
+    std::tie(Name, Sym) = Symtab.findMangled(E.Name);
+    if (!Sym) {
+      llvm::errs() << "exported symbol is not defined: " << E.Name << "\n";
+      return false;
+    }
+    if (E.Name != Name)
+      Symtab.rename(E.Name, Name);
+    E.Sym = Sym;
+  }
+
   // Make sure we have resolved all symbols.
   if (Symtab.reportRemainingUndefines())
     return false;
@@ -593,12 +607,9 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     writeImportLibrary();
 
   // Windows specific -- fix up dllexported symbols.
-  if (!Config->Exports.empty()) {
-    for (Export &E : Config->Exports)
-      E.Sym = Symtab.find(E.Name);
+  if (!Config->Exports.empty())
     if (fixupExports())
       return false;
-  }
 
   // Windows specific -- Create a side-by-side manifest file.
   if (Config->Manifest == Configuration::SideBySide)
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
index b823ce9eacb5..831ab7d8a7ee 100644
--- a/lld/COFF/SymbolTable.cpp
+++ b/lld/COFF/SymbolTable.cpp
@@ -15,6 +15,7 @@
 #include "llvm/LTO/LTOCodeGenerator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include <utility>
 
 using namespace llvm;
 
@@ -173,6 +174,28 @@ Defined *SymbolTable::find(StringRef Name) {
   return nullptr;
 }
 
+// Find a given symbol or its mangled symbol.
+std::pair<StringRef, Symbol *> SymbolTable::findMangled(StringRef S) {
+  auto It = Symtab.find(S);
+  if (It != Symtab.end()) {
+    Symbol *Sym = It->second;
+    if (isa<Defined>(Sym->Body))
+      return std::make_pair(S, Sym);
+  }
+
+  // In Microsoft ABI, a non-member function name is mangled this way.
+  std::string Prefix = ("?" + S + "@@Y").str();
+  for (auto I : Symtab) {
+    StringRef Name = I.first;
+    Symbol *Sym = I.second;
+    if (!Name.startswith(Prefix))
+      continue;
+    if (isa<Defined>(Sym->Body))
+      return std::make_pair(Name, Sym);
+  }
+  return std::make_pair(S, nullptr);
+}
+
 std::error_code SymbolTable::resolveLazy(StringRef Name) {
   auto It = Symtab.find(Name);
   if (It == Symtab.end())
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
index a1cf3408271f..62f14b09f525 100644
--- a/lld/COFF/SymbolTable.h
+++ b/lld/COFF/SymbolTable.h
@@ -52,6 +52,11 @@ public:
   // different symbol). Returns a nullptr if not found.
   Defined *find(StringRef Name);
 
+  // Find a symbol assuming that Name is a function name.
+  // Not only a given string but its mangled names (in MSVC C++ manner)
+  // will be searched.
+  std::pair<StringRef, Symbol *> findMangled(StringRef Name);
+
   // Windows specific -- `main` is not the only main function in Windows.
   // You can choose one from these four -- {w,}{WinMain,main}.
   // There are four different entry point functions for them,