[lld] r219180 - [PECOFF] Infer machine type from input object file

[Rui Ueyama]

Author: ruiu
Date: Mon Oct  6 20:38:46 2014
New Revision: 219180

URL: http://llvm.org/viewvc/llvm-project?rev=219180&view=rev
Log:
[PECOFF] Infer machine type from input object file

If /machine option is omitted, the linker needs to infer that from
input object files. This patch implements that.

Modified:
    lld/trunk/lib/Driver/WinLinkDriver.cpp
    lld/trunk/test/pecoff/machinetype.test

Modified: lld/trunk/lib/Driver/WinLinkDriver.cpp
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/lib/Driver/WinLinkDriver.cpp?rev=219180&r1=219179&r2=219180&view=diff
==============================================================================
--- lld/trunk/lib/Driver/WinLinkDriver.cpp (original)
+++ lld/trunk/lib/Driver/WinLinkDriver.cpp Mon Oct  6 20:38:46 2014
@@ -23,6 +23,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/Object/COFF.h"
 #include "llvm/Option/Arg.h"
 #include "llvm/Option/Option.h"
 #include "llvm/Support/CommandLine.h"
@@ -361,6 +362,25 @@ static bool parseManifestUAC(StringRef o
   }
 }
 
+// Returns the machine type (e.g. x86) of the given input file.
+// If the file is not COFF, returns false.
+static bool getMachineType(StringRef path, llvm::COFF::MachineTypes &result) {
+  llvm::sys::fs::file_magic fileType;
+  if (llvm::sys::fs::identify_magic(path, fileType))
+    return false;
+  if (fileType != llvm::sys::fs::file_magic::coff_object)
+    return false;
+  ErrorOr<std::unique_ptr<MemoryBuffer>> buf = MemoryBuffer::getFile(path);
+  if (!buf)
+    return false;
+  std::error_code ec;
+  llvm::object::COFFObjectFile obj(buf.get()->getMemBufferRef(), ec);
+  if (ec)
+    return false;
+  result = static_cast<llvm::COFF::MachineTypes>(obj.getMachine());
+  return true;
+}
+
 // Parse /export:entryname[=internalname][, at ordinal[,NONAME]][,DATA][,PRIVATE].
 //
 // MSDN doesn't say anything about /export:foo=bar style option or PRIVATE
@@ -854,9 +874,6 @@ bool WinLinkDriver::parse(int argc, cons
   // Handle /machine before parsing all the other options, as the target machine
   // type affects how to handle other options. For example, x86 needs the
   // leading underscore to mangle symbols, while x64 doesn't need it.
-  //
-  // TODO: If /machine option is missing, we probably should take a look at
-  // the magic byte of the first object file to set machine type.
   if (llvm::opt::Arg *inputArg = parsedArgs->getLastArg(OPT_machine)) {
     StringRef arg = inputArg->getValue();
     llvm::COFF::MachineTypes type = stringToMachineType(arg);
@@ -865,6 +882,25 @@ bool WinLinkDriver::parse(int argc, cons
       return false;
     }
     ctx.setMachineType(type);
+  } else {
+    // If /machine option is missing, we need to take a look at
+    // the magic byte of the first object file to infer machine type.
+    std::vector<StringRef> files;
+    for (auto arg : *parsedArgs)
+      if (arg->getOption().getID() == OPT_INPUT)
+        files.push_back(arg->getValue());
+    if (llvm::opt::Arg *arg = parsedArgs->getLastArg(OPT_DASH_DASH))
+      files.insert(files.end(), arg->getValues().begin(),
+                   arg->getValues().end());
+    for (StringRef path : files) {
+      llvm::COFF::MachineTypes type;
+      if (!getMachineType(path, type))
+        continue;
+      if (type == llvm::COFF::IMAGE_FILE_MACHINE_UNKNOWN)
+        continue;
+      ctx.setMachineType(type);
+      break;
+    }
   }
 
   // Handle /nodefaultlib:<lib>. The same option without argument is handled in

Modified: lld/trunk/test/pecoff/machinetype.test
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/test/pecoff/machinetype.test?rev=219180&r1=219179&r2=219180&view=diff
==============================================================================
--- lld/trunk/test/pecoff/machinetype.test (original)
+++ lld/trunk/test/pecoff/machinetype.test Mon Oct  6 20:38:46 2014
@@ -1,3 +1,13 @@
-# RUN: yaml2obj %p/Inputs/hello.obj.yaml > %t1.obj
-# RUN: yaml2obj %p/Inputs/machine-type-unknown.obj.yaml > %t2.obj
+# RUN: yaml2obj %p/Inputs/machine-type-unknown.obj.yaml > %t1.obj
+# RUN: yaml2obj %p/Inputs/hello.obj.yaml > %t2.obj
+# RUN: yaml2obj %p/Inputs/hello64.obj.yaml > %t3.obj
+
 # RUN: lld -flavor link /out:%t.exe /subsystem:console /force -- %t1.obj %t2.obj
+# RUN: llvm-readobj %t.exe | FileCheck -check-prefix=X86 %s
+
+X86: Arch: i386
+
+# RUN: lld -flavor link /out:%t.exe /subsystem:console /force -- %t1.obj %t3.obj
+# RUN: llvm-readobj %t.exe | FileCheck -check-prefix=X64 %s
+
+X64: Arch: x86_64