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 ↗

(On Diff #26805)

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 ↗

(On Diff #26805)

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

lld/trunk/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))

lld/trunk/COFF/SymbolTable.h

public:
 
   // Dump contents of the symbol table to stderr.
   void dump();
 
   // The writer needs to handle DLL import libraries specially in
   // 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);
   std::error_code addImport(ImportFile *File);
 
   std::error_code resolve(SymbolBody *Body);
   std::error_code addMemberFile(Lazy *Body);
   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;
   StringAllocator StringAlloc;
 };
 
 } // namespace pecoff
 } // namespace lld
 
 #endif

lld/trunk/COFF/Writer.cpp

for (std::unique_ptr<OutputSection> &Sec : OutputSections) {
     RVA += RoundUpToAlignment(Sec->getVirtualSize(), PageSize);
     FileOff += RoundUpToAlignment(Sec->getRawSize(), FileAlignment);
   }
   SizeOfImage = SizeOfHeaders + RoundUpToAlignment(RVA - 0x1000, PageSize);
   FileSize = SizeOfHeaders +
              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();
   auto *DOS = reinterpret_cast<dos_header *>(Buf);
   Buf += DOSStubSize;
   DOS->Magic[0] = 'M';
   DOS->Magic[1] = 'Z';
   DOS->AddressOfRelocationTable = sizeof(dos_header);
   DOS->AddressOfNewExeHeader = DOSStubSize;
 
   // Write PE magic
   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 |
        IMAGE_FILE_LARGE_ADDRESS_AWARE);
   COFF->SizeOfOptionalHeader =
       sizeof(pe32plus_header) + sizeof(data_directory) * NumberfOfDataDirectory;
 
   // Write PE header

lld/trunk/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