COFF: Infer entry symbol name if /entry is not given.

`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, {w,}{WinMain,main}CRTStartup, respectively.
The linker needs to choose the right one depending on which `main`
function is defined.

llvm-svn: 238667

lld/COFF/Config.h

@@ -26,7 +26,7 @@ public:
   llvm::COFF::MachineTypes MachineType = llvm::COFF::IMAGE_FILE_MACHINE_AMD64;
   bool Verbose = false;
   WindowsSubsystem Subsystem = llvm::COFF::IMAGE_SUBSYSTEM_WINDOWS_CUI;
-  std::string EntryName = "mainCRTStartup";
+  std::string EntryName;
 
   uint64_t ImageBase = 0x140000000;
   uint64_t StackReserve = 1024 * 1024;

lld/COFF/Driver.cpp

@@ -183,6 +183,20 @@ bool link(int Argc, const char *Argv[]) {
       return false;
     }
   }
+
+  // Windows specific -- If entry point name is not given, we need to
+  // infer that from user-defined entry name. The symbol table takes
+  // care of details.
+  if (Config->EntryName.empty()) {
+    auto EntryOrErr = Symtab.findDefaultEntry();
+    if (auto EC = EntryOrErr.getError()) {
+      llvm::errs() << EC.message() << "\n";
+      return false;
+    }
+    Config->EntryName = EntryOrErr.get();
+  }
+
+  // Make sure we have resolved all symbols.
   if (Symtab.reportRemainingUndefines())
     return false;
 

lld/COFF/SymbolTable.cpp

@@ -19,13 +19,9 @@ namespace lld {
 namespace coff {
 
 SymbolTable::SymbolTable() {
-  addInitialSymbol(new DefinedAbsolute("__ImageBase", Config->ImageBase));
-  addInitialSymbol(new Undefined(Config->EntryName));
-}
-
-void SymbolTable::addInitialSymbol(SymbolBody *Body) {
-  OwnedSymbols.push_back(std::unique_ptr<SymbolBody>(Body));
-  Symtab[Body->getName()] = new (Alloc) Symbol(Body);
+  addSymbol(new DefinedAbsolute("__ImageBase", Config->ImageBase));
+  if (!Config->EntryName.empty())
+    addSymbol(new Undefined(Config->EntryName));
 }
 
 std::error_code SymbolTable::addFile(std::unique_ptr<InputFile> File) {
@@ -155,11 +151,40 @@ std::vector<Chunk *> SymbolTable::getChunks() {
   return Res;
 }
 
-SymbolBody *SymbolTable::find(StringRef Name) {
+Defined *SymbolTable::find(StringRef Name) {
   auto It = Symtab.find(Name);
   if (It == Symtab.end())
     return nullptr;
-  return It->second->Body;
+  if (auto *Def = dyn_cast<Defined>(It->second->Body))
+    return Def;
+  return nullptr;
+}
+
+// Link default entry point name.
+ErrorOr<StringRef> SymbolTable::findDefaultEntry() {
+  static const char *Entries[][2] = {
+      {"mainCRTStartup", "mainCRTStartup"},
+      {"wmainCRTStartup", "wmainCRTStartup"},
+      {"WinMainCRTStartup", "WinMainCRTStartup"},
+      {"wWinMainCRTStartup", "wWinMainCRTStartup"},
+      {"main", "mainCRTStartup"},
+      {"wmain", "wmainCRTStartup"},
+      {"WinMain", "WinMainCRTStartup"},
+      {"wWinMain", "wWinMainCRTStartup"},
+  };
+  for (size_t I = 0; I < sizeof(Entries); ++I) {
+    if (!find(Entries[I][0]))
+      continue;
+    if (auto EC = addSymbol(new Undefined(Entries[I][1])))
+      return EC;
+    return StringRef(Entries[I][1]);
+  }
+  return make_dynamic_error_code("entry point must be defined");
+}
+
+std::error_code SymbolTable::addSymbol(SymbolBody *Body) {
+  OwningSymbols.push_back(std::unique_ptr<SymbolBody>(Body));
+  return resolve(Body);
 }
 
 void SymbolTable::dump() {

lld/COFF/SymbolTable.h

@@ -44,7 +44,16 @@ public:
   // returned symbol actually has the same name (because of various
   // mechanisms to allow aliases, a name can be resolved to a
   // different symbol). Returns a nullptr if not found.
-  SymbolBody *find(StringRef Name);
+  Defined *find(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,
+  // {w,}{WinMain,main}CRTStartup, respectively. The linker needs to
+  // choose the right one depending on which `main` function is defined.
+  // This function looks up the symbol table and resolve corresponding
+  // entry point name.
+  ErrorOr<StringRef> findDefaultEntry();
 
   // Dump contents of the symbol table to stderr.
   void dump();
@@ -63,11 +72,11 @@ private:
 
   std::error_code resolve(SymbolBody *Body);
   std::error_code addMemberFile(Lazy *Body);
-  void addInitialSymbol(SymbolBody *Body);
+  std::error_code addSymbol(SymbolBody *Body);
 
   std::unordered_map<StringRef, Symbol *> Symtab;
   std::vector<std::unique_ptr<ArchiveFile>> ArchiveFiles;
-  std::vector<std::unique_ptr<SymbolBody>> OwnedSymbols;
+  std::vector<std::unique_ptr<SymbolBody>> OwningSymbols;
   llvm::BumpPtrAllocator Alloc;
   StringAllocator StringAlloc;
 };

lld/test/COFF/entry-inference.test

@@ -0,0 +1,50 @@
+# RUN: sed -e s/ENTRYNAME/main/ %s | yaml2obj > %t.obj
+# RUN: not lld -flavor link2 /out:%t.exe %t.obj > %t.log 2>&1
+# RUN: FileCheck -check-prefix=MAIN %s < %t.log
+
+# RUN: sed s/ENTRYNAME/wmain/ %s | yaml2obj > %t.obj
+# RUN: not lld -flavor link2 /out:%t.exe %t.obj > %t.log 2>&1
+# RUN: FileCheck -check-prefix=WMAIN %s < %t.log
+
+# RUN: sed s/ENTRYNAME/WinMain/ %s | yaml2obj > %t.obj
+# RUN: not lld -flavor link2 /out:%t.exe %t.obj > %t.log 2>&1
+# RUN: FileCheck -check-prefix=WINMAIN %s < %t.log
+
+# RUN: sed s/ENTRYNAME/wWinMain/ %s | yaml2obj > %t.obj
+# RUN: not lld -flavor link2 /out:%t.exe %t.obj > %t.log 2>&1
+# RUN: FileCheck -check-prefix=WWINMAIN %s < %t.log
+
+# MAIN:     undefined symbol: mainCRTStartup
+# WMAIN:    undefined symbol: wmainCRTStartup
+# WINMAIN:  undefined symbol: WinMainCRTStartup
+# WWINMAIN: undefined symbol: wWinMainCRTStartup
+
+---
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  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:            ENTRYNAME
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...