[llvm] [ORC] Add absoluteSymbolsObjectLayer to expose absolute symbols to platform (PR #77008)

Ben Langmuir via llvm-commits llvm-commits at lists.llvm.org
Thu Jan 4 13:28:19 PST 2024


https://github.com/benlangmuir created https://github.com/llvm/llvm-project/pull/77008

Adds a variant of absoluteSymbols that creates a LinkGraph and exposes its symbols to the platform/orc runtime. This allows e.g. using __orc_rt_run_program to run a precompiled function that was found via dlsym. Ideally we would use this in llvm-jitlink's own search generator, but it will require more work to align with the Process/Platform JITDylib split, so not handled here.

As part of this change we need to handle LinkGraphs that only have absolute symbols.

>From eebaa0cc7eee03daf616b0c6b594641f10afbae4 Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir at apple.com>
Date: Wed, 3 Jan 2024 15:52:33 -0800
Subject: [PATCH] [ORC] Add absoluteSymbolsObjectLayer to expose absolute
 symbols to platform

Adds a variant of absoluteSymbols that creates a LinkGraph and exposes
its symbols to the platform/orc runtime. This allows e.g. using
__orc_rt_run_program to run a precompiled function that was found via
dlsym. Ideally we would use this in llvm-jitlink's own search generator,
but it will require more work to align with the Process/Platform JITDylib
split, so not handled here.

As part of this change we need to handle LinkGraphs that only have
absolute symbols.
---
 bolt/lib/Rewrite/JITLinkLinker.cpp            |  3 +-
 .../Orc/EPCDynamicLibrarySearchGenerator.h    | 21 ++--
 .../llvm/ExecutionEngine/Orc/ExecutionUtils.h | 17 +++-
 .../ExecutionEngine/Orc/ObjectLinkingLayer.h  | 11 +++
 .../JITLink/JITLinkGeneric.cpp                | 34 ++++---
 .../Orc/EPCDynamicLibrarySearchGenerator.cpp  |  9 +-
 .../ExecutionEngine/Orc/ExecutionUtils.cpp    | 12 ++-
 .../lib/ExecutionEngine/Orc/MachOPlatform.cpp |  2 +
 .../Orc/ObjectLinkingLayer.cpp                | 60 +++++++++++-
 .../Orc/ObjectLinkingLayerTest.cpp            | 95 +++++++++++++++++++
 10 files changed, 227 insertions(+), 37 deletions(-)

diff --git a/bolt/lib/Rewrite/JITLinkLinker.cpp b/bolt/lib/Rewrite/JITLinkLinker.cpp
index 10a3b6a0407ff3..66e129bf1d05db 100644
--- a/bolt/lib/Rewrite/JITLinkLinker.cpp
+++ b/bolt/lib/Rewrite/JITLinkLinker.cpp
@@ -173,7 +173,8 @@ struct JITLinkLinker::Context : jitlink::JITLinkContext {
 
   void notifyFinalized(
       jitlink::JITLinkMemoryManager::FinalizedAlloc Alloc) override {
-    Linker.Allocs.push_back(std::move(Alloc));
+    if (Alloc)
+      Linker.Allocs.push_back(std::move(Alloc));
     ++Linker.MM->ObjectsLoaded;
   }
 };
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h b/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
index 63797edec89e61..063f9be9b751d0 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
@@ -25,6 +25,8 @@ class ExecutorProcessControl;
 class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
 public:
   using SymbolPredicate = unique_function<bool(const SymbolStringPtr &)>;
+  using AbsoluteSymbolsMUFactory =
+      unique_function<std::unique_ptr<MaterializationUnit>(SymbolMap)>;
 
   /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
   /// library with the given handle.
@@ -32,24 +34,28 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
   /// If the Allow predicate is given then only symbols matching the predicate
   /// will be searched for. If the predicate is not given then all symbols will
   /// be searched for.
-  EPCDynamicLibrarySearchGenerator(ExecutionSession &ES,
-                                   tpctypes::DylibHandle H,
-                                   SymbolPredicate Allow = SymbolPredicate())
-      : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)) {}
+  EPCDynamicLibrarySearchGenerator(
+      ExecutionSession &ES, tpctypes::DylibHandle H,
+      SymbolPredicate Allow = SymbolPredicate(),
+      AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols)
+      : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)),
+        MUFactory(std::move(MUFactory)) {}
 
   /// Permanently loads the library at the given path and, on success, returns
   /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
   /// in the library. On failure returns the reason the library failed to load.
   static Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
   Load(ExecutionSession &ES, const char *LibraryPath,
-       SymbolPredicate Allow = SymbolPredicate());
+       SymbolPredicate Allow = SymbolPredicate(),
+       AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols);
 
   /// Creates a EPCDynamicLibrarySearchGenerator that searches for symbols in
   /// the target process.
   static Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
   GetForTargetProcess(ExecutionSession &ES,
-                      SymbolPredicate Allow = SymbolPredicate()) {
-    return Load(ES, nullptr, std::move(Allow));
+                      SymbolPredicate Allow = SymbolPredicate(),
+                      AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols) {
+    return Load(ES, nullptr, std::move(Allow), std::move(MUFactory));
   }
 
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
@@ -60,6 +66,7 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
   ExecutorProcessControl &EPC;
   tpctypes::DylibHandle H;
   SymbolPredicate Allow;
