[ORC][MachO] Support common load commands in the platform's mach-o header builder

benlangmuir · benlangmuir · commit b08aca7a4fb9 · 2024-01-18T09:14:15.000-08:00
Add a HeaderOptions struct that can be used to configure commonly-used
load commands LC_ID_DYLIB, LC_LOAD_DYLIB, and LC_RPATH when setupDylib
creates a mach-o header.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOBuilder.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOBuilder.h
@@ -104,6 +104,47 @@ struct MachOBuilderLoadCommand<MachO::LC_ID_DYLIB>
   std::string Name;
 };
 
+template <>
+struct MachOBuilderLoadCommand<MachO::LC_LOAD_DYLIB>
+    : public MachOBuilderLoadCommandImplBase<MachO::LC_LOAD_DYLIB> {
+
+  MachOBuilderLoadCommand(std::string Name, uint32_t Timestamp,
+                          uint32_t CurrentVersion,
+                          uint32_t CompatibilityVersion)
+      : MachOBuilderLoadCommandImplBase(
+            MachO::dylib{24, Timestamp, CurrentVersion, CompatibilityVersion}),
+        Name(std::move(Name)) {
+    cmdsize += (this->Name.size() + 1 + 3) & ~0x3;
+  }
+
+  size_t write(MutableArrayRef<char> Buf, size_t Offset,
+               bool SwapStruct) override {
+    Offset = writeMachOStruct(Buf, Offset, rawStruct(), SwapStruct);
+    strcpy(Buf.data() + Offset, Name.data());
+    return Offset + ((Name.size() + 1 + 3) & ~0x3);
+  }
+
+  std::string Name;
+};
+
+template <>
+struct MachOBuilderLoadCommand<MachO::LC_RPATH>
+    : public MachOBuilderLoadCommandImplBase<MachO::LC_RPATH> {
+  MachOBuilderLoadCommand(std::string Path)
+      : MachOBuilderLoadCommandImplBase(12u), Path(std::move(Path)) {
+    cmdsize += (this->Path.size() + 1 + 3) & ~0x3;
+  }
+
+  size_t write(MutableArrayRef<char> Buf, size_t Offset,
+               bool SwapStruct) override {
+    Offset = writeMachOStruct(Buf, Offset, rawStruct(), SwapStruct);
+    strcpy(Buf.data() + Offset, Path.data());
+    return Offset + ((Path.size() + 1 + 3) & ~0x3);
+  }
+
+  std::string Path;
+};
+
 // Builds MachO objects.
 template <typename MachOTraits> class MachOBuilder {
 private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
@@ -47,14 +47,38 @@ class MachOPlatform : public Platform {
     LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Callable)
   };
 
+  /// Configuration for the mach-o header of a JITDylib. Specify common load
+  /// commands that should be added to the header.
+  struct HeaderOptions {
+    /// A dylib for use with a dylib command (e.g. LC_ID_DYLIB, LC_LOAD_DYLIB).
+    struct Dylib {
+      std::string Name;
+      uint32_t Timestamp;
+      uint32_t CurrentVersion;
+      uint32_t CompatibilityVersion;
+    };
+
+    /// Override for LC_IC_DYLIB. If this is nullopt, {JD.getName(), 0, 0, 0}
+    /// will be used.
+    std::optional<Dylib> IDDylib;
+    /// List of LC_LOAD_DYLIBs.
+    std::vector<Dylib> LoadDylibs;
+    /// List of LC_RPATHs.
+    std::vector<std::string> RPaths;
+
+    HeaderOptions() = default;
+    HeaderOptions(Dylib D) : IDDylib(std::move(D)) {}
+  };
+
   /// Used by setupJITDylib to create MachO header MaterializationUnits for
   /// JITDylibs.
   using MachOHeaderMUBuilder =
