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

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;

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:

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)

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())

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,

lld/test/COFF/Inputs/export.yaml

@@ -48,4 +48,10 @@ symbols:
     SimpleType:      IMAGE_SYM_TYPE_NULL
     ComplexType:     IMAGE_SYM_DTYPE_NULL
     StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            '?mangled@@YAHXZ'
+    Value:           16
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
 ...

lld/test/COFF/dll.test

@@ -1,7 +1,8 @@
 # REQUIRES: winres
 
 # RUN: yaml2obj < %p/Inputs/export.yaml > %t.obj
-# RUN: lld -flavor link2 /out:%t.dll /dll %t.obj /export:exportfn1 /export:exportfn2
+# RUN: lld -flavor link2 /out:%t.dll /dll %t.obj /export:exportfn1 /export:exportfn2 \
+# RUN:   /export:mangled
 # RUN: llvm-objdump -p %t.dll | FileCheck -check-prefix=EXPORT %s
 
 EXPORT:      Export Table:
@@ -11,6 +12,7 @@ EXPORT-NEXT:       0        0
 EXPORT-NEXT:       1   0x1008  exportfn1
 EXPORT-NEXT:       2   0x1010  exportfn2
 EXPORT-NEXT:       3   0x1010  exportfn3
+EXPORT-NEXT:       4   0x1010  mangled
 
 # RUN: llvm-as -o %t.lto.obj %p/Inputs/export.ll
 # RUN: lld -flavor link2 /out:%t.lto.dll /dll %t.lto.obj /export:exportfn1 /export:exportfn2
@@ -28,7 +30,8 @@ EXPORT-LTO-NEXT:       3   0x1030  exportfn3
 # RUN: lld -flavor link2 /out:%t2.exe %t2.obj %t.lib
 # RUN: llvm-readobj -coff-imports %t2.exe | FileCheck -check-prefix=IMPORT %s
 
-# RUN: lld -flavor link2 /out:%t.dll /dll %t.obj /implib:%t2.lib /export:exportfn1 /export:exportfn2
+# RUN: lld -flavor link2 /out:%t.dll /dll %t.obj /implib:%t2.lib \
+# RUN:   /export:exportfn1 /export:exportfn2
 # RUN: lld -flavor link2 /out:%t2.exe %t2.obj %t2.lib
 # RUN: llvm-readobj -coff-imports %t2.exe | FileCheck -check-prefix=IMPORT %s