COFF: Simplify and rename findMangle. NFC.

Occasionally we have to resolve an undefined symbol to its
mangled symbol. Previously, we did that on calling side of
findMangle by explicitly updating SymbolBody.
In this patch, mangled symbols are handled as weak aliases
for undefined symbols.

llvm-svn: 241213

7 files changed, 51 insertions, 75 deletions

diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index c8554e1435aa..0ce60b714a09 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -22,13 +22,13 @@ namespace coff {
 
 using llvm::COFF::WindowsSubsystem;
 using llvm::StringRef;
-struct Symbol;
+class Undefined;
 
 // Represents an /export option.
 struct Export {
   StringRef Name;
   StringRef ExtName;
-  Symbol *Sym = nullptr;
+  Undefined *Sym = nullptr;
   uint16_t Ordinal = 0;
   bool Noname = false;
   bool Data = false;
@@ -42,7 +42,7 @@ struct Configuration {
   llvm::COFF::MachineTypes MachineType = llvm::COFF::IMAGE_FILE_MACHINE_AMD64;
   bool Verbose = false;
   WindowsSubsystem Subsystem = llvm::COFF::IMAGE_SUBSYSTEM_UNKNOWN;
-  StringRef EntryName;
+  Undefined *Entry = nullptr;
   bool NoEntry = false;
   std::string OutputFile;
   bool DoGC = true;
diff --git a/lld/COFF/DLL.cpp b/lld/COFF/DLL.cpp
index 5dda97288dc7..479e643a80e5 100644
--- a/lld/COFF/DLL.cpp
+++ b/lld/COFF/DLL.cpp
@@ -420,7 +420,7 @@ public:
 
   void writeTo(uint8_t *Buf) override {
     for (Export &E : Config->Exports) {
-      auto *D = cast<Defined>(E.Sym->Body);
+      auto *D = cast<Defined>(E.Sym->getReplacement());
       write32le(Buf + FileOff + E.Ordinal * 4, D->getRVA());
     }
   }
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 66009a203b31..b0dbf8b19624 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -202,9 +202,10 @@ void LinkerDriver::addLibSearchPaths() {
   }
 }
 
-void LinkerDriver::addUndefined(StringRef Sym) {
-  Symtab.addUndefined(Sym);
-  Config->GCRoots.insert(Sym);
+Undefined *LinkerDriver::addUndefined(StringRef Name) {
+  Undefined *U = Symtab.addUndefined(Name);
+  Config->GCRoots.insert(Name);
+  return U;
 }
 
 // Windows specific -- find default entry point name.
@@ -286,10 +287,8 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     Config->Force = true;
 
   // Handle /entry
-  if (auto *Arg = Args.getLastArg(OPT_entry)) {
-    Config->EntryName = Arg->getValue();
-    addUndefined(Config->EntryName);
-  }
+  if (auto *Arg = Args.getLastArg(OPT_entry))
+    Config->Entry = addUndefined(Arg->getValue());
 
   // Handle /noentry
   if (Args.hasArg(OPT_noentry)) {
@@ -304,10 +303,8 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
   if (Args.hasArg(OPT_dll)) {
     Config->DLL = true;
     Config->ManifestID = 2;
-    if (Config->EntryName.empty() && !Config->NoEntry) {
-      Config->EntryName = "_DllMainCRTStartup";
-      addUndefined("_DllMainCRTStartup");
-    }
+    if (Config->Entry == nullptr && !Config->NoEntry)
+      Config->Entry = addUndefined("_DllMainCRTStartup");
   }
 
   // Handle /fixed
@@ -544,14 +541,13 @@ 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.
-  if (Config->EntryName.empty() && !Config->NoEntry) {
+  if (Config->Entry == nullptr && !Config->NoEntry) {
     StringRef S = findDefaultEntry();
     if (S.empty()) {
       llvm::errs() << "entry point must be defined\n";
       return false;
     }
-    Config->EntryName = S;
-    addUndefined(S);
+    Config->Entry = addUndefined(S);
   }
 
   // Resolve auxiliary symbols until converge.
@@ -561,9 +557,16 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
   for (;;) {
     size_t Ver = Symtab.getVersion();
 
+    // Windows specific -- if entry point is not found,
+    // search for its mangled names.
+    if (Config->Entry)
+      Symtab.mangleMaybe(Config->Entry);
+
     // Windows specific -- Make sure we resolve all dllexported symbols.
-    for (Export &E : Config->Exports)
-      addUndefined(E.Name);
+    for (Export &E : Config->Exports) {
+      E.Sym = addUndefined(E.Name);
+      Symtab.mangleMaybe(E.Sym);
+    }
 
     // Add weak aliases. Weak aliases is a mechanism to give remaining
     // undefined symbols final chance to be resolved successfully.
@@ -585,30 +588,6 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
       break;
   }
 
-  // Windows specific -- if entry point is not found,
-  // search for its mangled names.
-  if (!Config->EntryName.empty() && !Symtab.find(Config->EntryName)) {
-    StringRef Name;
-    Symbol *Sym;
-    std::tie(Name, Sym) = Symtab.findMangled(Config->EntryName);
-    if (Sym)
-      Symtab.rename(Config->EntryName, Name);
-  }
-
-  // Windows specific -- resolve dllexported symbols.
-  for (Export &E : Config->Exports) {
-    StringRef Name;
-    Symbol *Sym;
-    std::tie(Name, Sym) = Symtab.findMangled(E.Name);
-    if (!Sym) {
-      llvm::errs() << "exported symbol is not defined: " << E.Name << "\n";
-      return false;
-    }
-    if (E.Name != Name)
-      Symtab.rename(E.Name, Name);
-    E.Sym = Sym;
-  }
-
   // Make sure we have resolved all symbols.
   if (Symtab.reportRemainingUndefines())
     return false;
diff --git a/lld/COFF/Driver.h b/lld/COFF/Driver.h
index c3e2186d5318..3bb5682bd1ac 100644
--- a/lld/COFF/Driver.h
+++ b/lld/COFF/Driver.h
@@ -92,7 +92,7 @@ private:
   std::vector<StringRef> SearchPaths;
   std::set<std::string> VisitedFiles;
 
-  void addUndefined(StringRef Sym);
+  Undefined *addUndefined(StringRef Sym);
 
   // Windows specific -- "main" is not the only main function in Windows.
   // You can choose one from these four -- {w,}{WinMain,main}.
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
index 51cf3de1d1d7..bc0d1e4ff9ee 100644
--- a/lld/COFF/SymbolTable.cpp
+++ b/lld/COFF/SymbolTable.cpp
@@ -25,8 +25,6 @@ namespace coff {
 
 SymbolTable::SymbolTable() {
   addSymbol(new (Alloc) DefinedAbsolute("__ImageBase", Config->ImageBase));
-  if (!Config->EntryName.empty())
-    addSymbol(new (Alloc) Undefined(Config->EntryName));
 }
 
 void SymbolTable::addFile(std::unique_ptr<InputFile> FileP) {
@@ -123,9 +121,11 @@ bool SymbolTable::reportRemainingUndefines() {
       continue;
     StringRef Name = Undef->getName();
     // The weak alias may have been resovled, so check for that.
-    if (auto *D = dyn_cast_or_null<Defined>(Undef->WeakAlias)) {
-      Sym->Body = D;
-      continue;
+    if (SymbolBody *Alias = Undef->WeakAlias) {
+      if (auto *D = dyn_cast<Defined>(Alias->getReplacement())) {
+        Sym->Body = D;
+        continue;
+      }
     }
     // If we can resolve a symbol by removing __imp_ prefix, do that.
     // This odd rule is for compatibility with MSVC linker.
@@ -245,34 +245,30 @@ Symbol *SymbolTable::findSymbol(StringRef Name) {
   return It->second;
 }
 
-// Find a given symbol or its mangled symbol.
-std::pair<StringRef, Symbol *> SymbolTable::findMangled(StringRef S) {
-  auto It = Symtab.find(S);
-  if (It != Symtab.end()) {
-    Symbol *Sym = It->second;
-    if (isa<Defined>(Sym->Body))
-      return std::make_pair(S, Sym);
-  }
+void SymbolTable::mangleMaybe(Undefined *U) {
+  if (U->WeakAlias)
+    return;
+  if (!isa<Undefined>(U->getReplacement()))
+    return;
 
   // In Microsoft ABI, a non-member function name is mangled this way.
-  std::string Prefix = ("?" + S + "@@Y").str();
+  std::string Prefix = ("?" + U->getName() + "@@Y").str();
   for (auto I : Symtab) {
     StringRef Name = I.first;
-    Symbol *Sym = I.second;
+    Symbol *New = I.second;
     if (!Name.startswith(Prefix))
       continue;
-    if (auto *B = dyn_cast<Lazy>(Sym->Body)) {
-      addMemberFile(B);
-      run();
-    }
-    if (isa<Defined>(Sym->Body))
-      return std::make_pair(Name, Sym);
+    U->WeakAlias = New->Body;
+    if (auto *L = dyn_cast<Lazy>(New->Body))
+      addMemberFile(L);
+    return;
   }
-  return std::make_pair(S, nullptr);
 }
 
-std::error_code SymbolTable::addUndefined(StringRef Name) {
-  return addSymbol(new (Alloc) Undefined(Name));
+Undefined *SymbolTable::addUndefined(StringRef Name) {
+  auto *U = new (Alloc) Undefined(Name);
+  addSymbol(U);
+  return U;
 }
 
 // Resolve To, and make From an alias to To.
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
index 8688e8e62311..b8f2fd88f1d9 100644
--- a/lld/COFF/SymbolTable.h
+++ b/lld/COFF/SymbolTable.h
@@ -61,10 +61,11 @@ public:
   Defined *find(StringRef Name);
   Symbol *findSymbol(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);
+  // Occasionally we have to resolve an undefined symbol to its
+  // mangled symbol. This function tries to find a mangled name
+  // for U from the symbol table, and if found, set the symbol as
+  // a weak alias for U.
+  void mangleMaybe(Undefined *U);
 
   // Print a layout map to OS.
   void printMap(llvm::raw_ostream &OS);
@@ -82,7 +83,7 @@ public:
   std::vector<ObjectFile *> ObjectFiles;
 
   // Creates an Undefined symbol for a given name.
-  std::error_code addUndefined(StringRef Name);
+  Undefined *addUndefined(StringRef Name);
 
   // Rename From -> To in the symbol table.
   std::error_code rename(StringRef From, StringRef To);
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 7d85e534cee3..0a1e051149df 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -348,7 +348,7 @@ void Writer::writeHeader() {
   PE->SizeOfImage = SizeOfImage;
   PE->SizeOfHeaders = SizeOfHeaders;
   if (!Config->NoEntry) {
-    Defined *Entry = cast<Defined>(Symtab->find(Config->EntryName));
+    Defined *Entry = cast<Defined>(Config->Entry->getReplacement());
     PE->AddressOfEntryPoint = Entry->getRVA();
   }
   PE->SizeOfStackReserve = Config->StackReserve;