-      unique_function<std::unique_ptr<MaterializationUnit>(MachOPlatform &MOP)>;
+      unique_function<std::unique_ptr<MaterializationUnit>(MachOPlatform &MOP,
+                                                           HeaderOptions Opts)>;
 
   /// Simple MachO header graph builder.
   static inline std::unique_ptr<MaterializationUnit>
-  buildSimpleMachOHeaderMU(MachOPlatform &MOP);
+  buildSimpleMachOHeaderMU(MachOPlatform &MOP, HeaderOptions Opts);
 
   /// Try to create a MachOPlatform instance, adding the ORC runtime to the
   /// given JITDylib.
@@ -97,13 +121,15 @@ class MachOPlatform : public Platform {
   static Expected<std::unique_ptr<MachOPlatform>>
   Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
          JITDylib &PlatformJD, std::unique_ptr<DefinitionGenerator> OrcRuntime,
+         HeaderOptions PlatformJDOpts = {},
          MachOHeaderMUBuilder BuildMachOHeaderMU = buildSimpleMachOHeaderMU,
          std::optional<SymbolAliasMap> RuntimeAliases = std::nullopt);
 
   /// Construct using a path to the ORC runtime.
   static Expected<std::unique_ptr<MachOPlatform>>
   Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
          JITDylib &PlatformJD, const char *OrcRuntimePath,
+         HeaderOptions PlatformJDOpts = {},
          MachOHeaderMUBuilder BuildMachOHeaderMU = buildSimpleMachOHeaderMU,
          std::optional<SymbolAliasMap> RuntimeAliases = std::nullopt);
 
@@ -115,6 +141,11 @@ class MachOPlatform : public Platform {
   }
 
   Error setupJITDylib(JITDylib &JD) override;
+
+  /// Install any platform-specific symbols (e.g. `__dso_handle`) and create a
+  /// mach-o header based on the given options.
+  Error setupJITDylib(JITDylib &JD, HeaderOptions Opts);
+
   Error teardownJITDylib(JITDylib &JD) override;
   Error notifyAdding(ResourceTracker &RT,
                      const MaterializationUnit &MU) override;
