Move machine type auto-detection to SymbolTable.

[pcc]

Summary

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

[ruiu]

LGTM with following change.

I first thought that it's invalid if all files are archive, but then realized that that's valid.
An object file defining an entry point function can be in an archive.

COFF/SymbolTable.cpp

174–184

I'd make this a non-member function in Writer.cpp which takes vector<unique_ptr<ObjectFile>>& as a parameter, as this function is not that related to symbol resolution. It feels that it belongs Writer. Maybe you need to expose ObjectFiles, but that's fine.

And in the Writer, you can return Config->MachineType as a default value instead of UNKNOWN.

[pcc]

COFF/SymbolTable.cpp

174–184

Config->MachineType overrides the detected value, doesn't it? That's what the existing logic seems to implement.

COFF/DriverUtils.cpp

for (StringRef Dir : splitPathList(*Env)) {
     SmallString<128> Path = Dir;
     llvm::sys::path::append(Path, Filename);
     if (llvm::sys::fs::exists(Path.str()))
       return Path.str();
   }
   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)) {
     StringRef S(Arg->getValue());
     MachineTypes MT = StringSwitch<MachineTypes>(S.lower())
                           .Case("arm", IMAGE_FILE_MACHINE_ARMNT)
                           .Case("x64", IMAGE_FILE_MACHINE_AMD64)
                           .Case("x86", IMAGE_FILE_MACHINE_I386)
                           .Default(IMAGE_FILE_MACHINE_UNKNOWN);
     if (MT == IMAGE_FILE_MACHINE_UNKNOWN)
       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;
 }
 
 // Parses a string in the form of "<integer>[,<integer>]".
 std::error_code parseNumbers(StringRef Arg, uint64_t *Addr, uint64_t *Size) {
   StringRef S1, S2;
   std::tie(S1, S2) = Arg.split(',');
   if (S1.getAsInteger(0, *Addr))

COFF/SymbolTable.h

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);
 
   // Dump contents of the symbol table to stderr.
   void dump();
 
+  // Try to infer the machine type from the object files known to the symbol
+  // table.
+  llvm::COFF::MachineTypes getInferredMachineType();
+
   // The writer needs to handle DLL import libraries specially in
   // order to create the import descriptor table.
   std::vector<std::unique_ptr<ImportFile>> ImportFiles;
 
 private:
   std::error_code addObject(ObjectFile *File);
   std::error_code addArchive(ArchiveFile *File);
   std::error_code addImport(ImportFile *File);

COFF/SymbolTable.cpp

void SymbolTable::dump() {
   for (auto &P : Symtab) {
     Symbol *Ref = P.second;
     if (auto *Body = dyn_cast<Defined>(Ref->Body))
       llvm::dbgs() << Twine::utohexstr(Config->ImageBase + Body->getRVA())
                    << " " << Body->getName() << "\n";
   }
 }
 
+llvm::COFF::MachineTypes SymbolTable::getInferredMachineType() {
+  for (std::unique_ptr<ObjectFile> &File : ObjectFiles) {
+    // Try to infer machine type from the magic byte of the object file.
+    auto MT =
+        static_cast<llvm::COFF::MachineTypes>(File->getCOFFObj()->getMachine());
+    if (MT != llvm::COFF::IMAGE_FILE_MACHINE_UNKNOWN)
+      return MT;
+  }
+  return llvm::COFF::IMAGE_FILE_MACHINE_UNKNOWN;
+}
+
 } // namespace coff
 } // namespace lld

COFF/Writer.cpp

void Writer::writeHeader() {
   DOS->Magic[1] = 'Z';
   DOS->AddressOfRelocationTable = sizeof(dos_header);
   DOS->AddressOfNewExeHeader = DOSStubSize;
 
   // Write PE magic
   memcpy(Buf, PEMagic, sizeof(PEMagic));
   Buf += sizeof(PEMagic);
 
+  MachineTypes MachineType = Config->MachineType;
+  if (MachineType == IMAGE_FILE_MACHINE_UNKNOWN)
+    MachineType = Symtab->getInferredMachineType();
+
   // 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 |
        IMAGE_FILE_LARGE_ADDRESS_AWARE);
   COFF->SizeOfOptionalHeader =
       sizeof(pe32plus_header) + sizeof(data_directory) * NumberfOfDataDirectory;
 
   // Write PE header

test/COFF/machine.test

 # 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
 
 # RUN: lld -flavor link2 /entry:main /machine:x86 /out:%t.exe \
 # RUN:   %p/Inputs/ret42.obj
 # RUN: llvm-readobj -file-headers %t.exe | FileCheck -check-prefix=I386 %s
 
 I386: Machine: IMAGE_FILE_MACHINE_I386