+  AbsoluteSymbolsMUFactory MUFactory;
 };
 
 } // end namespace orc
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
index 6a43376a5bd909..e3c4f22f23e29a 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -216,6 +216,8 @@ class ItaniumCXAAtExitSupport {
 class DynamicLibrarySearchGenerator : public DefinitionGenerator {
 public:
   using SymbolPredicate = std::function<bool(const SymbolStringPtr &)>;
+  using AbsoluteSymbolsMUFactory =
+      unique_function<std::unique_ptr<MaterializationUnit>(SymbolMap)>;
 
   /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
   /// given sys::DynamicLibrary.
@@ -223,22 +225,26 @@ class DynamicLibrarySearchGenerator : public DefinitionGenerator {
   /// If the Allow predicate is given then only symbols matching the predicate
   /// will be searched for. If the predicate is not given then all symbols will
   /// be searched for.
-  DynamicLibrarySearchGenerator(sys::DynamicLibrary Dylib, char GlobalPrefix,
-                                SymbolPredicate Allow = SymbolPredicate());
+  DynamicLibrarySearchGenerator(
+      sys::DynamicLibrary Dylib, char GlobalPrefix,
+      SymbolPredicate Allow = SymbolPredicate(),
+      AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols);
 
   /// Permanently loads the library at the given path and, on success, returns
   /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
   /// in the library. On failure returns the reason the library failed to load.
   static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   Load(const char *FileName, char GlobalPrefix,
-       SymbolPredicate Allow = SymbolPredicate());
+       SymbolPredicate Allow = SymbolPredicate(),
+       AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols);
 
   /// Creates a DynamicLibrarySearchGenerator that searches for symbols in
   /// the current process.
   static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   GetForCurrentProcess(char GlobalPrefix,
-                       SymbolPredicate Allow = SymbolPredicate()) {
-    return Load(nullptr, GlobalPrefix, std::move(Allow));
+                       SymbolPredicate Allow = SymbolPredicate(),
+                       AbsoluteSymbolsMUFactory MUFactory = absoluteSymbols) {
+    return Load(nullptr, GlobalPrefix, std::move(Allow), std::move(MUFactory));
   }
 
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
@@ -248,6 +254,7 @@ class DynamicLibrarySearchGenerator : public DefinitionGenerator {
 private:
   sys::DynamicLibrary Dylib;
   SymbolPredicate Allow;
+  AbsoluteSymbolsMUFactory MUFactory;
   char GlobalPrefix;
 };
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
index 34f2e0789462c9..a0557de49e98dd 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
@@ -128,6 +128,11 @@ class ObjectLinkingLayer : public RTTIExtends<ObjectLinkingLayer, ObjectLayer>,
     return *this;
   }
 
+  /// Create a \c MaterializationUnit for the given \c LinkGraph using this
+  /// object layer.
+  std::unique_ptr<MaterializationUnit>
+  createLinkGraphMaterializationUnit(std::unique_ptr<jitlink::LinkGraph> G);
+
   /// Add a LinkGraph to the JITDylib targeted by the given tracker.
   Error add(ResourceTrackerSP, std::unique_ptr<jitlink::LinkGraph> G);
 
@@ -223,6 +228,12 @@ class EHFrameRegistrationPlugin : public ObjectLinkingLayer::Plugin {
   DenseMap<ResourceKey, std::vector<ExecutorAddrRange>> EHFrameRanges;
 };
 
+/// Create a \c MaterializationUnit defining the given symbols. Similar to
+/// \c absoluteSymbols, but providing the symbols through the
+/// \c ObjectLinkingLayer so that the \c Platform handles them.
+std::unique_ptr<MaterializationUnit>
+absoluteSymbolsObjectLayer(ObjectLinkingLayer &OLL, SymbolMap Symbols);
+
 } // end namespace orc
 } // end namespace llvm
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index 5361272ae79ef7..5494697119fb94 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -48,15 +48,22 @@ void JITLinkerBase::linkPhase1(std::unique_ptr<JITLinkerBase> Self) {
   if (auto Err = runPasses(Passes.PostPrunePasses))
     return Ctx->notifyFailed(std::move(Err));
 
-  Ctx->getMemoryManager().allocate(
-      Ctx->getJITLinkDylib(), *G,
-      [S = std::move(Self)](AllocResult AR) mutable {
-        // FIXME: Once MSVC implements c++17 order of evaluation rules for calls
-        // this can be simplified to
-        //          S->linkPhase2(std::move(S), std::move(AR));
-        auto *TmpSelf = S.get();
-        TmpSelf->linkPhase2(std::move(S), std::move(AR));
-      });
+  auto LinkPhase2 = [S = std::move(Self)](AllocResult AR) mutable {
+    // FIXME: Once MSVC implements c++17 order of evaluation rules for calls
+    // this can be simplified to
+    //          S->linkPhase2(std::move(S), std::move(AR));
+    auto *TmpSelf = S.get();
+    TmpSelf->linkPhase2(std::move(S), std::move(AR));
+  };
+
+  if (G->allocActions().empty() && llvm::all_of(G->sections(), [](Section &S) {
+        return S.getMemLifetime() == orc::MemLifetime::NoAlloc;
+      })) {
+    LinkPhase2(nullptr);
+  } else {
+    Ctx->getMemoryManager().allocate(Ctx->getJITLinkDylib(), *G,
+                                     std::move(LinkPhase2));
+  }
 }
 
 void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
