[PECOFF] Find symbols with @number suffix for dllexported symbols

As written in the comment in this patch, symbol names specified with
/export option is resolved in a special way; for /export:foo, linker
finds a foo@<number> symbol if such symbols exists.

On Windows, a function in stdcall calling convention is mangled with
a leading underscore and following "@" and numbers. This name
mangling is kind of automatic, so you can sometimes omit _ and @number
when specifying a symbol. /export option is that case.

Previously, if a file in an archive file foo.lib provides a symbol
_fn@8, and /export:fn is specified, LLD failed to resolve the symbol.
It only tried to find _fn, and failed to find _fn@8. With this patch,
_fn@8 will be searched on the second iteration.

Differential Revision: http://reviews.llvm.org/D3736

llvm-svn: 208754

4 files changed, 148 insertions, 1 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
index 5ff8bc647bd..a6f69392b98 100644
--- a/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
+++ b/lld/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h
@@ -14,6 +14,8 @@
 #include "lld/ReaderWriter/Simple.h"
 #include "llvm/Support/Allocator.h"
 
+#include <mutex>
+
 namespace lld {
 namespace pecoff {
 
@@ -86,6 +88,22 @@ private:
   atom_collection_vector<AbsoluteAtom> _absoluteAtoms;
 };
 
+// A file to make Resolver to resolve a symbol TO instead of a symbol FROM,
+// using fallback mechanism for an undefined symbol. One can virtually rename an
+// undefined symbol using this file.
+class SymbolRenameFile : public SimpleFile {
+public:
+  SymbolRenameFile(StringRef from, StringRef to)
+      : SimpleFile("<symbol-rename>"), _to(*this, to),
+        _from(*this, from, &_to) {
+    addAtom(_from);
+  };
+
+private:
+  COFFUndefinedAtom _to;
+  COFFUndefinedAtom _from;
+};
+
 } // anonymous namespace
 
 // A virtual file containing absolute symbol __ImageBase. __ImageBase (or
@@ -141,5 +159,94 @@ private:
   mutable llvm::BumpPtrAllocator _alloc;
 };
 
+// A ExportedSymbolRenameFile is a virtual archive file for dllexported symbols.
+//
+// One usually has to specify the exact symbol name to resolve it. That's true
+// in most cases for PE/COFF, except the one described below.
+//
+// DLLExported symbols can be specified using a module definition file. In a
+// file, one can write an EXPORT directive followed by symbol names. Such
+// symbols may not be fully decorated -- one can omit "@" and the following
+// number suffix for the stdcall function.
+//
+// If a symbol FOO is specified to be dllexported by a module definition file,
+// linker has to search not only for FOO but also for FOO@[0-9]+. This ambiguous
+// matching semantics does not fit well with Resolver.
+//
+// We could probably modify Resolver to resolve ambiguous symbols, but I think
+// we don't want to do that because it'd be rarely used, and only this Windows
+// specific feature would use it. It's probably not a good idea to make the core
+// linker to be able to deal with it.
+//
+// So, instead of tweaking Resolver, we chose to do some hack here. An
+// ExportedSymbolRenameFile maintains a set containing all possibly defined
+// symbol names. That set would be a union of (1) all the defined symbols that
+// are already parsed and read and (2) all the defined symbols in archive files
+// that are not yet be parsed.
+//
+// If Resolver asks this file to return an atom for a dllexported symbol, find()
+// looks up the set, doing ambiguous matching. If there's a symbol with @
+// prefix, it returns an atom to rename the dllexported symbol, hoping that
+// Resolver will find the new symbol with atsign from an archive file at the
+// next visit.
+class ExportedSymbolRenameFile : public VirtualArchiveLibraryFile {
+public:
+  ExportedSymbolRenameFile(const PECOFFLinkingContext &ctx)
+      : VirtualArchiveLibraryFile("<export>") {
+    for (const PECOFFLinkingContext::ExportDesc &desc : ctx.getDllExports())
+      _exportedSyms.insert(desc.name);
+  }
+
+  void addResolvableSymbols(File *file) {
+    std::lock_guard<std::mutex> lock(_mutex);
+    if (_seen.count(file) > 0)
+      return;
+    _seen.insert(file);
+    if (auto *archive = dyn_cast<ArchiveLibraryFile>(file)) {
+      for (const std::string &sym : archive->getDefinedSymbols())
+        _defined.insert(sym);
+      return;
+    }
+    for (const DefinedAtom *atom : file->defined())
+      if (!atom->name().empty())
+        _defined.insert(atom->name());
+  }
+
+  const File *find(StringRef sym, bool dataSymbolOnly) const override {
+    if (_exportedSyms.count(sym) == 0)
+      return nullptr;
+    std::string replace;
+    if (!findSymbolWithAtsignSuffix(sym.str(), replace))
+      return nullptr;
+    return new (_alloc) SymbolRenameFile(sym, replace);
+  }
+
+private:
+  // Find a symbol that starts with a given symbol name followed
+  // by @number suffix.
+  bool findSymbolWithAtsignSuffix(std::string sym, std::string &res) const {
+    sym.append("@");
+    auto it = _defined.lower_bound(sym);
+    for (auto e = _defined.end(); it != e; ++it) {
+      if (!StringRef(*it).startswith(sym))
+        return false;
+      if (it->size() == sym.size())
+        continue;
+      StringRef suffix = it->substr(sym.size());
+      if (suffix.find_first_not_of("0123456789") != StringRef::npos)
+        continue;
+      res = *it;
+      return true;
+    }
+    return false;
+  }
+
+  std::set<std::string> _exportedSyms;
+  std::set<std::string> _defined;
+  std::set<File *> _seen;
+  mutable std::mutex _mutex;
+  mutable llvm::BumpPtrAllocator _alloc;
+};
+
 } // end namespace pecoff
 } // end namespace lld
