[ORC][MachO] Expose SimpleMachOHeaderMU and related utilities.

SimpleMachOHeaderMU can be used as a convenient base for classes that create
custom MachO headers. Instances of these custom classes can be used by passing
a MachOHeaderMUBuilder using the MachOPlatform extensions added in ef314d3.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h

@@ -52,8 +52,9 @@ class MachOPlatform : public Platform {
   using MachOHeaderMUBuilder =
       unique_function<std::unique_ptr<MaterializationUnit>(MachOPlatform &MOP)>;
 
-  static std::unique_ptr<MaterializationUnit>
-  defaultMachOHeaderBuilder(MachOPlatform &MOP);
+  /// Simple MachO header graph builder.
+  static inline std::unique_ptr<MaterializationUnit>
+  buildSimpleMachOHeaderMU(MachOPlatform &MOP);
 
   /// Try to create a MachOPlatform instance, adding the ORC runtime to the
   /// given JITDylib.
@@ -96,14 +97,14 @@ class MachOPlatform : public Platform {
   static Expected<std::unique_ptr<MachOPlatform>>
   Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
          JITDylib &PlatformJD, std::unique_ptr<DefinitionGenerator> OrcRuntime,
-         MachOHeaderMUBuilder BuildMachOHeaderMU = defaultMachOHeaderBuilder,
+         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,
-         MachOHeaderMUBuilder BuildMachOHeaderMU = defaultMachOHeaderBuilder,
+         MachOHeaderMUBuilder BuildMachOHeaderMU = buildSimpleMachOHeaderMU,
          std::optional<SymbolAliasMap> RuntimeAliases = std::nullopt);
 
   ExecutionSession &getExecutionSession() const { return ES; }
@@ -332,6 +333,49 @@ class MachOPlatform : public Platform {
   std::atomic<BootstrapInfo *> Bootstrap;
 };
 
+// Generates a MachO header.
+class SimpleMachOHeaderMU : public MaterializationUnit {
+public:
+  SimpleMachOHeaderMU(MachOPlatform &MOP, SymbolStringPtr HeaderStartSymbol);
+  StringRef getName() const override { return "MachOHeaderMU"; }
+  void materialize(std::unique_ptr<MaterializationResponsibility> R) override;
+  void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override;
+
+protected:
+  virtual jitlink::Block &createHeaderBlock(JITDylib &JD, jitlink::LinkGraph &G,
+                                            jitlink::Section &HeaderSection);
+
+private:
+  struct HeaderSymbol {
+    const char *Name;
+    uint64_t Offset;
+  };
+
+  static constexpr HeaderSymbol AdditionalHeaderSymbols[] = {
+      {"___mh_executable_header", 0}};
+
+  void addMachOHeader(JITDylib &JD, jitlink::LinkGraph &G,
+                      const SymbolStringPtr &InitializerSymbol);
+  static MaterializationUnit::Interface
+  createHeaderInterface(MachOPlatform &MOP,
+                        const SymbolStringPtr &HeaderStartSymbol);
+
+  MachOPlatform &MOP;
+};
+
+/// Simple MachO header graph builder.
+inline std::unique_ptr<MaterializationUnit>
+MachOPlatform::buildSimpleMachOHeaderMU(MachOPlatform &MOP) {
+  return std::make_unique<SimpleMachOHeaderMU>(MOP, MOP.MachOHeaderStartSymbol);
+}
+
+struct MachOHeaderInfo {
+  size_t PageSize = 0;
+  uint32_t CPUType = 0;
+  uint32_t CPUSubType = 0;
+};
+MachOHeaderInfo getMachOHeaderInfoFromTriple(const Triple &TT);
+
 } // end namespace orc
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -116,101 +116,6 @@ std::unique_ptr<jitlink::LinkGraph> createPlatformGraph(MachOPlatform &MOP,
                                               jitlink::getGenericEdgeKindName);
 }
 
