[ORC] Add MachO universal binary support to StaticLibraryDefinitionGenerator.

Add a new overload of StaticLibraryDefinitionGenerator::Load that takes a triple
argument and supports loading archives from MachO universal binaries in addition
to regular archives.

The LLI tool is updated to use this overload.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -367,6 +367,14 @@ class StaticLibraryDefinitionGenerator : public JITDylib::DefinitionGenerator {
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Load(ObjectLayer &L, const char *FileName);
 
+  /// Try to create a StaticLibraryDefinitionGenerator from the given path.
+  ///
+  /// This call will succeed if the file at the given path is a static library
+  /// or a MachO universal binary containing a static library that is compatible
+  /// with the given triple. Otherwise it will return an error.
+  static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+  Load(ObjectLayer &L, const char *FileName, const Triple &TT);
+
   /// Try to create a StaticLibrarySearchGenerator from the given memory buffer.
   /// This call will succeed if the buffer contains a valid archive, otherwise
   /// it will return an error.

llvm/include/llvm/Object/MachOUniversal.h

@@ -90,25 +90,14 @@ class MachOUniversalBinary : public Binary {
       else // Parent->getMagic() == MachO::FAT_MAGIC_64
         return Header64.reserved;
     }
+    Triple getTriple() const {
+      return MachOObjectFile::getArchTriple(getCPUType(), getCPUSubType());
+    }
     std::string getArchFlagName() const {
       const char *McpuDefault, *ArchFlag;
-      if (Parent->getMagic() == MachO::FAT_MAGIC) {
-        Triple T =
-            MachOObjectFile::getArchTriple(Header.cputype, Header.cpusubtype,
-                                           &McpuDefault, &ArchFlag);
-      } else { // Parent->getMagic() == MachO::FAT_MAGIC_64
-        Triple T =
-            MachOObjectFile::getArchTriple(Header64.cputype,
-                                           Header64.cpusubtype,
-                                           &McpuDefault, &ArchFlag);
-      }
-      if (ArchFlag) {
-        std::string ArchFlagName(ArchFlag);
-        return ArchFlagName;
-      } else {
-        std::string ArchFlagName("");
-        return ArchFlagName;
-      }
+      MachOObjectFile::getArchTriple(getCPUType(), getCPUSubType(),
+                                     &McpuDefault, &ArchFlag);
+      return ArchFlag ? ArchFlag : std::string();
     }
 
     Expected<std::unique_ptr<MachOObjectFile>> getAsObjectFile() const;

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -13,6 +13,8 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Object/MachOUniversal.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Target/TargetMachine.h"
 
@@ -302,6 +304,51 @@ StaticLibraryDefinitionGenerator::Load(ObjectLayer &L, const char *FileName) {
   return Create(L, std::move(*ArchiveBuffer));
 }
 
+Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+StaticLibraryDefinitionGenerator::Load(ObjectLayer &L, const char *FileName,
+                                       const Triple &TT) {
+  auto B = object::createBinary(FileName);
+  if (!B)
+    return B.takeError();
+
+  // If this is a regular archive then create an instance from it.
+  if (isa<object::Archive>(B->getBinary()))
+    return Create(L, std::move(B->takeBinary().second));
+
+  // If this is a universal binary then search for a slice matching the given
+  // Triple.
+  if (auto *UB = cast<object::MachOUniversalBinary>(B->getBinary())) {
+    for (const auto &Obj : UB->objects()) {
+      auto ObjTT = Obj.getTriple();
+      if (ObjTT.getArch() == TT.getArch() &&
+          ObjTT.getSubArch() == TT.getSubArch() &&
+          ObjTT.getVendor() == TT.getVendor()) {
+        // We found a match. Create an instance from a buffer covering this
+        // slice.
+        auto SliceBuffer = MemoryBuffer::getFileSlice(FileName, Obj.getSize(),
+                                                      Obj.getOffset());
+        if (!SliceBuffer)
+          return make_error<StringError>(
+              Twine("Could not create buffer for ") + TT.str() + " slice of " +
+                  FileName + ": [ " + formatv("{0:x}", Obj.getOffset()) +
+                  " .. " + formatv("{0:x}", Obj.getOffset() + Obj.getSize()) +
+                  ": " + SliceBuffer.getError().message(),
+              SliceBuffer.getError());
+        return Create(L, std::move(*SliceBuffer));
+      }
+    }
+
+    return make_error<StringError>(Twine("Universal binary ") + FileName +
+                                       " does not contain a slice for " +
+                                       TT.str(),
+                                   inconvertibleErrorCode());
+  }
+
+  return make_error<StringError>(Twine("Unrecognized file type for ") +
+                                     FileName,
+                                 inconvertibleErrorCode());
+}
+
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
 StaticLibraryDefinitionGenerator::Create(
     ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer) {

llvm/test/ExecutionEngine/OrcLazy/macho-universal-static-library-support.ll

@@ -0,0 +1,26 @@
+; REQUIRES: x86_64-apple, x86-registered-target, arm-registered-target
+;
+; RUN: rm -rf %t && mkdir -p %t
+; RUN: llc -filetype=obj -mtriple=x86_64-apple-macosx -o %t/foo.x86_64.o \
+; RUN:   %p/Inputs/foo-return-i32-0.ll
+; RUN: llvm-ar r %t/foo.x86_64.a %t/foo.x86_64.o
+; RUN: llc -filetype=obj -mtriple=arm-apple-ios -o %t/foo.arm.o \
+; RUN:   %p/Inputs/foo-return-i32-0.ll
+; RUN: llvm-ar r %t/foo.arm.a %t/foo.arm.o
+; RUN: llvm-lipo -create %t/foo.x86_64.a %t/foo.arm.a -output %t/foo.a
+; RUN: lli -jit-kind=orc-lazy -extra-archive %t/foo.a %s
+;
+; Check that MachO universal binaries containing archives work.
+; This test compiles two copies of a simple int foo() function that returns
+; zero, one copy for x86_64 and one for arm. It then puts each of these in an
+; archive and combines these two archives into a macho universal binary.
+; Finally we execute a main function that references foo to ensure that the
+; x86-64 copy is correctly found and linked.
+
+declare i32 @foo()
+
+define i32 @main(i32 %argc, i8** nocapture readnone %argv) {
+entry:
+  %0 = call i32 @foo()
+  ret i32 %0
+}

llvm/tools/lli/lli.cpp

@@ -966,7 +966,7 @@ int runOrcLazyJIT(const char *ProgName) {
       auto JDItr = std::prev(IdxToDylib.lower_bound(EAIdx));
       auto &JD = *JDItr->second;
       JD.addGenerator(ExitOnErr(orc::StaticLibraryDefinitionGenerator::Load(
-          J->getObjLinkingLayer(), EAItr->c_str())));
+          J->getObjLinkingLayer(), EAItr->c_str(), *TT)));
     }
   }
 

llvm/utils/lit/lit/llvm/config.py

@@ -87,6 +87,7 @@ def __init__(self, lit_config, config):
 
         if target_triple:
             if re.match(r'^x86_64.*-apple', target_triple):
+                features.add('x86_64-apple')
                 host_cxx = getattr(config, 'host_cxx', None)
                 if 'address' in sanitizers and self.get_clang_has_lsan(host_cxx, target_triple):
                     self.with_environment(