4 files changed, 68 insertions, 66 deletions

diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 12cff8a4f9af..8f34d1f98c1d 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -207,15 +207,30 @@ void LinkerDriver::addUndefined(StringRef Sym) {
   Config->GCRoots.insert(Sym);
 }
 
-static WindowsSubsystem inferSubsystem() {
+// Windows specific -- find default entry point name.
+StringRef LinkerDriver::findDefaultEntry() {
+  // User-defined main functions and their corresponding entry points.
+  static const char *Entries[][2] = {
+      {"main", "mainCRTStartup"},
+      {"wmain", "wmainCRTStartup"},
+      {"WinMain", "WinMainCRTStartup"},
+      {"wWinMain", "wWinMainCRTStartup"},
+  };
+  for (auto E : Entries) {
+    if (Symtab.findLazy(E[0]))
+      return E[1];
+  }
+  return "";
+}
+
+WindowsSubsystem LinkerDriver::inferSubsystem() {
   if (Config->DLL)
     return IMAGE_SUBSYSTEM_WINDOWS_GUI;
-  return StringSwitch<WindowsSubsystem>(Config->EntryName)
-      .Case("mainCRTStartup", IMAGE_SUBSYSTEM_WINDOWS_CUI)
-      .Case("wmainCRTStartup", IMAGE_SUBSYSTEM_WINDOWS_CUI)
-      .Case("WinMainCRTStartup", IMAGE_SUBSYSTEM_WINDOWS_GUI)
-      .Case("wWinMainCRTStartup", IMAGE_SUBSYSTEM_WINDOWS_GUI)
-      .Default(IMAGE_SUBSYSTEM_UNKNOWN);
+  if (Symtab.find("main") || Symtab.find("wmain"))
+    return IMAGE_SUBSYSTEM_WINDOWS_CUI;
+  if (Symtab.find("WinMain") || Symtab.find("wWinMain"))
+    return IMAGE_SUBSYSTEM_WINDOWS_GUI;
+  return IMAGE_SUBSYSTEM_UNKNOWN;
 }
 
 bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
@@ -522,6 +537,18 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     return false;
   }
 
+  // Windows specific -- If entry point name is not given, we need to
+  // infer that from user-defined entry name.
+  if (Config->EntryName.empty() && !Config->NoEntry) {
+    StringRef S = findDefaultEntry();
+    if (S.empty()) {
+      llvm::errs() << "entry point must be defined\n";
+      return false;
+    }
+    Config->EntryName = S;
+    addUndefined(S);
+  }
+
   // Resolve auxiliary symbols until converge.
   // (Trying to resolve a symbol may trigger a Lazy symbol to load a new file.
   // A new file may contain a directive section to add new command line options.
@@ -545,19 +572,6 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
       }
     }
 
