[ORC][llvm-jitlink] Add support for forced loading of archive members.

This patch adds support for forced loading of archive members, similar to the
behavior of the -all_load and -ObjC options in ld64. To enable this, the
StaticLibraryDefinitionGenerator class constructors are extended with a
VisitMember callback that is called on each member file in the archive at
generator construction time. This callback can be used to unconditionally add
the member file to a JITDylib at that point.

To test this the llvm-jitlink utility is extended with -all_load (all platforms)
and -ObjC (darwin only) options. Since we can't refer to symbols in the test
objects directly (these would always cause the member to be linked in, even
without the new flags) we instead test side-effects of force loading: execution
of constructors and registration of Objective-C metadata.

rdar://134446111

compiler-rt/test/orc/TestCases/Darwin/Generic/Inputs/EmptyClassFoo.m

@@ -0,0 +1,7 @@
+#include <Foundation/Foundation.h>
+
+@interface Foo : NSObject
+@end
+
+@implementation Foo
+@end

compiler-rt/test/orc/TestCases/Darwin/Generic/llvm-jitlink-force-link-objc.m

@@ -0,0 +1,14 @@
+// RUN: rm -rf %t && mkdir -p %t
+// RUN: %clang -c -o %t/EmptyClassFoo.o  %S/Inputs/EmptyClassFoo.m
+// RUN: ar r %t/libFooClass.a %t/EmptyClassFoo.o
+// RUN: %clang -c -o %t/force-objc.o %s
+// RUN: %llvm_jitlink -ObjC %t/force-objc.o -L%t -lFooClass
+//
+// REQUIRES: system-darwin && host-arch-compatible
+
+id objc_getClass(const char *name);
+
+int main(int argc, char *argv[]) {
+  // Return succeess if we find Foo, error otherwise.
+  return objc_getClass("Foo") ? 0 : 1;
+}

compiler-rt/test/orc/TestCases/Linux/Generic/Inputs/SetGlobalIntXInConstructor.cpp

@@ -0,0 +1,12 @@
+extern "C" int x;
+
+namespace {
+
+struct Init {
+public:
+  Init() { x = 1; }
+};
+
+Init SetX;
+
+} // namespace

compiler-rt/test/orc/TestCases/Linux/Generic/llvm-jitlink-all-load.c

@@ -0,0 +1,14 @@
+// Check that the -all_load flag to llvm-jitlink causes all objects from
+// archives to be loaded, regardless of whether or not they're referenced.
+//
+// RUN: rm -rf %t && mkdir -p %t
+// RUN: %clangxx -c -o %t/SetX.o  %S/Inputs/SetGlobalIntXInConstructor.cpp
+// RUN: ar r %t/libSetX.a %t/SetX.o
+// RUN: %clang -c -o %t/all_load.o %s
+// RUN: %llvm_jitlink -all_load %t/all_load.o -L%t -lSetX
+//
+// REQUIRES: system-darwin && host-arch-compatible
+
+int x = 0;
+
+int main(int argc, char *argv[]) { return x == 1 ? 0 : 1; }

compiler-rt/test/orc/lit.cfg.py

@@ -14,6 +14,8 @@ host_arch_compatible = config.target_arch == config.host_arch
 
 if config.host_arch == "x86_64h" and config.target_arch == "x86_64":
     host_arch_compatible = True
+if host_arch_compatible:
+    config.available_features.add("host-arch-compatible")
 config.test_target_is_host_executable = (
     config.target_os == config.host_os and host_arch_compatible
 )
@@ -71,9 +73,10 @@ config.substitutions.append(
         (lli + " -jit-kind=orc -jit-linker=jitlink -orc-runtime=" + orc_rt_path),
     )
 )
+config.substitutions.append(("%ar", "ar"))
 
 # Default test suffixes.
-config.suffixes = [".c", ".cpp", ".S", ".ll", ".test"]
+config.suffixes = [".c", ".cpp", ".m", ".S", ".ll", ".test"]
 
 # Exclude Inputs directories.
 config.excludes = ["Inputs"]

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