@@ -163,13 +170,18 @@ void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self,
   if (auto Err = runPasses(Passes.PostFixupPasses))
     return abandonAllocAndBailOut(std::move(Self), std::move(Err));
 
-  Alloc->finalize([S = std::move(Self)](FinalizeResult FR) mutable {
+  auto LinkPhase4 = [S = std::move(Self)](FinalizeResult FR) mutable {
     // FIXME: Once MSVC implements c++17 order of evaluation rules for calls
     // this can be simplified to
     //          S->linkPhase2(std::move(S), std::move(AR));
     auto *TmpSelf = S.get();
     TmpSelf->linkPhase4(std::move(S), std::move(FR));
-  });
+  };
+
+  if (Alloc)
+    Alloc->finalize(std::move(LinkPhase4));
+  else
+    LinkPhase4(FinalizeResult{JITLinkMemoryManager::FinalizedAlloc{}});
 }
 
 void JITLinkerBase::linkPhase4(std::unique_ptr<JITLinkerBase> Self,
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
index 2e2e7a9c5f32e8..69d52511281cb3 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
@@ -14,13 +14,14 @@ namespace orc {
 Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
 EPCDynamicLibrarySearchGenerator::Load(ExecutionSession &ES,
                                        const char *LibraryPath,
-                                       SymbolPredicate Allow) {
+                                       SymbolPredicate Allow,
+                                       AbsoluteSymbolsMUFactory MUFactory) {
   auto Handle = ES.getExecutorProcessControl().loadDylib(LibraryPath);
   if (!Handle)
     return Handle.takeError();
 
-  return std::make_unique<EPCDynamicLibrarySearchGenerator>(ES, *Handle,
-                                                            std::move(Allow));
+  return std::make_unique<EPCDynamicLibrarySearchGenerator>(
+      ES, *Handle, std::move(Allow), std::move(MUFactory));
 }
 
 Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
@@ -62,7 +63,7 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     return Error::success();
 
   // Define resolved symbols.
-  return JD.define(absoluteSymbols(std::move(NewSymbols)));
+  return JD.define(MUFactory(std::move(NewSymbols)));
 }
 
 } // end namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
index 8d5608cc4d4cfd..cb7333d205665f 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -218,19 +218,21 @@ void ItaniumCXAAtExitSupport::runAtExits(void *DSOHandle) {
 }
 
 DynamicLibrarySearchGenerator::DynamicLibrarySearchGenerator(
-    sys::DynamicLibrary Dylib, char GlobalPrefix, SymbolPredicate Allow)
+    sys::DynamicLibrary Dylib, char GlobalPrefix, SymbolPredicate Allow,
+    AbsoluteSymbolsMUFactory MUFactory)
     : Dylib(std::move(Dylib)), Allow(std::move(Allow)),
