COFF: Move machine type auto-detection to SymbolTable.

The new mechanism is less code, and fixes the case where all inputs
are archives.

Differential Revision: http://reviews.llvm.org/D10136

llvm-svn: 238618

lld/COFF/DriverUtils.cpp

@@ -87,24 +87,6 @@ std::string findFile(StringRef Filename) {
   return Filename;
 }
 
-// Peeks at the file header to get architecture (e.g. i386 or AMD64).
-// Returns "unknown" if it's not a valid object file.
-static MachineTypes getFileMachineType(StringRef Path) {
-  file_magic Magic;
-  if (identify_magic(Path, Magic))
-    return IMAGE_FILE_MACHINE_UNKNOWN;
-  if (Magic != file_magic::coff_object)
-    return IMAGE_FILE_MACHINE_UNKNOWN;
-  ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr = MemoryBuffer::getFile(Path);
-  if (BufOrErr.getError())
-    return IMAGE_FILE_MACHINE_UNKNOWN;
-  std::error_code EC;
-  llvm::object::COFFObjectFile Obj(BufOrErr.get()->getMemBufferRef(), EC);
-  if (EC)
-    return IMAGE_FILE_MACHINE_UNKNOWN;
-  return static_cast<MachineTypes>(Obj.getMachine());
-}
-
 // Returns /machine's value.
 ErrorOr<MachineTypes> getMachineType(llvm::opt::InputArgList *Args) {
   if (auto *Arg = Args->getLastArg(OPT_machine)) {
@@ -118,13 +100,6 @@ ErrorOr<MachineTypes> getMachineType(llvm::opt::InputArgList *Args) {
       return make_dynamic_error_code("unknown /machine argument" + S);
     return MT;
   }
-  // If /machine option is missing, we need to take a look at
-  // the magic byte of the first object file to infer machine type.
-  for (auto *Arg : Args->filtered(OPT_INPUT)) {
-    MachineTypes MT = getFileMachineType(Arg->getValue());
-    if (MT != IMAGE_FILE_MACHINE_UNKNOWN)
-      return MT;
-  }
   return IMAGE_FILE_MACHINE_UNKNOWN;
 }
 

lld/COFF/SymbolTable.h

@@ -53,6 +53,9 @@ class SymbolTable {
   // order to create the import descriptor table.
   std::vector<std::unique_ptr<ImportFile>> ImportFiles;
 
+  // The writer needs to infer the machine type from the object files.
+  std::vector<std::unique_ptr<ObjectFile>> ObjectFiles;
+
 private:
   std::error_code addObject(ObjectFile *File);
   std::error_code addArchive(ArchiveFile *File);
@@ -63,7 +66,6 @@ class SymbolTable {
   void addInitialSymbol(SymbolBody *Body);
 
   std::unordered_map<StringRef, Symbol *> Symtab;
-  std::vector<std::unique_ptr<ObjectFile>> ObjectFiles;
   std::vector<std::unique_ptr<ArchiveFile>> ArchiveFiles;
   std::vector<std::unique_ptr<SymbolBody>> OwnedSymbols;
   llvm::BumpPtrAllocator Alloc;

lld/COFF/Writer.cpp

@@ -227,6 +227,17 @@ void Writer::assignAddresses() {
              RoundUpToAlignment(FileOff - SizeOfHeaders, FileAlignment);
 }
 
+static MachineTypes
+inferMachineType(std::vector<std::unique_ptr<ObjectFile>> &ObjectFiles) {
+  for (std::unique_ptr<ObjectFile> &File : ObjectFiles) {
+    // Try to infer machine type from the magic byte of the object file.
+    auto MT = static_cast<MachineTypes>(File->getCOFFObj()->getMachine());
+    if (MT != IMAGE_FILE_MACHINE_UNKNOWN)
+      return MT;
+  }
+  return IMAGE_FILE_MACHINE_UNKNOWN;
+}
+
 void Writer::writeHeader() {
   // Write DOS stub
   uint8_t *Buf = Buffer->getBufferStart();
@@ -241,10 +252,15 @@ void Writer::writeHeader() {
   memcpy(Buf, PEMagic, sizeof(PEMagic));
   Buf += sizeof(PEMagic);
 
+  // Determine machine type, infer if needed. TODO: diagnose conflicts.
+  MachineTypes MachineType = Config->MachineType;
+  if (MachineType == IMAGE_FILE_MACHINE_UNKNOWN)
+    MachineType = inferMachineType(Symtab->ObjectFiles);
+
   // Write COFF header
   auto *COFF = reinterpret_cast<coff_file_header *>(Buf);
   Buf += sizeof(*COFF);
-  COFF->Machine = Config->MachineType;
+  COFF->Machine = MachineType;
   COFF->NumberOfSections = OutputSections.size();
   COFF->Characteristics =
       (IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_RELOCS_STRIPPED |

lld/test/COFF/machine.test

@@ -1,6 +1,15 @@
 # RUN: lld -flavor link2 /entry:main /machine:x64 /out:%t.exe \
 # RUN:   %p/Inputs/ret42.obj
 # RUN: llvm-readobj -file-headers %t.exe | FileCheck -check-prefix=AMD64 %s
+# RUN: lld -flavor link2 /entry:main /machine:x64 /out:%t.exe \
+# RUN:   %p/Inputs/ret42.lib
+# RUN: llvm-readobj -file-headers %t.exe | FileCheck -check-prefix=AMD64 %s
+# RUN: lld -flavor link2 /entry:main /out:%t.exe \
+# RUN:   %p/Inputs/ret42.obj
+# RUN: llvm-readobj -file-headers %t.exe | FileCheck -check-prefix=AMD64 %s
+# RUN: lld -flavor link2 /entry:main /out:%t.exe \
+# RUN:   %p/Inputs/ret42.lib
+# RUN: llvm-readobj -file-headers %t.exe | FileCheck -check-prefix=AMD64 %s
 
 AMD64: Machine: IMAGE_FILE_MACHINE_AMD64