@@ -268,11 +268,21 @@ private:
 /// the containing object being added to the JITDylib.
 class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
 public:
-  // Interface builder function for objects loaded from this archive.
+  /// Interface builder function for objects loaded from this archive.
   using GetObjectFileInterface =
       unique_function<Expected<MaterializationUnit::Interface>(
           ExecutionSession &ES, MemoryBufferRef ObjBuffer)>;
 
+  /// Callback for visiting archive members at construction time.
+  /// Con be used to pre-load archive members.
+  using VisitMembersFunction = unique_function<Error(MemoryBufferRef)>;
+
+  /// A VisitMembersFunction that unconditionally loads all object files from
+  /// the archive.
+  /// Archive members that are not valid object files will be skipped.
+  static VisitMembersFunction loadAllObjectFileMembers(ObjectLayer &L,
+                                                       JITDylib &JD);
+
   /// Try to create a StaticLibraryDefinitionGenerator from the given path.
   ///
   /// This call will succeed if the file at the given path is a static library
@@ -280,6 +290,7 @@ public:
   /// with the ExecutionSession's triple. Otherwise it will return an error.
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Load(ObjectLayer &L, const char *FileName,
+       VisitMembersFunction VisitMembers = VisitMembersFunction(),
        GetObjectFileInterface GetObjFileInterface = GetObjectFileInterface());
 
   /// Try to create a StaticLibrarySearchGenerator from the given memory buffer
@@ -287,6 +298,7 @@ public:
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Create(ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
          std::unique_ptr<object::Archive> Archive,
+         VisitMembersFunction VisitMembers = VisitMembersFunction(),
          GetObjectFileInterface GetObjFileInterface = GetObjectFileInterface());
 
   /// Try to create a StaticLibrarySearchGenerator from the given memory buffer.
@@ -298,6 +310,7 @@ public:
   /// with the ExecutionSession's triple. Otherwise it will return an error.
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Create(ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
+         VisitMembersFunction VisitMembers = VisitMembersFunction(),
          GetObjectFileInterface GetObjFileInterface = GetObjectFileInterface());
 
   /// Returns a list of filenames of dynamic libraries that this archive has

llvm/include/llvm/ExecutionEngine/Orc/MachO.h

@@ -28,6 +28,9 @@ class MachOUniversalBinary;
 
 namespace orc {
 
+class JITDylib;
+class ObjectLayer;
+
 /// Check that the given buffer contains a MachO object file compatible with the
 /// given triple.
 /// ObjIsSlice should be set to true if Obj is a slice of a universal binary
@@ -72,6 +75,20 @@ getMachOSliceRangeForTriple(object::MachOUniversalBinary &UB, const Triple &TT);
 Expected<std::pair<size_t, size_t>>
 getMachOSliceRangeForTriple(MemoryBufferRef UBBuf, const Triple &TT);
 
+/// For use with StaticLibraryDefinitionGenerators.
+class ForceLoadMachOArchiveMembers {
+public:
+  ForceLoadMachOArchiveMembers(ObjectLayer &L, JITDylib &JD, bool ObjCOnly)
+      : L(L), JD(JD), ObjCOnly(ObjCOnly) {}
+
+  Error operator()(MemoryBufferRef MemberBuf);
+
+private:
+  ObjectLayer &L;
+  JITDylib &JD;
+  bool ObjCOnly;
+};
+
 } // namespace orc
 } // namespace llvm
 

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -273,9 +273,24 @@ Error DynamicLibrarySearchGenerator::tryToGenerate(
   return JD.define(absoluteSymbols(std::move(NewSymbols)));
 }
 
