COFF: Infer machine type earlier than before.

Previously, we infer machine type at the very end of linking after
all symbols are resolved. That's actually too late because machine
type affects how we mangle symbols (whether or not we need to
add "_").

For example, /entry:foo adds "_foo" to the symbol table if x86 but
"foo" if x64.

This patch moves the code to infer machine type, so that machine
type is inferred based on input files given via the command line
(but not based on .directives files).

llvm-svn: 241843

5 files changed, 90 insertions, 49 deletions

diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 522bc9bf0f9f..01280b68f4d7 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -209,6 +209,7 @@ Undefined *LinkerDriver::addUndefined(StringRef Name) {
 
 // Symbol names are mangled by appending "_" prefix on x86.
 StringRef LinkerDriver::mangle(StringRef Sym) {
+  assert(Config->MachineType != IMAGE_FILE_MACHINE_UNKNOWN);
   if (Config->MachineType == IMAGE_FILE_MACHINE_I386)
     return Alloc.save("_" + Sym);
   return Sym;
@@ -293,8 +294,9 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     Config->Force = true;
 
   // Handle /entry
+  StringRef Entry;
   if (auto *Arg = Args.getLastArg(OPT_entry))
-    Config->Entry = addUndefined(mangle(Arg->getValue()));
+    Entry = Arg->getValue();
 
   // Handle /debug
   if (Args.hasArg(OPT_debug))
@@ -314,8 +316,8 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     Config->DLL = true;
     Config->ImageBase = 0x180000000U;
     Config->ManifestID = 2;
-    if (Config->Entry == nullptr && !Config->NoEntry)
-      Config->Entry = addUndefined("_DllMainCRTStartup");
+    if (Entry.empty() && !Config->NoEntry)
+      Entry = "_DllMainCRTStartup";
   }
 
   // Handle /fixed
@@ -497,21 +499,21 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
 
   // Create a list of input files. Files can be given as arguments
   // for /defaultlib option.
-  std::vector<StringRef> InputPaths;
-  std::vector<MemoryBufferRef> Inputs;
+  std::vector<StringRef> Paths;
+  std::vector<MemoryBufferRef> MBs;
   for (auto *Arg : Args.filtered(OPT_INPUT))
     if (Optional<StringRef> Path = findFile(Arg->getValue()))
-      InputPaths.push_back(*Path);
+      Paths.push_back(*Path);
   for (auto *Arg : Args.filtered(OPT_defaultlib))
     if (Optional<StringRef> Path = findLib(Arg->getValue()))
-      InputPaths.push_back(*Path);
-  for (StringRef Path : InputPaths) {
+      Paths.push_back(*Path);
+  for (StringRef Path : Paths) {
     ErrorOr<MemoryBufferRef> MBOrErr = openFile(Path);
     if (auto EC = MBOrErr.getError()) {
       llvm::errs() << "cannot open " << Path << ": " << EC.message() << "\n";
       return false;
     }
-    Inputs.push_back(MBOrErr.get());
+    MBs.push_back(MBOrErr.get());
   }
 
   // Windows specific -- Create a resource file containing a manifest file.
@@ -520,42 +522,72 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     if (MBOrErr.getError())
       return false;
     std::unique_ptr<MemoryBuffer> MB = std::move(MBOrErr.get());
-    Inputs.push_back(MB->getMemBufferRef());
+    MBs.push_back(MB->getMemBufferRef());
     OwningMBs.push_back(std::move(MB)); // take ownership
   }
 
   // Windows specific -- Input files can be Windows resource files (.res files).
   // We invoke cvtres.exe to convert resource files to a regular COFF file
   // then link the result file normally.
+  std::vector<MemoryBufferRef> Resources;
   auto NotResource = [](MemoryBufferRef MB) {
     return identify_magic(MB.getBuffer()) != file_magic::windows_resource;
   };
-  auto It = std::stable_partition(Inputs.begin(), Inputs.end(), NotResource);
-  if (It != Inputs.end()) {
-    std::vector<MemoryBufferRef> Files(It, Inputs.end());
-    auto MBOrErr = convertResToCOFF(Files);
-    if (MBOrErr.getError())
-      return false;
-    std::unique_ptr<MemoryBuffer> MB = std::move(MBOrErr.get());
-    Inputs.erase(It, Inputs.end());
-    Inputs.push_back(MB->getMemBufferRef());
-    OwningMBs.push_back(std::move(MB)); // take ownership
+  auto It = std::stable_partition(MBs.begin(), MBs.end(), NotResource);
+  if (It != MBs.end()) {
+    Resources.insert(Resources.end(), It, MBs.end());
+    MBs.erase(It, MBs.end());
   }
 
-  Symtab.addAbsolute(mangle("__ImageBase"), Config->ImageBase);
-
   // Read all input files given via the command line. Note that step()
   // doesn't read files that are specified by directive sections.
-  for (MemoryBufferRef MB : Inputs)
+  for (MemoryBufferRef MB : MBs)
     Symtab.addFile(createFile(MB));
   if (auto EC = Symtab.step()) {
     llvm::errs() << EC.message() << "\n";
     return false;
   }
 
+  // Determine machine type and check if all object files are
+  // for the same CPU type. Note that this needs to be done before
+  // any call to mangle().
+  for (std::unique_ptr<InputFile> &File : Symtab.getFiles()) {
+    MachineTypes MT = File->getMachineType();
+    if (MT == IMAGE_FILE_MACHINE_UNKNOWN)
+      continue;
+    if (Config->MachineType == IMAGE_FILE_MACHINE_UNKNOWN) {
+      Config->MachineType = MT;
+      continue;
+    }
+    if (Config->MachineType != MT) {
+      llvm::errs() << File->getShortName() << ": machine type "
+                   << machineTypeToStr(MT) << " conflicts with "
+                   << machineTypeToStr(Config->MachineType) << "\n";
+      return false;
+    }
+  }
+  if (Config->MachineType == IMAGE_FILE_MACHINE_UNKNOWN) {
+    llvm::errs() << "warning: /machine is not specified. x64 is assumed.\n";
+    Config->MachineType = IMAGE_FILE_MACHINE_AMD64;
+  }
+
+  // Windows specific -- Convert Windows resource files to a COFF file.
+  if (!Resources.empty()) {
+    auto MBOrErr = convertResToCOFF(Resources);
+    if (MBOrErr.getError())
+      return false;
+    std::unique_ptr<MemoryBuffer> MB = std::move(MBOrErr.get());
+    Symtab.addFile(createFile(MB->getMemBufferRef()));
+    OwningMBs.push_back(std::move(MB)); // take ownership
+  }
+
+  if (!Entry.empty())
+    Config->Entry = addUndefined(mangle(Entry));
+  Symtab.addAbsolute(mangle("__ImageBase"), Config->ImageBase);
+
   // Windows specific -- If entry point name is not given, we need to
   // infer that from user-defined entry name.
-  if (Config->Entry == nullptr && !Config->NoEntry) {
+  if (Entry.empty() && !Config->NoEntry) {
     StringRef S = findDefaultEntry();
     if (S.empty()) {
       llvm::errs() << "entry point must be defined\n";
@@ -566,7 +598,7 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
       llvm::outs() << "Entry name inferred: " << S << "\n";
   }
 
-  // Read as much files as we can.
+  // Read as much files as we can from directives sections.
   if (auto EC = Symtab.run()) {
     llvm::errs() << EC.message() << "\n";
     return false;
@@ -630,28 +662,6 @@ bool LinkerDriver::link(llvm::ArrayRef<const char *> ArgsArr) {
     }
   }
 
-  // Check if all object files are for the same CPU type and
-  // compatible with /machine option (if given).
-  for (ObjectFile *File : Symtab.ObjectFiles) {
-    auto MT = static_cast<MachineTypes>(File->getCOFFObj()->getMachine());
-    if (MT == IMAGE_FILE_MACHINE_UNKNOWN)
-      continue;
-    if (Config->MachineType == IMAGE_FILE_MACHINE_UNKNOWN) {
-      Config->MachineType = MT;
-      continue;
-    }
-    if (Config->MachineType != MT) {
-      llvm::errs() << File->getShortName() << ": machine type "
-                   << machineTypeToStr(MT) << " conflicts with "
-                   << machineTypeToStr(Config->MachineType) << "\n";
-      return false;
-    }
-  }
-  if (Config->MachineType == IMAGE_FILE_MACHINE_UNKNOWN) {
-    llvm::errs() << "machine type must be specified\n";
-    return false;
-  }
-
   // Windows specific -- when we are creating a .dll file, we also
   // need to create a .lib file.
   if (!Config->Exports.empty())
diff --git a/lld/COFF/InputFiles.cpp b/lld/COFF/InputFiles.cpp
index a0211623f62d..cbc239388863 100644
--- a/lld/COFF/InputFiles.cpp
+++ b/lld/COFF/InputFiles.cpp
@@ -24,8 +24,9 @@ using namespace llvm::COFF;
 using namespace llvm::object;
 using namespace llvm::support::endian;
 using llvm::RoundUpToAlignment;
-using llvm::sys::fs::identify_magic;
+using llvm::Triple;
 using llvm::sys::fs::file_magic;
+using llvm::sys::fs::identify_magic;
 
 namespace lld {
 namespace coff {
@@ -242,6 +243,12 @@ Defined *ObjectFile::createDefined(COFFSymbolRef Sym, const void *AuxP,
   return B;
 }
 
+MachineTypes ObjectFile::getMachineType() {
+  if (COFFObj)
+    return static_cast<MachineTypes>(COFFObj->getMachine());
+  return IMAGE_FILE_MACHINE_UNKNOWN;
+}
+
 std::error_code ImportFile::parse() {
   const char *Buf = MB.getBufferStart();
   const char *End = MB.getBufferEnd();
@@ -319,5 +326,20 @@ std::error_code BitcodeFile::parse() {
   return std::error_code();
 }
 
+MachineTypes BitcodeFile::getMachineType() {
+  if (!M)
+    return IMAGE_FILE_MACHINE_UNKNOWN;
+  switch (Triple(M->getTargetTriple()).getArch()) {
+  case Triple::x86_64:
+    return IMAGE_FILE_MACHINE_AMD64;
+  case Triple::x86:
+    return IMAGE_FILE_MACHINE_I386;
+  case Triple::arm:
+    return IMAGE_FILE_MACHINE_ARMNT;
+  default:
+    return IMAGE_FILE_MACHINE_UNKNOWN;
+  }
+}
+
 } // namespace coff
 } // namespace lld
diff --git a/lld/COFF/InputFiles.h b/lld/COFF/InputFiles.h
index 98a162482b90..0d768c52ea26 100644
--- a/lld/COFF/InputFiles.h
+++ b/lld/COFF/InputFiles.h
@@ -24,6 +24,8 @@ namespace lld {
 namespace coff {
 
 using llvm::LTOModule;
+using llvm::COFF::IMAGE_FILE_MACHINE_UNKNOWN;
+using llvm::COFF::MachineTypes;
 using llvm::object::Archive;
 using llvm::object::COFFObjectFile;
 using llvm::object::COFFSymbolRef;
@@ -51,6 +53,9 @@ public:
   // file is broken.
   virtual std::error_code parse() = 0;
 
+  // Returns the CPU type this file was compiled to.
+  virtual MachineTypes getMachineType() { return IMAGE_FILE_MACHINE_UNKNOWN; }
+
   // Returns a short, human-friendly filename. If this is a member of
   // an archive file, a returned value includes parent's filename.
   // Used for logging or debugging.
@@ -112,6 +117,7 @@ public:
   explicit ObjectFile(MemoryBufferRef M) : InputFile(ObjectKind, M) {}
   static bool classof(const InputFile *F) { return F->kind() == ObjectKind; }
   std::error_code parse() override;
+  MachineTypes getMachineType() override;
   std::vector<Chunk *> &getChunks() { return Chunks; }
   std::vector<SymbolBody *> &getSymbols() override { return SymbolBodies; }
 
@@ -181,6 +187,7 @@ public:
   explicit BitcodeFile(MemoryBufferRef M) : InputFile(BitcodeKind, M) {}
   static bool classof(const InputFile *F) { return F->kind() == BitcodeKind; }
   std::vector<SymbolBody *> &getSymbols() override { return SymbolBodies; }
+  MachineTypes getMachineType() override;
 
   LTOModule *getModule() const { return M.get(); }
   LTOModule *releaseModule() { return M.release(); }
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
index 44c7ce4a7b45..8fe7bbb5e6aa 100644
--- a/lld/COFF/SymbolTable.h
+++ b/lld/COFF/SymbolTable.h
@@ -42,6 +42,7 @@ struct Symbol;
 class SymbolTable {
 public:
   void addFile(std::unique_ptr<InputFile> File);
+  std::vector<std::unique_ptr<InputFile>> &getFiles() { return Files; }
   std::error_code step();
   std::error_code run();
   bool queueEmpty();
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 23ce8485581b..ca50e8bdce51 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -592,7 +592,8 @@ OutputSection *Writer::createSection(StringRef Name) {
 // Dest is .reloc section. Add contents to that section.
 void Writer::addBaserels(OutputSection *Dest) {
   std::vector<uint32_t> V;
-  Defined *ImageBase = cast<Defined>(Symtab->find("__ImageBase")->Body);
+  StringRef Name = Config->is64() ? "__ImageBase" : "___ImageBase";
+  Defined *ImageBase = cast<Defined>(Symtab->find(Name)->Body);
   for (OutputSection *Sec : OutputSections) {
     if (Sec == Dest)
       continue;