[PECOFF] Process Import Name/Type field in the import library.

This patch removes hacky mangle() function, which strips all decorations
uncondtitionally. LLD now interprets Import Name/Type field in the import
library properly as described in the Microsoft PE/COFF Spec.

llvm-svn: 187388

6 files changed, 101 insertions, 47 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/Atoms.h b/lld/lib/ReaderWriter/PECOFF/Atoms.h
index 0db4f3cd8e7..b2b2c737480 100644
--- a/lld/lib/ReaderWriter/PECOFF/Atoms.h
+++ b/lld/lib/ReaderWriter/PECOFF/Atoms.h
@@ -264,17 +264,23 @@ private:
 // to an import address table entry in Idata pass.
 class COFFSharedLibraryAtom : public SharedLibraryAtom {
 public:
-  COFFSharedLibraryAtom(const File &file, uint16_t hint,
-                        StringRef symbolName, StringRef loadName)
-      : _file(file), _hint(hint), _unmangledName(symbolName),
-        _loadName(loadName), _mangledName(addImpPrefix(symbolName)),
+  COFFSharedLibraryAtom(const File &file, uint16_t hint, StringRef symbolName,
+                        StringRef importName, StringRef dllName)
+      : _file(file), _hint(hint), _mangledName(addImpPrefix(symbolName)),
+        _importName(importName), _dllName(dllName),
         _importTableEntry(nullptr) {}
 
   virtual const File &file() const { return _file; }
   uint16_t hint() const { return _hint; }
+
+  /// Returns the symbol name to be used by the core linker.
   virtual StringRef name() const { return _mangledName; }
-  virtual StringRef unmangledName() const { return _unmangledName; }
-  virtual StringRef loadName() const { return _loadName; }
+
+  /// Returns the symbol name to be used in the import description table in the
+  /// COFF header.
+  virtual StringRef importName() const { return _importName; }
+
+  virtual StringRef loadName() const { return _dllName; }
   virtual bool canBeNullAtRuntime() const { return false; }
 
   void setImportTableEntry(const DefinedAtom *atom) {
@@ -296,9 +302,9 @@ private:
 
   const File &_file;
   uint16_t _hint;
-  StringRef _unmangledName;
-  StringRef _loadName;
   std::string _mangledName;
+  std::string _importName;
+  StringRef _dllName;
   const DefinedAtom *_importTableEntry;
 };
 
diff --git a/lld/lib/ReaderWriter/PECOFF/IdataPass.h b/lld/lib/ReaderWriter/PECOFF/IdataPass.h
index 88a84d378d3..07d619696e4 100644
--- a/lld/lib/ReaderWriter/PECOFF/IdataPass.h
+++ b/lld/lib/ReaderWriter/PECOFF/IdataPass.h
@@ -116,41 +116,41 @@ private:
 /// loader can find the symbol quickly.
 class HintNameAtom : public IdataAtom {
 public:
-  HintNameAtom(Context &ctx, uint16_t hint, StringRef name)
-      : IdataAtom(ctx.file, assembleRawContent(hint, name)), _name(name) {
+  HintNameAtom(Context &ctx, uint16_t hint, StringRef importName)
+      : IdataAtom(ctx.file, assembleRawContent(hint, importName)),
+        _importName(importName) {
     ctx.hintNameAtoms.push_back(this);
   }
 
-  StringRef getContentString() { return _name; }
+  StringRef getContentString() { return _importName; }
 
 private:
   // The first two bytes of the content is a hint, followed by a null-terminated
   // symbol name. The total size needs to be multiple of 2.
-  vector<uint8_t> assembleRawContent(uint16_t hint, StringRef name) {
-    name = unmangle(name);
-    size_t size = llvm::RoundUpToAlignment(sizeof(hint) + name.size() + 1, 2);
+  vector<uint8_t> assembleRawContent(uint16_t hint, StringRef importName) {
+    size_t size = llvm::RoundUpToAlignment(sizeof(hint) + importName.size() + 1, 2);
     vector<uint8_t> ret(size);
-    ret[name.size()] = 0;
-    ret[name.size() - 1] = 0;
+    ret[importName.size()] = 0;
+    ret[importName.size() - 1] = 0;
     *reinterpret_cast<llvm::support::ulittle16_t *>(&ret[0]) = hint;
-    std::memcpy(&ret[2], name.data(), name.size());
+    std::memcpy(&ret[2], importName.data(), importName.size());
     return ret;
   }
 
-  /// Undo name mangling. In Windows, the symbol name for function is encoded
-  /// as "_name@X", where X is the number of bytes of the arguments.
-  StringRef unmangle(StringRef mangledName) {
-    assert(mangledName.startswith("_"));
-    return mangledName.substr(1).split('@').first;
-  }
-
-  StringRef _name;
+  StringRef _importName;
 };
 
 class ImportTableEntryAtom : public IdataAtom {
 public:
-  explicit ImportTableEntryAtom(Context &ctx)
-      : IdataAtom(ctx.file, vector<uint8_t>(4, 0)) {}
+  explicit ImportTableEntryAtom(Context &ctx, uint32_t contents)
+      : IdataAtom(ctx.file, assembleRawContent(contents)) {}
+
+private:
+  vector<uint8_t> assembleRawContent(uint32_t contents) {
+    vector<uint8_t> ret(4);
+    *reinterpret_cast<llvm::support::ulittle32_t *>(&ret[0]) = contents;
+    return ret;
+  }
 };
 
 /// An ImportDirectoryAtom includes information to load a DLL, including a DLL
@@ -194,21 +194,29 @@ private:
                            const vector<COFFSharedLibraryAtom *> &sharedAtoms,
                            bool shouldAddReference,
                            vector<ImportTableEntryAtom *> &ret) const {
-    for (COFFSharedLibraryAtom *shared : sharedAtoms) {
-      HintNameAtom *hintName = createHintNameAtom(ctx, shared);
-      ImportTableEntryAtom *entry = new (_alloc) ImportTableEntryAtom(ctx);
-      addDir32NBReloc(entry, hintName);
+    for (COFFSharedLibraryAtom *atom : sharedAtoms) {
+      ImportTableEntryAtom *entry = nullptr;
+      if (atom->importName().empty()) {
+        // Import by ordinal
+        uint32_t hint = (1U << 31) | atom->hint();
+        entry = new (_alloc) ImportTableEntryAtom(ctx, hint);
+      } else {
+        // Import by name
+        entry = new (_alloc) ImportTableEntryAtom(ctx, 0);
+        HintNameAtom *hintName = createHintNameAtom(ctx, atom);
+        addDir32NBReloc(entry, hintName);
+      }
       ret.push_back(entry);
       if (shouldAddReference)
-        shared->setImportTableEntry(entry);
+        atom->setImportTableEntry(entry);
     }
     // Add the NULL entry.
-    ret.push_back(new (_alloc) ImportTableEntryAtom(ctx));
+    ret.push_back(new (_alloc) ImportTableEntryAtom(ctx, 0));
   }
 
   HintNameAtom *createHintNameAtom(
       Context &ctx, const COFFSharedLibraryAtom *atom) const {
-    return new (_alloc) HintNameAtom(ctx, atom->hint(), atom->unmangledName());
+    return new (_alloc) HintNameAtom(ctx, atom->hint(), atom->importName());
   }
 
   mutable llvm::BumpPtrAllocator _alloc;
diff --git a/lld/lib/ReaderWriter/PECOFF/ReaderImportHeader.cpp b/lld/lib/ReaderWriter/PECOFF/ReaderImportHeader.cpp
index 35432cd5a9f..480cae4d53d 100644
--- a/lld/lib/ReaderWriter/PECOFF/ReaderImportHeader.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/ReaderImportHeader.cpp
@@ -184,10 +184,20 @@ public:
     StringRef symbolName(buf + sizeof(COFF::ImportHeader));
     StringRef dllName(buf + sizeof(COFF::ImportHeader) + symbolName.size() + 1);
 
+    // TypeInfo is a bitfield. The least significant 2 bits are import
+    // type, followed by 3 bit import name type.
+    uint16_t typeInfo = *reinterpret_cast<const support::ulittle16_t *>(
+        buf + offsetof(COFF::ImportHeader, TypeInfo));
+    int type = typeInfo & 0x3;
+    int nameType = (typeInfo >> 2) & 0x7;
+
+    // Symbol name used by the linker may be different from the symbol name used
+    // by the loader. The latter may lack symbol decorations, or may not even
+    // have name if it's imported by ordinal.
+    StringRef importName = symbolNameToImportName(symbolName, nameType);
+
     const COFFSharedLibraryAtom *dataAtom = addSharedLibraryAtom(
-        hint, symbolName, dllName);
-    int type = *reinterpret_cast<const support::ulittle16_t *>(
-        buf + offsetof(COFF::ImportHeader, TypeInfo)) >> 14;
+        hint, symbolName, importName, dllName);
     if (type == llvm::COFF::IMPORT_CODE)
       addDefinedAtom(symbolName, dllName, dataAtom);
 
@@ -213,10 +223,11 @@ public:
   virtual const TargetInfo &getTargetInfo() const { return _targetInfo; }
 
 private:
-  const COFFSharedLibraryAtom *addSharedLibraryAtom(
-      uint16_t hint, StringRef symbolName, StringRef dllName) {
+  const COFFSharedLibraryAtom *
+  addSharedLibraryAtom(uint16_t hint, StringRef symbolName,
+                       StringRef importName, StringRef dllName) {
     auto *atom = new (_alloc) COFFSharedLibraryAtom(
-        *this, hint, symbolName, dllName);
+        *this, hint, symbolName, importName, dllName);
     _sharedLibraryAtoms._atoms.push_back(atom);
     return atom;
   }
@@ -235,6 +246,32 @@ private:
   atom_collection_vector<SharedLibraryAtom> _sharedLibraryAtoms;
   const TargetInfo &_targetInfo;
   mutable llvm::BumpPtrAllocator _alloc;
+
+  // Convert the given symbol name to the import symbol name exported by the
+  // DLL.
+  StringRef symbolNameToImportName(StringRef symbolName, int nameType) const {
+    StringRef ret;
+    switch (nameType) {
+    case llvm::COFF::IMPORT_ORDINAL:
+      // The import is by ordinal. No symbol name will be used to identify the
+      // item in the DLL. Only its ordinal will be used.
+      return "";
+    case llvm::COFF::IMPORT_NAME:
+      // The import name in this case is identical to the symbol name.
+      return symbolName;
+    case llvm::COFF::IMPORT_NAME_NOPREFIX:
+      // The import name is the symbol name without leading ?, @ or _.
+      ret = symbolName.ltrim("?@_");
+      break;
+    case llvm::COFF::IMPORT_NAME_UNDECORATE:
+      // Similar to NOPREFIX, but we also need to truncate at the first @.
+      ret = symbolName.ltrim("?@_");
+      ret = ret.substr(0, ret.find('@'));
+      break;
+    }
+    std::string *str = new (_alloc) std::string(ret);
+    return *str;
+  }
 };
 
 } // end anonymous namespace