+StaticLibraryDefinitionGenerator::VisitMembersFunction
+StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(ObjectLayer &L,
+                                                           JITDylib &JD) {
+  return [&](MemoryBufferRef Buf) -> Error {
+    switch (identify_magic(Buf.getBuffer())) {
+    case file_magic::elf_relocatable:
+    case file_magic::macho_object:
+    case file_magic::coff_object:
+      return L.add(JD, MemoryBuffer::getMemBuffer(Buf));
+    default:
+      return Error::success();
+    }
+  };
+}
+
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
 StaticLibraryDefinitionGenerator::Load(
-    ObjectLayer &L, const char *FileName,
+    ObjectLayer &L, const char *FileName, VisitMembersFunction VisitMembers,
     GetObjectFileInterface GetObjFileInterface) {
 
   const auto &TT = L.getExecutionSession().getTargetTriple();
@@ -283,17 +298,33 @@ StaticLibraryDefinitionGenerator::Load(
   if (!Linkable)
     return Linkable.takeError();
 
-  return Create(L, std::move(Linkable->first), std::move(GetObjFileInterface));
+  return Create(L, std::move(Linkable->first), std::move(VisitMembers),
+                std::move(GetObjFileInterface));
 }
 
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
 StaticLibraryDefinitionGenerator::Create(
     ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
-    std::unique_ptr<object::Archive> Archive,
+    std::unique_ptr<object::Archive> Archive, VisitMembersFunction VisitMembers,
     GetObjectFileInterface GetObjFileInterface) {
 
   Error Err = Error::success();
 
+  if (VisitMembers) {
+    for (auto Child : Archive->children(Err)) {
+      if (auto ChildBuf = Child.getMemoryBufferRef()) {
+        if (auto Err2 = VisitMembers(*ChildBuf))
+          return std::move(Err2);
+      } else {
+        // We silently allow non-object archive members. This matches the
+        // behavior of ld.
+        consumeError(ChildBuf.takeError());
+      }
+    }
+    if (Err)
+      return std::move(Err);
+  }
+
   std::unique_ptr<StaticLibraryDefinitionGenerator> ADG(
       new StaticLibraryDefinitionGenerator(
           L, std::move(ArchiveBuffer), std::move(Archive),
@@ -308,6 +339,7 @@ StaticLibraryDefinitionGenerator::Create(
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
 StaticLibraryDefinitionGenerator::Create(
     ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
+    VisitMembersFunction VisitMembers,
     GetObjectFileInterface GetObjFileInterface) {
 
   auto B = object::createBinary(ArchiveBuffer->getMemBufferRef());
@@ -319,7 +351,7 @@ StaticLibraryDefinitionGenerator::Create(
     return Create(L, std::move(ArchiveBuffer),
                   std::unique_ptr<object::Archive>(
                       static_cast<object::Archive *>(B->release())),
-                  std::move(GetObjFileInterface));
+                  std::move(VisitMembers), std::move(GetObjFileInterface));
 
   // If this is a universal binary then search for a slice matching the given
   // Triple.
@@ -341,7 +373,7 @@ StaticLibraryDefinitionGenerator::Create(
       return Archive.takeError();
 
     return Create(L, std::move(ArchiveBuffer), std::move(*Archive),
-                  std::move(GetObjFileInterface));
+                  std::move(VisitMembers), std::move(GetObjFileInterface));
   }
 
   return make_error<StringError>(Twine("Unrecognized file type for ") +

llvm/lib/ExecutionEngine/Orc/MachO.cpp

@@ -10,6 +10,7 @@
 
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/BinaryFormat/MachO.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/Object/MachOUniversal.h"
 #include "llvm/Support/FileSystem.h"
 
@@ -228,5 +229,39 @@ getMachOSliceRangeForTriple(MemoryBufferRef UBBuf, const Triple &TT) {
   return getMachOSliceRangeForTriple(**UB, TT);
 }
 
+Error ForceLoadMachOArchiveMembers::operator()(MemoryBufferRef MemberBuf) {
+  if (!ObjCOnly)
+    return L.add(JD, MemoryBuffer::getMemBuffer(MemberBuf));
+
+  // We need to check whether this archive member contains any Objective-C
+  // or Swift metadata.
+
+  auto Obj = object::ObjectFile::createObjectFile(MemberBuf);
+  if (!Obj) {
+    // We silently ignore invalid files.
+    consumeError(Obj.takeError());
+    return Error::success();
+  }
+
+  if (auto *MachOObj = dyn_cast<object::MachOObjectFile>(&**Obj)) {
+    // Load the object if any recognized special section is present.
+    for (auto Sec : MachOObj->sections()) {
+      auto SegName =
+          MachOObj->getSectionFinalSegmentName(Sec.getRawDataRefImpl());
+      if (auto SecName = Sec.getName()) {
+        if (*SecName == "__objc_classlist" || *SecName == "__objc_protolist" ||
+            *SecName == "__objc_clsrolist" || *SecName == "__objc_catlist" ||
+            *SecName == "__objc_catlist2" || *SecName == "__objc_nlcatlist" ||
+            (SegName == "__TEXT" && (*SecName).starts_with("__swift") &&
+             *SecName != "__swift_modhash"))
+          return L.add(JD, MemoryBuffer::getMemBuffer(MemberBuf));
+      } else
+        return SecName.takeError();
+    }
+  }
+
+  return Error::success();
+}
+
 } // End namespace orc.
 } // End namespace llvm.

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -30,6 +30,7 @@
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LoadLinkableFile.h"
+#include "llvm/ExecutionEngine/Orc/MachO.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
@@ -265,6 +266,17 @@ static cl::opt<std::string>
     OverrideTriple("triple", cl::desc("Override target triple detection"),
                    cl::init(""), cl::cat(JITLinkCategory));
 
+static cl::opt<bool> AllLoad("all_load",
+                             cl::desc("Load all members of static archives"),
+                             cl::init(false), cl::cat(JITLinkCategory));
+
+static cl::opt<bool> ForceLoadObjC(
+    "ObjC",
+    cl::desc("Load all members of static archives that implement "
+             "Objective-C classes or categories, or Swift structs, "
+             "classes or extensions"),
+    cl::init(false), cl::cat(JITLinkCategory));
+
 static ExitOnError ExitOnErr;
 
 static LLVM_ATTRIBUTE_USED void linkComponents() {
@@ -1957,10 +1969,9 @@ static Error addLibraries(Session &S,
              });
 
   // 3. Process library loads.
-  auto AddArchive = [&](const char *Path, const LibraryLoad &LL)
+  auto AddArchive = [&](JITDylib &JD, const char *Path, const LibraryLoad &LL)
       -> Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>> {
-    unique_function<Expected<MaterializationUnit::Interface>(
-        ExecutionSession & ES, MemoryBufferRef ObjBuffer)>
+    StaticLibraryDefinitionGenerator::GetObjectFileInterface
         GetObjFileInterface;
     switch (LL.Modifier) {
     case LibraryLoad::Standard:
@@ -1970,8 +1981,17 @@ static Error addLibraries(Session &S,
       GetObjFileInterface = getObjectFileInterfaceHidden;
       break;
     }
+
+    StaticLibraryDefinitionGenerator::VisitMembersFunction VisitMembers;
+    if (AllLoad)
+      VisitMembers = StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(
+          S.ObjLayer, JD);
+    else if (S.ES.getTargetTriple().isOSBinFormatMachO() && ForceLoadObjC)
+      VisitMembers = ForceLoadMachOArchiveMembers(S.ObjLayer, JD, true);
+
     auto G = StaticLibraryDefinitionGenerator::Load(
-        S.ObjLayer, Path, std::move(GetObjFileInterface));
+        S.ObjLayer, Path, std::move(VisitMembers),
+        std::move(GetObjFileInterface));
     if (!G)
       return G.takeError();
 
@@ -2009,7 +2029,7 @@ static Error addLibraries(Session &S,
     }
 
     if (LL.IsPath) {
-      auto G = AddArchive(LL.LibName.c_str(), LL);
+      auto G = AddArchive(JD, LL.LibName.c_str(), LL);
       if (!G)
         return createFileError(LL.LibName, G.takeError());
       JD.addGenerator(std::move(*G));
@@ -2065,7 +2085,7 @@ static Error addLibraries(Session &S,
         }
         case file_magic::archive:
         case file_magic::macho_universal_binary: {
-          auto G = AddArchive(LibPath.data(), LL);
+          auto G = AddArchive(JD, LibPath.data(), LL);
           if (!G)
             return G.takeError();
           JD.addGenerator(std::move(*G));