-    // 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() && !Config->NoEntry) {
-      auto EntryOrErr = Symtab.findDefaultEntry();
-      if (auto EC = EntryOrErr.getError()) {
-        llvm::errs() << EC.message() << "\n";
-        return false;
-      }
-      Config->EntryName = EntryOrErr.get();
-      addUndefined(Config->EntryName);
-    }
-
     if (auto EC = Symtab.run()) {
       llvm::errs() << EC.message() << "\n";
       return false;
diff --git a/lld/COFF/Driver.h b/lld/COFF/Driver.h
index 509edb757a6b..e861dead817f 100644
--- a/lld/COFF/Driver.h
+++ b/lld/COFF/Driver.h
@@ -94,6 +94,16 @@ private:
 
   void addUndefined(StringRef Sym);
 
+  // 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.
+  StringRef findDefaultEntry();
+  WindowsSubsystem inferSubsystem();
+
   // Driver is the owner of all opened files.
   // InputFiles have MemoryBufferRefs to them.
   std::vector<std::unique_ptr<MemoryBuffer>> OwningMBs;
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
index 831ab7d8a7ee..4b6555210d1a 100644
--- a/lld/COFF/SymbolTable.cpp
+++ b/lld/COFF/SymbolTable.cpp
@@ -174,6 +174,25 @@ Defined *SymbolTable::find(StringRef Name) {
   return nullptr;
 }
 
+// Find a given symbol. If a lazy symbol is found,
+// resolve that before returning.
+Defined *SymbolTable::findLazy(StringRef Name) {
+  auto It = Symtab.find(Name);
+  if (It == Symtab.end())
+    return nullptr;
+  Symbol *Sym = It->second;
+  if (auto *B = dyn_cast<Defined>(Sym->Body))
+    return B;
+  if (auto *B = dyn_cast<Lazy>(Sym->Body)) {
+    if (addMemberFile(B))
+      return nullptr;
+    if (run())
+      return nullptr;
+    return cast<Defined>(Sym->Body);
+  }
+  return nullptr;
+}
+
 // Find a given symbol or its mangled symbol.
 std::pair<StringRef, Symbol *> SymbolTable::findMangled(StringRef S) {
   auto It = Symtab.find(S);
@@ -196,41 +215,6 @@ std::pair<StringRef, Symbol *> SymbolTable::findMangled(StringRef S) {
   return std::make_pair(S, nullptr);
 }
 
-std::error_code SymbolTable::resolveLazy(StringRef Name) {
-  auto It = Symtab.find(Name);
-  if (It == Symtab.end())
-    return std::error_code();
-  if (auto *B = dyn_cast<Lazy>(It->second->Body)) {
-    if (auto EC = addMemberFile(B))
-      return EC;
-    return run();
-  }
-  return std::error_code();
-}
-
-// Windows specific -- Link default entry point name.
-ErrorOr<StringRef> SymbolTable::findDefaultEntry() {
-  // User-defined main functions and their corresponding entry points.
-  static const char *Entries[][2] = {
-      {"main", "mainCRTStartup"},
-      {"wmain", "wmainCRTStartup"},
-      {"WinMain", "WinMainCRTStartup"},
-      {"wWinMain", "wWinMainCRTStartup"},
-  };
-  for (auto E : Entries) {
-    resolveLazy(E[1]);
-    if (find(E[1]))
-      return StringRef(E[1]);
-    if (!find(E[0]))
-      continue;
-    if (auto EC = resolve(new (Alloc) Undefined(E[1])))
-      return EC;
-    return StringRef(E[1]);
-  }
-  llvm::errs() << "entry point must be defined\n";
-  return make_error_code(LLDError::InvalidOption);
-}
-
 std::error_code SymbolTable::addUndefined(StringRef Name) {
   return resolve(new (Alloc) Undefined(Name));
 }
@@ -248,7 +232,10 @@ std::error_code SymbolTable::rename(StringRef From, StringRef To) {
   SymbolBody *Body = new (Alloc) Undefined(To);
   if (auto EC = resolve(Body))
     return EC;
-  Sym->Body = Body->getReplacement();
+  SymbolBody *Repl = Body->getReplacement();
+  if (isa<Undefined>(Repl))
+    return std::error_code();
+  Sym->Body = Repl;
   Body->setBackref(Sym);
   ++Version;
   return std::error_code();
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
index 62f14b09f525..b7c8888238c8 100644
--- a/lld/COFF/SymbolTable.h
+++ b/lld/COFF/SymbolTable.h
@@ -51,21 +51,13 @@ public:
   // mechanisms to allow aliases, a name can be resolved to a
   // different symbol). Returns a nullptr if not found.
   Defined *find(StringRef Name);
+  Defined *findLazy(StringRef Name);
 
   // Find a symbol assuming that Name is a function name.
   // Not only a given string but its mangled names (in MSVC C++ manner)
   // will be searched.
   std::pair<StringRef, Symbol *> findMangled(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();
-
   // Print a layout map to OS.
   void printMap(llvm::raw_ostream &OS);
 
@@ -92,7 +84,6 @@ public:
 
 private:
   std::error_code resolve(SymbolBody *Body);
-  std::error_code resolveLazy(StringRef Name);
   std::error_code addMemberFile(Lazy *Body);
   ErrorOr<ObjectFile *> createLTOObject(llvm::LTOCodeGenerator *CG);