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

[ruiu]

Summary

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.

[Bigcheese]

Would it be better to lazily initialize the _defined member when find is called? Or is this going to almost always be used?

[ruiu]

It'll rarely be used, so lazy initialization is a good idea. I'll do that.

[ruiu]

Closed by commit rL208754 (authored by @ruiu).

lld/trunk/lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.h

 //===- lib/ReaderWriter/PECOFF/LinkerGeneratedSymbolFile.cpp --------------===//
 //
 //                             The LLVM Linker
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Atoms.h"
 
 #include "lld/Core/ArchiveLibraryFile.h"
 #include "lld/ReaderWriter/PECOFFLinkingContext.h"
 #include "lld/ReaderWriter/Simple.h"
 #include "llvm/Support/Allocator.h"
 
+#include <mutex>
+
 namespace lld {
 namespace pecoff {
 
 namespace {
 
 /// The defined atom for dllexported symbols with __imp_ prefix.
 class ImpPointerAtom : public COFFLinkerInternalAtom {
 public:

 
 private:
   atom_collection_vector<DefinedAtom> _definedAtoms;
   atom_collection_vector<UndefinedAtom> _undefinedAtoms;
   atom_collection_vector<SharedLibraryAtom> _sharedLibraryAtoms;
   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
 // ___ImageBase on x86) is a linker-generated symbol whose address is the same
 // as the image base address.
 class LinkerGeneratedSymbolFile : public SimpleFile {
 public:
   LinkerGeneratedSymbolFile(const PECOFFLinkingContext &ctx)
public:
   }
 
 private:
   std::string _prefix;
   mutable uint64_t _ordinal;
   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

lld/trunk/lib/ReaderWriter/PECOFF/PECOFFLinkingContext.cpp

std::unique_ptr<File> PECOFFLinkingContext::createEntrySymbolFile() const {
   return LinkingContext::createEntrySymbolFile("<command line option /entry>");
 }
 
 std::unique_ptr<File> PECOFFLinkingContext::createUndefinedSymbolFile() const {
   return LinkingContext::createUndefinedSymbolFile(
       "<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(
       new pecoff::LinkerGeneratedSymbolFile(*this));
   fileNode->appendInputFile(std::move(linkerGeneratedSymFile));
   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;
 }
 
 /// Returns the section name in the resulting executable.
 ///
 /// Sections in object files are usually output to the executable with the same
 /// name, but you can rename by command line option. /merge:from=to makes the
 /// linker to combine "from" section contents to "to" section in the

lld/trunk/test/pecoff/Inputs/export.obj.yaml

- Name:            exportfn3@256
     ComplexType:     IMAGE_SYM_DTYPE_NULL
     StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
   - Name:            _exportfn6
     Value:           16
     SectionNumber:   1
     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
 ...

lld/trunk/test/pecoff/export.test

 
 CHECK4:      Export Table:
 CHECK4:      DLL name: export.test.tmp4.dll
 CHECK4:       Ordinal      RVA  Name
 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