-// Generates a MachO header.
-class MachOHeaderMaterializationUnit : public MaterializationUnit {
-public:
-  MachOHeaderMaterializationUnit(MachOPlatform &MOP,
-                                 SymbolStringPtr HeaderStartSymbol)
-      : MaterializationUnit(
-            createHeaderInterface(MOP, std::move(HeaderStartSymbol))),
-        MOP(MOP) {}
-
-  StringRef getName() const override { return "MachOHeaderMU"; }
-
-  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
-    auto G = createPlatformGraph(MOP, "<MachOHeaderMU>");
-    addMachOHeader(*G, MOP, R->getInitializerSymbol());
-    MOP.getObjectLinkingLayer().emit(std::move(R), std::move(G));
-  }
-
-  void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {}
-
-  static void addMachOHeader(jitlink::LinkGraph &G, MachOPlatform &MOP,
-                             const SymbolStringPtr &InitializerSymbol) {
-    auto &HeaderSection = G.createSection("__header", MemProt::Read);
-    auto &HeaderBlock = createHeaderBlock(G, HeaderSection);
-
-    // Init symbol is header-start symbol.
-    G.addDefinedSymbol(HeaderBlock, 0, *InitializerSymbol,
-                       HeaderBlock.getSize(), jitlink::Linkage::Strong,
-                       jitlink::Scope::Default, false, true);
-    for (auto &HS : AdditionalHeaderSymbols)
-      G.addDefinedSymbol(HeaderBlock, HS.Offset, HS.Name, HeaderBlock.getSize(),
-                         jitlink::Linkage::Strong, jitlink::Scope::Default,
-                         false, true);
-  }
-
-private:
-  struct HeaderSymbol {
-    const char *Name;
-    uint64_t Offset;
-  };
-
-  static constexpr HeaderSymbol AdditionalHeaderSymbols[] = {
-      {"___mh_executable_header", 0}};
-
-  static jitlink::Block &createHeaderBlock(jitlink::LinkGraph &G,
-                                           jitlink::Section &HeaderSection) {
-    MachO::mach_header_64 Hdr;
-    Hdr.magic = MachO::MH_MAGIC_64;
-    switch (G.getTargetTriple().getArch()) {
-    case Triple::aarch64:
-      Hdr.cputype = MachO::CPU_TYPE_ARM64;
-      Hdr.cpusubtype = MachO::CPU_SUBTYPE_ARM64_ALL;
-      break;
-    case Triple::x86_64:
-      Hdr.cputype = MachO::CPU_TYPE_X86_64;
-      Hdr.cpusubtype = MachO::CPU_SUBTYPE_X86_64_ALL;
-      break;
-    default:
-      llvm_unreachable("Unrecognized architecture");
-    }
-    Hdr.filetype = MachO::MH_DYLIB; // Custom file type?
-    Hdr.ncmds = 0;
-    Hdr.sizeofcmds = 0;
-    Hdr.flags = 0;
-    Hdr.reserved = 0;
-
-    if (G.getEndianness() != llvm::endianness::native)
-      MachO::swapStruct(Hdr);
-
-    auto HeaderContent = G.allocateContent(
-        ArrayRef<char>(reinterpret_cast<const char *>(&Hdr), sizeof(Hdr)));
-
-    return G.createContentBlock(HeaderSection, HeaderContent, ExecutorAddr(), 8,
-                                0);
-  }
-
-  static MaterializationUnit::Interface
-  createHeaderInterface(MachOPlatform &MOP,
-                        const SymbolStringPtr &HeaderStartSymbol) {
-    SymbolFlagsMap HeaderSymbolFlags;
-
-    HeaderSymbolFlags[HeaderStartSymbol] = JITSymbolFlags::Exported;
-    for (auto &HS : AdditionalHeaderSymbols)
-      HeaderSymbolFlags[MOP.getExecutionSession().intern(HS.Name)] =
-          JITSymbolFlags::Exported;
-
-    return MaterializationUnit::Interface(std::move(HeaderSymbolFlags),
-                                          HeaderStartSymbol);
-  }
-
-  MachOPlatform &MOP;
-};
-
-constexpr MachOHeaderMaterializationUnit::HeaderSymbol
-    MachOHeaderMaterializationUnit::AdditionalHeaderSymbols[];
-
 // Creates a Bootstrap-Complete LinkGraph to run deferred actions.
 class MachOPlatformCompleteBootstrapMaterializationUnit
     : public MaterializationUnit {
@@ -350,12 +255,6 @@ struct ObjCImageInfoFlags {
 namespace llvm {
 namespace orc {
 
-std::unique_ptr<MaterializationUnit>
-MachOPlatform::defaultMachOHeaderBuilder(MachOPlatform &MOP) {
-  return std::make_unique<MachOHeaderMaterializationUnit>(
-      MOP, SymbolStringPtr(MOP.getMachOHeaderStartSymbol()));
-}
-
 Expected<std::unique_ptr<MachOPlatform>>
 MachOPlatform::Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
                       JITDylib &PlatformJD,
@@ -598,8 +497,7 @@ MachOPlatform::MachOPlatform(
   //    the support methods callable. The bootstrap is now complete.
 
   // Step (1) Add header materialization unit and request.
-  if ((Err = PlatformJD.define(std::make_unique<MachOHeaderMaterializationUnit>(
-           *this, MachOHeaderStartSymbol))))
+  if ((Err = PlatformJD.define(this->BuildMachOHeaderMU(*this))))
     return;
   if ((Err = ES.lookup(&PlatformJD, MachOHeaderStartSymbol).takeError()))
     return;
@@ -1767,5 +1665,97 @@ Error MachOPlatform::MachOPlatformPlugin::addSymbolTableRegistration(
 
   return Error::success();
 }
+
+template <typename MachOTraits>
+jitlink::Block &createTrivialHeaderBlock(MachOPlatform &MOP,
+                                         jitlink::LinkGraph &G,
+                                         jitlink::Section &HeaderSection) {
+  auto HdrInfo =
+      getMachOHeaderInfoFromTriple(MOP.getExecutionSession().getTargetTriple());
+  MachOBuilder<MachOTraits> B(HdrInfo.PageSize);
+
+  B.Header.filetype = MachO::MH_DYLIB;
+  B.Header.cputype = HdrInfo.CPUType;
+  B.Header.cpusubtype = HdrInfo.CPUSubType;
+
+  auto HeaderContent = G.allocateBuffer(B.layout());
+  B.write(HeaderContent);
+
+  return G.createContentBlock(HeaderSection, HeaderContent, ExecutorAddr(), 8,
+                              0);
+}
+
+SimpleMachOHeaderMU::SimpleMachOHeaderMU(MachOPlatform &MOP,
+                                         SymbolStringPtr HeaderStartSymbol)
+    : MaterializationUnit(
+          createHeaderInterface(MOP, std::move(HeaderStartSymbol))),
+      MOP(MOP) {}
+
+void SimpleMachOHeaderMU::materialize(
+    std::unique_ptr<MaterializationResponsibility> R) {
+  auto G = createPlatformGraph(MOP, "<MachOHeaderMU>");
+  addMachOHeader(R->getTargetJITDylib(), *G, R->getInitializerSymbol());
+  MOP.getObjectLinkingLayer().emit(std::move(R), std::move(G));
+}
+
+void SimpleMachOHeaderMU::discard(const JITDylib &JD,
+                                  const SymbolStringPtr &Sym) {}
+
+void SimpleMachOHeaderMU::addMachOHeader(
+    JITDylib &JD, jitlink::LinkGraph &G,
+    const SymbolStringPtr &InitializerSymbol) {
+  auto &HeaderSection = G.createSection("__header", MemProt::Read);
+  auto &HeaderBlock = createHeaderBlock(JD, G, HeaderSection);
+
+  // Init symbol is header-start symbol.
+  G.addDefinedSymbol(HeaderBlock, 0, *InitializerSymbol, HeaderBlock.getSize(),
+                     jitlink::Linkage::Strong, jitlink::Scope::Default, false,
+                     true);
+  for (auto &HS : AdditionalHeaderSymbols)
+    G.addDefinedSymbol(HeaderBlock, HS.Offset, HS.Name, HeaderBlock.getSize(),
+                       jitlink::Linkage::Strong, jitlink::Scope::Default, false,
+                       true);
+}
+
+jitlink::Block &
+SimpleMachOHeaderMU::createHeaderBlock(JITDylib &JD, jitlink::LinkGraph &G,
+                                       jitlink::Section &HeaderSection) {
+  switch (MOP.getExecutionSession().getTargetTriple().getArch()) {
+  case Triple::aarch64:
+  case Triple::x86_64:
+    return createTrivialHeaderBlock<MachO64LE>(MOP, G, HeaderSection);
+  default:
+    llvm_unreachable("Unsupported architecture");
+  }
+}
+
+MaterializationUnit::Interface SimpleMachOHeaderMU::createHeaderInterface(
+    MachOPlatform &MOP, const SymbolStringPtr &HeaderStartSymbol) {
+  SymbolFlagsMap HeaderSymbolFlags;
+
+  HeaderSymbolFlags[HeaderStartSymbol] = JITSymbolFlags::Exported;
+  for (auto &HS : AdditionalHeaderSymbols)
+    HeaderSymbolFlags[MOP.getExecutionSession().intern(HS.Name)] =
+        JITSymbolFlags::Exported;
+
+  return MaterializationUnit::Interface(std::move(HeaderSymbolFlags),
+                                        HeaderStartSymbol);
+}
+
+MachOHeaderInfo getMachOHeaderInfoFromTriple(const Triple &TT) {
+  switch (TT.getArch()) {
+  case Triple::aarch64:
+    return {/* PageSize   = */ 16 * 1024,
+            /* CPUType    = */ MachO::CPU_TYPE_ARM64,
+            /* CPUSubType = */ MachO::CPU_SUBTYPE_ARM64_ALL};
+  case Triple::x86_64:
+    return {/* PageSize   = */ 4 * 1024,
+            /* CPUType    = */ MachO::CPU_TYPE_X86_64,
+            /* CPUSubType = */ MachO::CPU_SUBTYPE_X86_64_ALL};
+  default:
+    llvm_unreachable("Unrecognized architecture");
+  }
+}
+
 } // End namespace orc.
 } // End namespace llvm.