@@ -258,6 +289,7 @@ class MachOPlatform : public Platform {
   MachOPlatform(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
                 JITDylib &PlatformJD,
                 std::unique_ptr<DefinitionGenerator> OrcRuntimeGenerator,
+                HeaderOptions PlatformJDOpts,
                 MachOHeaderMUBuilder BuildMachOHeaderMU, Error &Err);
 
   // Associate MachOPlatform JIT-side runtime support functions with handlers.
@@ -336,7 +368,8 @@ class MachOPlatform : public Platform {
 // Generates a MachO header.
 class SimpleMachOHeaderMU : public MaterializationUnit {
 public:
-  SimpleMachOHeaderMU(MachOPlatform &MOP, SymbolStringPtr HeaderStartSymbol);
+  SimpleMachOHeaderMU(MachOPlatform &MOP, SymbolStringPtr HeaderStartSymbol,
+                      MachOPlatform::HeaderOptions Opts);
   StringRef getName() const override { return "MachOHeaderMU"; }
   void materialize(std::unique_ptr<MaterializationResponsibility> R) override;
   void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override;
@@ -346,6 +379,7 @@ class SimpleMachOHeaderMU : public MaterializationUnit {
                                             jitlink::Section &HeaderSection);
 
   MachOPlatform &MOP;
+  MachOPlatform::HeaderOptions Opts;
 
 private:
   struct HeaderSymbol {
@@ -365,8 +399,10 @@ class SimpleMachOHeaderMU : public MaterializationUnit {
 
 /// Simple MachO header graph builder.
 inline std::unique_ptr<MaterializationUnit>
-MachOPlatform::buildSimpleMachOHeaderMU(MachOPlatform &MOP) {
-  return std::make_unique<SimpleMachOHeaderMU>(MOP, MOP.MachOHeaderStartSymbol);
+MachOPlatform::buildSimpleMachOHeaderMU(MachOPlatform &MOP,
+                                        HeaderOptions Opts) {
+  return std::make_unique<SimpleMachOHeaderMU>(MOP, MOP.MachOHeaderStartSymbol,
+                                               std::move(Opts));
 }
 
 struct MachOHeaderInfo {
diff --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -255,12 +255,11 @@ struct ObjCImageInfoFlags {
 namespace llvm {
 namespace orc {
 
-Expected<std::unique_ptr<MachOPlatform>>
-MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
-                      JITDylib &PlatformJD,
-                      std::unique_ptr<DefinitionGenerator> OrcRuntime,
-                      MachOHeaderMUBuilder BuildMachOHeaderMU,
-                      std::optional<SymbolAliasMap> RuntimeAliases) {
+Expected<std::unique_ptr<MachOPlatform>> MachOPlatform::Create(
+    ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+    JITDylib &PlatformJD, std::unique_ptr<DefinitionGenerator> OrcRuntime,
+    HeaderOptions PlatformJDOpts, MachOHeaderMUBuilder BuildMachOHeaderMU,
+    std::optional<SymbolAliasMap> RuntimeAliases) {
 
   // If the target is not supported then bail out immediately.
   if (!supportedTarget(ES.getTargetTriple()))
@@ -290,9 +289,9 @@ MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
 
   // Create the instance.
   Error Err = Error::success();
-  auto P = std::unique_ptr<MachOPlatform>(
-      new MachOPlatform(ES, ObjLinkingLayer, PlatformJD, std::move(OrcRuntime),
-                        std::move(BuildMachOHeaderMU), Err));
+  auto P = std::unique_ptr<MachOPlatform>(new MachOPlatform(
+      ES, ObjLinkingLayer, PlatformJD, std::move(OrcRuntime),
+      std::move(PlatformJDOpts), std::move(BuildMachOHeaderMU), Err));
   if (Err)
     return std::move(Err);
   return std::move(P);
@@ -301,6 +300,7 @@ MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
 Expected<std::unique_ptr<MachOPlatform>>
 MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
                       JITDylib &PlatformJD, const char *OrcRuntimePath,
+                      HeaderOptions PlatformJDOpts,
                       MachOHeaderMUBuilder BuildMachOHeaderMU,
                       std::optional<SymbolAliasMap> RuntimeAliases) {
 
@@ -312,11 +312,16 @@ MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
 
   return Create(ES, ObjLinkingLayer, PlatformJD,
                 std::move(*OrcRuntimeArchiveGenerator),
-                std::move(BuildMachOHeaderMU), std::move(RuntimeAliases));
+                std::move(PlatformJDOpts), std::move(BuildMachOHeaderMU),
+                std::move(RuntimeAliases));
 }
 
 Error MachOPlatform::setupJITDylib(JITDylib &JD) {
-  if (auto Err = JD.define(BuildMachOHeaderMU(*this)))
+  return setupJITDylib(JD, /*Opts=*/{});
+}
+
+Error MachOPlatform::setupJITDylib(JITDylib &JD, HeaderOptions Opts) {
+  if (auto Err = JD.define(BuildMachOHeaderMU(*this, std::move(Opts))))
     return Err;
 
   return ES.lookup({&JD}, MachOHeaderStartSymbol).takeError();
@@ -432,7 +437,8 @@ MachOPlatform::MachOPlatform(
     ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
     JITDylib &PlatformJD,
     std::unique_ptr<DefinitionGenerator> OrcRuntimeGenerator,
-    MachOHeaderMUBuilder BuildMachOHeaderMU, Error &Err)
+    HeaderOptions PlatformJDOpts, MachOHeaderMUBuilder BuildMachOHeaderMU,
+    Error &Err)
     : ES(ES), PlatformJD(PlatformJD), ObjLinkingLayer(ObjLinkingLayer),
       BuildMachOHeaderMU(std::move(BuildMachOHeaderMU)) {
   ErrorAsOutParameter _(&Err);
@@ -497,7 +503,8 @@ MachOPlatform::MachOPlatform(
   //    the support methods callable. The bootstrap is now complete.
 
   // Step (1) Add header materialization unit and request.
-  if ((Err = PlatformJD.define(this->BuildMachOHeaderMU(*this))))
+  if ((Err = PlatformJD.define(
+           this->BuildMachOHeaderMU(*this, std::move(PlatformJDOpts)))))
     return;
   if ((Err = ES.lookup(&PlatformJD, MachOHeaderStartSymbol).takeError()))
     return;
@@ -1669,9 +1676,10 @@ Error MachOPlatform::MachOPlatformPlugin::addSymbolTableRegistration(
 }
 
 template <typename MachOTraits>
-jitlink::Block &createTrivialHeaderBlock(MachOPlatform &MOP,
-                                         jitlink::LinkGraph &G,
-                                         jitlink::Section &HeaderSection) {
+jitlink::Block &createHeaderBlock(MachOPlatform &MOP,
+                                  const MachOPlatform::HeaderOptions &Opts,
+                                  JITDylib &JD, jitlink::LinkGraph &G,
+                                  jitlink::Section &HeaderSection) {
   auto HdrInfo =
       getMachOHeaderInfoFromTriple(MOP.getExecutionSession().getTargetTriple());
   MachOBuilder<MachOTraits> B(HdrInfo.PageSize);
@@ -1680,6 +1688,19 @@ jitlink::Block &createTrivialHeaderBlock(MachOPlatform &MOP,
   B.Header.cputype = HdrInfo.CPUType;
   B.Header.cpusubtype = HdrInfo.CPUSubType;
 
+  if (Opts.IDDylib)
+    B.template addLoadCommand<MachO::LC_ID_DYLIB>(
+        Opts.IDDylib->Name, Opts.IDDylib->Timestamp,
+        Opts.IDDylib->CurrentVersion, Opts.IDDylib->CompatibilityVersion);
+  else
+    B.template addLoadCommand<MachO::LC_ID_DYLIB>(JD.getName(), 0, 0, 0);
+
+  for (auto &D : Opts.LoadDylibs)
+    B.template addLoadCommand<MachO::LC_LOAD_DYLIB>(
+        D.Name, D.Timestamp, D.CurrentVersion, D.CompatibilityVersion);
+  for (auto &P : Opts.RPaths)
+    B.template addLoadCommand<MachO::LC_RPATH>(P);
+
   auto HeaderContent = G.allocateBuffer(B.layout());
   B.write(HeaderContent);
 
@@ -1688,10 +1709,11 @@ jitlink::Block &createTrivialHeaderBlock(MachOPlatform &MOP,
 }
 
 SimpleMachOHeaderMU::SimpleMachOHeaderMU(MachOPlatform &MOP,
-                                         SymbolStringPtr HeaderStartSymbol)
+                                         SymbolStringPtr HeaderStartSymbol,
+                                         MachOPlatform::HeaderOptions Opts)
     : MaterializationUnit(
           createHeaderInterface(MOP, std::move(HeaderStartSymbol))),
-      MOP(MOP) {}
+      MOP(MOP), Opts(std::move(Opts)) {}
 
 void SimpleMachOHeaderMU::materialize(
     std::unique_ptr<MaterializationResponsibility> R) {
@@ -1725,7 +1747,7 @@ SimpleMachOHeaderMU::createHeaderBlock(JITDylib &JD, jitlink::LinkGraph &G,
   switch (MOP.getExecutionSession().getTargetTriple().getArch()) {
   case Triple::aarch64:
   case Triple::x86_64:
-    return createTrivialHeaderBlock<MachO64LE>(MOP, G, HeaderSection);
+    return ::createHeaderBlock<MachO64LE>(MOP, Opts, JD, G, HeaderSection);
   default:
     llvm_unreachable("Unsupported architecture");
   }