-      GlobalPrefix(GlobalPrefix) {}
+      MUFactory(std::move(MUFactory)), GlobalPrefix(GlobalPrefix) {}
 
 Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
 DynamicLibrarySearchGenerator::Load(const char *FileName, char GlobalPrefix,
-                                    SymbolPredicate Allow) {
+                                    SymbolPredicate Allow,
+                                    AbsoluteSymbolsMUFactory MUFactory) {
   std::string ErrMsg;
   auto Lib = sys::DynamicLibrary::getPermanentLibrary(FileName, &ErrMsg);
   if (!Lib.isValid())
     return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
   return std::make_unique<DynamicLibrarySearchGenerator>(
-      std::move(Lib), GlobalPrefix, std::move(Allow));
+      std::move(Lib), GlobalPrefix, std::move(Allow), std::move(MUFactory));
 }
 
 Error DynamicLibrarySearchGenerator::tryToGenerate(
@@ -261,7 +263,7 @@ Error DynamicLibrarySearchGenerator::tryToGenerate(
   if (NewSymbols.empty())
     return Error::success();
 
-  return JD.define(absoluteSymbols(std::move(NewSymbols)));
+  return JD.define(MUFactory(std::move(NewSymbols)));
 }
 
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
diff --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
index 9057300bf04352..6c17f14aa4c7c5 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -1608,6 +1608,8 @@ Error MachOPlatform::MachOPlatformPlugin::prepareSymbolTableRegistration(
     SmallVector<jitlink::Symbol *> SymsToProcess;
     for (auto *Sym : G.defined_symbols())
       SymsToProcess.push_back(Sym);
+    for (auto *Sym : G.absolute_symbols())
+      SymsToProcess.push_back(Sym);
 
     for (auto *Sym : SymsToProcess) {
       if (!Sym->hasName())
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 3d77f82e65696a..db0157db896098 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -12,7 +12,9 @@
 #include "llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h"
+#include "llvm/Support/Endian.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include <atomic>
 #include <string>
 #include <vector>
 
@@ -93,15 +95,20 @@ class LinkGraphMaterializationUnit : public MaterializationUnit {
 
     Interface LGI;
 
-    for (auto *Sym : G.defined_symbols()) {
+    auto AddSymbol = [&](Symbol *Sym) {
       // Skip local symbols.
       if (Sym->getScope() == Scope::Local)
-        continue;
+        return;
       assert(Sym->hasName() && "Anonymous non-local symbol?");
 
       LGI.SymbolFlags[ES.intern(Sym->getName())] =
           getJITSymbolFlagsForSymbol(*Sym);
-    }
+    };
+
+    for (auto *Sym : G.defined_symbols())
+      AddSymbol(Sym);
+    for (auto *Sym : G.absolute_symbols())
+      AddSymbol(Sym);
 
     if (hasInitializerSection(G))
       LGI.InitSymbol = makeInitSymbol(ES, G);
@@ -654,10 +661,16 @@ ObjectLinkingLayer::~ObjectLinkingLayer() {
   getExecutionSession().deregisterResourceManager(*this);
 }
 
+std::unique_ptr<MaterializationUnit>
+ObjectLinkingLayer::createLinkGraphMaterializationUnit(
+    std::unique_ptr<jitlink::LinkGraph> G) {
+  return LinkGraphMaterializationUnit::Create(*this, std::move(G));
+}
+
 Error ObjectLinkingLayer::add(ResourceTrackerSP RT,
                               std::unique_ptr<LinkGraph> G) {
   auto &JD = RT->getJITDylib();
-  return JD.define(LinkGraphMaterializationUnit::Create(*this, std::move(G)),
+  return JD.define(createLinkGraphMaterializationUnit(std::move(G)),
                    std::move(RT));
 }
 
@@ -705,6 +718,9 @@ Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
   if (Err)
     return Err;
 
+  if (!FA)
+    return Error::success();
+
   return MR.withResourceKeyDo(
       [&](ResourceKey K) { Allocs[K].push_back(std::move(FA)); });
 }
@@ -850,5 +866,41 @@ void EHFrameRegistrationPlugin::notifyTransferringResources(
   }
 }
 
+std::unique_ptr<MaterializationUnit>
+absoluteSymbolsObjectLayer(ObjectLinkingLayer &OLL, SymbolMap Symbols) {
+  const Triple &TT = OLL.getExecutionSession().getTargetTriple();
+  unsigned PointerSize;
+  endianness Endianness =
+      TT.isLittleEndian() ? endianness::little : endianness::big;
+  switch (TT.getArch()) {
+  case Triple::aarch64:
+  case llvm::Triple::riscv64:
+  case Triple::x86_64:
+    PointerSize = 8;
+    break;
+  case llvm::Triple::arm:
+  case llvm::Triple::riscv32:
+  case llvm::Triple::x86:
+    PointerSize = 4;
+    break;
+  default:
+    llvm::report_fatal_error("unhandled target architecture");
+  }
+
+  static std::atomic<uint64_t> Counter = {0};
+  auto Index = Counter.fetch_add(1, std::memory_order_relaxed);
+  auto G = std::make_unique<LinkGraph>(
+      "<Absolute Symbols " + std::to_string(Index) + ">", TT, PointerSize,
+      Endianness, /*GetEdgeKindName=*/nullptr);
+  for (auto &[Name, Def] : Symbols) {
+    auto &Sym =
+        G->addAbsoluteSymbol(*Name, Def.getAddress(), /*Size=*/0,
+                             Linkage::Strong, Scope::Default, /*IsLive=*/true);
+    Sym.setCallable(Def.getFlags().isCallable());
+  }
+
+  return OLL.createLinkGraphMaterializationUnit(std::move(G));
+}
+
 } // End namespace orc.
 } // End namespace llvm.
diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 91659240c9d6dc..a19b303df6ab81 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -10,6 +10,11 @@
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
+#include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
@@ -173,4 +178,94 @@ TEST_F(ObjectLinkingLayerTest, HandleErrorDuringPostAllocationPass) {
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor"), Failed());
 }
 
+TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
+  class TestEPC : public UnsupportedExecutorProcessControl {
+  public:
+    TestEPC()
+        : UnsupportedExecutorProcessControl(nullptr, nullptr,
+                                            "x86_64-apple-darwin") {}
+
+    Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override {
+      return ExecutorAddr::fromPtr((void *)nullptr);
+    }
+
+    Expected<std::vector<tpctypes::LookupResult>>
+    lookupSymbols(ArrayRef<LookupRequest> Request) override {
+      std::vector<ExecutorSymbolDef> Result;
+      EXPECT_EQ(Request.size(), 1u);
+      for (auto &LR : Request) {
+        EXPECT_EQ(LR.Symbols.size(), 1u);
+        for (auto &Sym : LR.Symbols) {
+          if (*Sym.first == "_testFunc") {
+            ExecutorSymbolDef Def{ExecutorAddr::fromPtr((void *)0x1000),
+                                  JITSymbolFlags::Exported};
+            Result.push_back(Def);
+          } else {
+            ADD_FAILURE() << "unexpected symbol request " << *Sym.first;
+          }
+        }
+      }
+      return std::vector<tpctypes::LookupResult>{1, Result};
+    }
+  };
+
+  ExecutionSession ES{std::make_unique<TestEPC>()};
+  JITDylib &JD = ES.createBareJITDylib("main");
+  ObjectLinkingLayer ObjLinkingLayer{
+      ES, std::make_unique<InProcessMemoryManager>(4096)};
+
+  auto G = EPCDynamicLibrarySearchGenerator::GetForTargetProcess(
+      ES, {}, [&](SymbolMap Syms) {
+        return absoluteSymbolsObjectLayer(ObjLinkingLayer, std::move(Syms));
+      });
+  ASSERT_THAT_EXPECTED(G, Succeeded());
+  JD.addGenerator(std::move(*G));
+
+  class CheckDefs : public ObjectLinkingLayer::Plugin {
+  public:
+    ~CheckDefs() { EXPECT_TRUE(SawSymbolDef); }
+
+    void modifyPassConfig(MaterializationResponsibility &MR,
+                          jitlink::LinkGraph &G,
+                          jitlink::PassConfiguration &Config) override {
+      Config.PostAllocationPasses.push_back([this](LinkGraph &G) {
+        unsigned SymCount = 0;
+        for (Symbol *Sym : G.absolute_symbols()) {
+          SymCount += 1;
+          if (!Sym->hasName()) {
+            ADD_FAILURE() << "unexpected unnamed symbol";
+            continue;
+          }
+          if (Sym->getName() == "_testFunc")
+            SawSymbolDef = true;
+          else
+            ADD_FAILURE() << "unexpected symbol " << Sym->getName();
+        }
+        EXPECT_EQ(SymCount, 1u);
+        return Error::success();
+      });
+    }
+
+    Error notifyFailed(MaterializationResponsibility &MR) override {
+      return Error::success();
+    }
+
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+      return Error::success();
+    }
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                     ResourceKey SrcKey) override {
+      llvm_unreachable("unexpected resource transfer");
+    }
+
+  private:
+    bool SawSymbolDef = false;
+  };
+
+  ObjLinkingLayer.addPlugin(std::make_unique<CheckDefs>());
+
+  EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_testFunc"), Succeeded());
+  EXPECT_THAT_ERROR(ES.endSession(), Succeeded());
+}
+
 } // end anonymous namespace



More information about the llvm-commits mailing list