diff --git a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
index 099b472446a..a327192f5b2 100644
--- a/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp
@@ -101,8 +101,23 @@ std::unique_ptr<File> PECOFFLinkingContext::createUndefinedSymbolFile() const {
       "<command line option /include>");
 }
 
+namespace {
+// As per policy, we cannot use std::function.
+class ObserverCallback {
+public:
+  explicit ObserverCallback(pecoff::ExportedSymbolRenameFile *f)
+      : _renameFile(f) {}
+
+  void operator()(File *file) { _renameFile->addResolvableSymbols(file); }
+
+private:
+  pecoff::ExportedSymbolRenameFile *_renameFile;
+};
+} // end anonymous namespace
+
 bool PECOFFLinkingContext::createImplicitFiles(
-    std::vector<std::unique_ptr<File> > &) const {
+    std::vector<std::unique_ptr<File>> &) const {
+  // Create a file for __ImageBase.
   std::unique_ptr<SimpleFileNode> fileNode(
       new SimpleFileNode("Implicit Files"));
   std::unique_ptr<File> linkerGeneratedSymFile(
@@ -111,11 +126,21 @@ bool PECOFFLinkingContext::createImplicitFiles(
   getInputGraph().insertElementAt(std::move(fileNode),
                                   InputGraph::Position::END);
 
+  // Create a file for _imp_ symbols.
   std::unique_ptr<SimpleFileNode> impFileNode(new SimpleFileNode("imp"));
   impFileNode->appendInputFile(
       std::unique_ptr<File>(new pecoff::LocallyImportedSymbolFile(*this)));
   getInputGraph().insertElementAt(std::move(impFileNode),
                                   InputGraph::Position::END);
+
+  // Create a file for dllexported symbols.
+  std::unique_ptr<SimpleFileNode> exportNode(new SimpleFileNode("<export>"));
+  pecoff::ExportedSymbolRenameFile *renameFile =
+      new pecoff::ExportedSymbolRenameFile(*this);
+  exportNode->appendInputFile(std::unique_ptr<File>(renameFile));
+  getLibraryGroup()->addFile(std::move(exportNode));
+  getInputGraph().registerObserver(
+      *(new (_allocator) ObserverCallback(renameFile)));
   return true;
 }
 
diff --git a/lld/test/pecoff/Inputs/export.obj.yaml b/lld/test/pecoff/Inputs/export.obj.yaml
index 9256fae0867..524a5489602 100644
--- a/lld/test/pecoff/Inputs/export.obj.yaml
+++ b/lld/test/pecoff/Inputs/export.obj.yaml
@@ -54,4 +54,10 @@ symbols:
     SimpleType:      IMAGE_SYM_TYPE_NULL
     ComplexType:     IMAGE_SYM_DTYPE_NULL
     StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            _exportfn7@8
+    Value:           16
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
 ...
diff --git a/lld/test/pecoff/export.test b/lld/test/pecoff/export.test
index 03d08981866..7d45b362fd5 100644
--- a/lld/test/pecoff/export.test
+++ b/lld/test/pecoff/export.test
@@ -41,3 +41,12 @@ CHECK4-NEXT:        5   0x2008  exportfn1
 CHECK4-NEXT:        6   0x2010  exportfn2
 CHECK4-NEXT:        7   0x2010  exportfn5
 CHECK4-NEXT:        8   0x2010  exportfn3@256
+
+# RUN: lld -flavor link /out:%t5.dll /dll /entry:init \
+# RUN:   /export:exportfn7 -- %t.obj
+# RUN: llvm-objdump -p %t5.dll | FileCheck -check-prefix=CHECK5 %s
+
+CHECK5:      Export Table:
+CHECK5:      DLL name: export.test.tmp5.dll
+CHECK5:       Ordinal      RVA  Name
+CHECK5-NEXT:        1   0x2010  exportfn7