diff --git a/lld/test/pecoff/Inputs/vars.c b/lld/test/pecoff/Inputs/vars.c
index 6532824fa3f..7d2ab9bfdea 100644
--- a/lld/test/pecoff/Inputs/vars.c
+++ b/lld/test/pecoff/Inputs/vars.c
@@ -1,8 +1,11 @@
 // cl.exe /c vars.c
-// link.exe /debug /nodefaultlib /entry:dllmain vars.obj
-__declspec(dllexport) int var = 3;
+// link /debug /dll /nodefaultlib /entry:dllmain /export:var,@1,NONAME,DATA /export:fn vars.obj
 
-__declspec(dllexport) int fn(void) {
+// will be exported by ordinal
+int var = 3;
+
+// will be exported by name
+int fn(void) {
   return 4;
 }
 
diff --git a/lld/test/pecoff/Inputs/vars.lib b/lld/test/pecoff/Inputs/vars.lib
index 6ed368e1b31..329362bc520 100644
--- a/lld/test/pecoff/Inputs/vars.lib
+++ b/lld/test/pecoff/Inputs/vars.lib
diff --git a/lld/test/pecoff/importlib.test b/lld/test/pecoff/importlib.test
index cd3da3accd5..9d015c12221 100644
--- a/lld/test/pecoff/importlib.test
+++ b/lld/test/pecoff/importlib.test
@@ -1,5 +1,5 @@
-# Verify that lld can handle .lib files. "main.obj" refers _val1 and
-# _val2 that are defined in "vars.lib".
+# Verify that lld can handle .lib files. "main.obj" refers "var" and
+# "fn" defined in "vars.lib".
 #
 # RUN: yaml2obj %p/Inputs/vars-main.obj.yaml > %t.obj
 #
@@ -22,8 +22,8 @@ CHECK: Disassembly of section .text:
 CHECK: .text:
 CHECK:     1000:       55                    pushl   %ebp
 CHECK:     1001:       8b ec                 movl    %esp, %ebp
-CHECK:     1003:       ff 15 4c 20 40 00     calll   *4202572
-CHECK:     1009:       8b 0d 50 20 40 00     movl    4202576, %ecx
+CHECK:     1003:       ff 15 40 20 40 00     calll   *4202560
+CHECK:     1009:       8b 0d 44 20 40 00     movl    4202564, %ecx
 CHECK:     100f:       03 01                 addl    (%ecx), %eax
 CHECK:     1011:       5d                    popl    %ebp
 CHECK:     1012:       c3                    ret