[lld] r242110 - COFF: Fix entry name inference for x86.

[Rui Ueyama]

Author: ruiu
Date: Mon Jul 13 21:58:13 2015
New Revision: 242110

URL: http://llvm.org/viewvc/llvm-project?rev=242110&view=rev
Log:
COFF: Fix entry name inference for x86.

Entry name selection rule is already complicated on x64, but it's more
complicated on x86 because of the underscore name mangling scheme.

If one of _main, _main@<number> (a C function) or ?main@@... (a C++ function)
is defined, entry name is _mainCRTStartup. If _wmain, _wmain@<number or
?wmain@@... is defined, entry name is _wmainCRTStartup. And so on.

Added:
    lld/trunk/test/COFF/entry-inference32.test
Modified:
    lld/trunk/COFF/Driver.cpp
    lld/trunk/COFF/SymbolTable.cpp
    lld/trunk/COFF/SymbolTable.h

Modified: lld/trunk/COFF/Driver.cpp
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/COFF/Driver.cpp?rev=242110&r1=242109&r2=242110&view=diff
==============================================================================
--- lld/trunk/COFF/Driver.cpp (original)
+++ lld/trunk/COFF/Driver.cpp Mon Jul 13 21:58:13 2015
@@ -227,8 +227,8 @@ StringRef LinkerDriver::findDefaultEntry
       {"wWinMain", "wWinMainCRTStartup"},
   };
   for (auto E : Entries) {
-    Symbol *Sym = Symtab.find(mangle(E[0]));
-    if (Sym && !isa<Undefined>(Sym->Body))
+    StringRef Entry = Symtab.findMangle(mangle(E[0]));
+    if (!Entry.empty() && !isa<Undefined>(Symtab.find(Entry)->Body))
       return mangle(E[1]);
   }
   return "";

Modified: lld/trunk/COFF/SymbolTable.cpp
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/COFF/SymbolTable.cpp?rev=242110&r1=242109&r2=242110&view=diff
==============================================================================
--- lld/trunk/COFF/SymbolTable.cpp (original)
+++ lld/trunk/COFF/SymbolTable.cpp Mon Jul 13 21:58:13 2015
@@ -274,28 +274,39 @@ Symbol *SymbolTable::find(StringRef Name
   return It->second;
 }
 
+StringRef SymbolTable::findByPrefix(StringRef Prefix) {
+  for (auto Pair : Symtab) {
+    StringRef Name = Pair.first;
+    if (Name.startswith(Prefix))
+      return Name;
+  }
+  return "";
+}
+
+StringRef SymbolTable::findMangle(StringRef Name) {
+  if (Symbol *Sym = find(Name))
+    if (!isa<Undefined>(Sym->Body))
+      return Name;
+  if (Config->is64())
+    return findByPrefix(("?" + Name + "@@Y").str());
+  if (!Name.startswith("_"))
+    return "";
+  // Search for x86 C function.
+  StringRef S = findByPrefix((Name + "@").str());
+  if (!S.empty())
+    return S;
+  // Search for x86 C++ non-member function.
+  return findByPrefix(("?" + Name.substr(1) + "@@Y").str());
+}
+
 void SymbolTable::mangleMaybe(Undefined *U) {
   if (U->WeakAlias)
     return;
   if (!isa<Undefined>(U->repl()))
     return;
-
-  // In Microsoft ABI, a non-member function name is mangled this way.
-  std::string Prefix;
-  if (Config->is64()) {
-    Prefix = ("?" + U->getName() + "@@Y").str();
-  } else {
-    if (!U->getName().startswith("_"))
-      return;
-    Prefix = ("?" + U->getName().substr(1) + "@@Y").str();
-  }
-  for (auto Pair : Symtab) {
-    StringRef Name = Pair.first;
-    if (!Name.startswith(Prefix))
-      continue;
-    U->WeakAlias = addUndefined(Name);
-    return;
-  }
+  StringRef Alias = findMangle(U->getName());
+  if (!Alias.empty())
+    U->WeakAlias = addUndefined(Alias);
 }
 
 Undefined *SymbolTable::addUndefined(StringRef Name) {

Modified: lld/trunk/COFF/SymbolTable.h
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/COFF/SymbolTable.h?rev=242110&r1=242109&r2=242110&view=diff
==============================================================================
--- lld/trunk/COFF/SymbolTable.h (original)
+++ lld/trunk/COFF/SymbolTable.h Mon Jul 13 21:58:13 2015
@@ -62,6 +62,7 @@ public:
   // for U from the symbol table, and if found, set the symbol as
   // a weak alias for U.
   void mangleMaybe(Undefined *U);
+  StringRef findMangle(StringRef Name);
 
   // Print a layout map to OS.
   void printMap(llvm::raw_ostream &OS);
@@ -92,6 +93,7 @@ private:
   std::error_code addSymbol(SymbolBody *New);
   void addLazy(Lazy *New, std::vector<Symbol *> *Accum);
   Symbol *insert(SymbolBody *New);
+  StringRef findByPrefix(StringRef Prefix);
 
   std::error_code addMemberFile(Lazy *Body);
   ErrorOr<ObjectFile *> createLTOObject(llvm::LTOCodeGenerator *CG);

Added: lld/trunk/test/COFF/entry-inference32.test
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/test/COFF/entry-inference32.test?rev=242110&view=auto
==============================================================================
--- lld/trunk/test/COFF/entry-inference32.test (added)
+++ lld/trunk/test/COFF/entry-inference32.test Mon Jul 13 21:58:13 2015
@@ -0,0 +1,35 @@
+# RUN: yaml2obj < %s > %t.obj
+# RUN: not lld -flavor link2 /out:%t.exe %t.obj /verbose > %t.log 2>&1
+# RUN: FileCheck %s < %t.log
+
+# CHECK: Entry name inferred: _WinMainCRTStartup
+
+---
+header:
+  Machine:         IMAGE_FILE_MACHINE_I386
+  Characteristics: []
+sections:
+  - Name:            .text
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       4
+    SectionData:     B82A000000C3
+symbols:
+  - Name:            .text
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          6
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          0
+  - Name:            _WinMain at 16
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...