[ORC] Move EPC load-dylib and lookup operations into their own class.

This keeps common operations together, and should make it easier to write
re-usable dylib managers in the future (e.g. a DylibManager that uses
the EPC's remote-execution APIs to implement load and lookup).

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h

@@ -15,6 +15,7 @@
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/DylibManager.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
@@ -32,7 +33,6 @@ namespace llvm {
 namespace orc {
 
 class ExecutionSession;
-class SymbolLookupSet;
 
 /// ExecutorProcessControl supports interaction with a JIT target process.
 class ExecutorProcessControl {
@@ -172,14 +172,6 @@ class ExecutorProcessControl {
     }
   };
 
-  /// A pair of a dylib and a set of symbols to be looked up.
-  struct LookupRequest {
-    LookupRequest(tpctypes::DylibHandle Handle, const SymbolLookupSet &Symbols)
-        : Handle(Handle), Symbols(Symbols) {}
-    tpctypes::DylibHandle Handle;
-    const SymbolLookupSet &Symbols;
-  };
-
   /// Contains the address of the dispatch function and context that the ORC
   /// runtime can use to call functions in the JIT.
   struct JITDispatchInfo {
@@ -229,6 +221,12 @@ class ExecutorProcessControl {
     return *MemMgr;
   }
 
+  /// Return the DylibManager for the target process.
+  DylibManager &getDylibMgr() const {
+    assert(DylibMgr && "No DylibMgr object set");
+    return *DylibMgr;
+  }
+
   /// Returns the bootstrap map.
   const StringMap<std::vector<char>> &getBootstrapMap() const {
     return BootstrapMap;
@@ -277,38 +275,6 @@ class ExecutorProcessControl {
     return Error::success();
   }
 
-  /// Load the dynamic library at the given path and return a handle to it.
-  /// If LibraryPath is null this function will return the global handle for
-  /// the target process.
-  virtual Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) = 0;
-
-  /// Search for symbols in the target process.
-  ///
-  /// The result of the lookup is a 2-dimensional array of target addresses
-  /// that correspond to the lookup order. If a required symbol is not
-  /// found then this method will return an error. If a weakly referenced
-  /// symbol is not found then it be assigned a '0' value.
-  Expected<std::vector<tpctypes::LookupResult>>
-  lookupSymbols(ArrayRef<LookupRequest> Request) {
-    std::promise<MSVCPExpected<std::vector<tpctypes::LookupResult>>> RP;
-    auto RF = RP.get_future();
-    lookupSymbolsAsync(Request,
-                       [&RP](auto Result) { RP.set_value(std::move(Result)); });
-    return RF.get();
-  }
-
-  using SymbolLookupCompleteFn =
-      unique_function<void(Expected<std::vector<tpctypes::LookupResult>>)>;
-
-  /// Search for symbols in the target process.
-  ///
-  /// The result of the lookup is a 2-dimensional array of target addresses
-  /// that correspond to the lookup order. If a required symbol is not
-  /// found then this method will return an error. If a weakly referenced
-  /// symbol is not found then it be assigned a '0' value.
-  virtual void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
-                                  SymbolLookupCompleteFn F) = 0;
-
   /// Run function with a main-like signature.
   virtual Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                                       ArrayRef<std::string> Args) = 0;
@@ -426,6 +392,7 @@ class ExecutorProcessControl {
   JITDispatchInfo JDI;
   MemoryAccess *MemAccess = nullptr;
   jitlink::JITLinkMemoryManager *MemMgr = nullptr;
+  DylibManager *DylibMgr = nullptr;
   StringMap<std::vector<char>> BootstrapMap;
   StringMap<ExecutorAddr> BootstrapSymbols;
 };
@@ -474,15 +441,6 @@ class UnsupportedExecutorProcessControl : public ExecutorProcessControl,
     this->MemAccess = this;
   }
 
-  Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override {
-    llvm_unreachable("Unsupported");
-  }
-
-  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
-                          SymbolLookupCompleteFn F) override {
-    llvm_unreachable("Unsupported");
-  }
-
   Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                               ArrayRef<std::string> Args) override {
     llvm_unreachable("Unsupported");
@@ -507,7 +465,8 @@ class UnsupportedExecutorProcessControl : public ExecutorProcessControl,
 
 /// A ExecutorProcessControl implementation targeting the current process.
 class SelfExecutorProcessControl : public ExecutorProcessControl,
-                                   private InProcessMemoryAccess {
+                                   private InProcessMemoryAccess,
+                                   private DylibManager {
 public:
   SelfExecutorProcessControl(
       std::shared_ptr<SymbolStringPool> SSP, std::unique_ptr<TaskDispatcher> D,
@@ -524,11 +483,6 @@ class SelfExecutorProcessControl : public ExecutorProcessControl,
          std::unique_ptr<TaskDispatcher> D = nullptr,
          std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr = nullptr);
 
-  Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
-
-  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
-                          SymbolLookupCompleteFn F) override;
-
   Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                               ArrayRef<std::string> Args) override;
 
@@ -547,6 +501,11 @@ class SelfExecutorProcessControl : public ExecutorProcessControl,
   jitDispatchViaWrapperFunctionManager(void *Ctx, const void *FnTag,
                                        const char *Data, size_t Size);
 
+  Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
+
+  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                          SymbolLookupCompleteFn F) override;
+
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   char GlobalManglingPrefix = 0;
 };

llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h

@@ -29,7 +29,8 @@ namespace llvm {
 namespace orc {
 
 class SimpleRemoteEPC : public ExecutorProcessControl,
-                        public SimpleRemoteEPCTransportClient {
+                        public SimpleRemoteEPCTransportClient,
+                        private DylibManager {
 public:
   /// A setup object containing callbacks to construct a memory manager and
   /// memory access object. Both are optional. If not specified,
@@ -69,11 +70,6 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   SimpleRemoteEPC &operator=(SimpleRemoteEPC &&) = delete;
   ~SimpleRemoteEPC();
 
-  Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
-
-  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
-                          SymbolLookupCompleteFn F) override;
-
   Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                               ArrayRef<std::string> Args) override;
 
@@ -96,7 +92,9 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
 private:
   SimpleRemoteEPC(std::shared_ptr<SymbolStringPool> SSP,
                   std::unique_ptr<TaskDispatcher> D)
-    : ExecutorProcessControl(std::move(SSP), std::move(D)) {}
+      : ExecutorProcessControl(std::move(SSP), std::move(D)) {
+    this->DylibMgr = this;
+  }
 
   static Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
   createDefaultMemoryManager(SimpleRemoteEPC &SREPC);
@@ -119,6 +117,11 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   uint64_t getNextSeqNo() { return NextSeqNo++; }
   void releaseSeqNo(uint64_t SeqNo) {}
 
+  Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
+
+  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                          SymbolLookupCompleteFn F) override;
+
   using PendingCallWrapperResultsMap =
     DenseMap<uint64_t, IncomingWFRHandler>;
 
@@ -131,7 +134,7 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   std::unique_ptr<MemoryAccess> OwnedMemAccess;
 
-  std::unique_ptr<EPCGenericDylibManager> DylibMgr;
+  std::unique_ptr<EPCGenericDylibManager> EPCDylibMgr;
   ExecutorAddr RunAsMainAddr;
   ExecutorAddr RunAsVoidFunctionAddr;
   ExecutorAddr RunAsIntFunctionAddr;

llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp

@@ -22,7 +22,7 @@ Expected<std::unique_ptr<EPCDebugObjectRegistrar>> createJITLoaderGDBRegistrar(
   auto &EPC = ES.getExecutorProcessControl();
 
   if (!RegistrationFunctionDylib) {
-    if (auto D = EPC.loadDylib(nullptr))
+    if (auto D = EPC.getDylibMgr().loadDylib(nullptr))
       RegistrationFunctionDylib = *D;
     else
       return D.takeError();
@@ -36,8 +36,8 @@ Expected<std::unique_ptr<EPCDebugObjectRegistrar>> createJITLoaderGDBRegistrar(
   SymbolLookupSet RegistrationSymbols;
   RegistrationSymbols.add(RegisterFn);
 
-  auto Result =
-      EPC.lookupSymbols({{*RegistrationFunctionDylib, RegistrationSymbols}});
+  auto Result = EPC.getDylibMgr().lookupSymbols(
+      {{*RegistrationFunctionDylib, RegistrationSymbols}});
   if (!Result)
     return Result.takeError();
 

llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp

@@ -19,7 +19,8 @@ Expected<std::unique_ptr<EPCDynamicLibrarySearchGenerator>>
 EPCDynamicLibrarySearchGenerator::Load(
     ExecutionSession &ES, const char *LibraryPath, SymbolPredicate Allow,
     AddAbsoluteSymbolsFn AddAbsoluteSymbols) {
-  auto Handle = ES.getExecutorProcessControl().loadDylib(LibraryPath);
+  auto Handle =
+      ES.getExecutorProcessControl().getDylibMgr().loadDylib(LibraryPath);
   if (!Handle)
     return Handle.takeError();
 
@@ -48,10 +49,11 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     LookupSymbols.add(KV.first, SymbolLookupFlags::WeaklyReferencedSymbol);
   }
 
-  ExecutorProcessControl::LookupRequest Request(H, LookupSymbols);
+  DylibManager::LookupRequest Request(H, LookupSymbols);
   // Copy-capture LookupSymbols, since LookupRequest keeps a reference.
-  EPC.lookupSymbolsAsync(Request, [this, &JD, LS = std::move(LS),
-                                   LookupSymbols](auto Result) mutable {
+  EPC.getDylibMgr().lookupSymbolsAsync(Request, [this, &JD, LS = std::move(LS),
+                                                 LookupSymbols](
+                                                    auto Result) mutable {
     if (!Result) {
       LLVM_DEBUG({
         dbgs() << "EPCDynamicLibrarySearchGenerator lookup failed due to error";

llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp

@@ -41,17 +41,17 @@ class TrivialSPSSequenceSerialization<SPSRemoteSymbolLookupSetElement,
 
 template <>
 class SPSSerializationTraits<SPSRemoteSymbolLookup,
-                             ExecutorProcessControl::LookupRequest> {
+                             DylibManager::LookupRequest> {
   using MemberSerialization =
       SPSArgList<SPSExecutorAddr, SPSRemoteSymbolLookupSet>;
 
 public:
-  static size_t size(const ExecutorProcessControl::LookupRequest &LR) {
+  static size_t size(const DylibManager::LookupRequest &LR) {
     return MemberSerialization::size(ExecutorAddr(LR.Handle), LR.Symbols);
   }
 
   static bool serialize(SPSOutputBuffer &OB,
-                        const ExecutorProcessControl::LookupRequest &LR) {
+                        const DylibManager::LookupRequest &LR) {
     return MemberSerialization::serialize(OB, ExecutorAddr(LR.Handle),
                                           LR.Symbols);
   }

llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp

@@ -21,6 +21,8 @@
 namespace llvm {
 namespace orc {
 
+DylibManager::~DylibManager() = default;
+
 ExecutorProcessControl::MemoryAccess::~MemoryAccess() = default;
 
 ExecutorProcessControl::~ExecutorProcessControl() = default;
@@ -41,6 +43,7 @@ SelfExecutorProcessControl::SelfExecutorProcessControl(
   this->PageSize = PageSize;
   this->MemMgr = OwnedMemMgr.get();
   this->MemAccess = this;
+  this->DylibMgr = this;
   this->JDI = {ExecutorAddr::fromPtr(jitDispatchViaWrapperFunctionManager),
                ExecutorAddr::fromPtr(this)};
   if (this->TargetTriple.isOSBinFormatMachO())
@@ -86,7 +89,7 @@ SelfExecutorProcessControl::loadDylib(const char *DylibPath) {
 
 void SelfExecutorProcessControl::lookupSymbolsAsync(
     ArrayRef<LookupRequest> Request,
-    ExecutorProcessControl::SymbolLookupCompleteFn Complete) {
+    DylibManager::SymbolLookupCompleteFn Complete) {
   std::vector<tpctypes::LookupResult> R;
 
   for (auto &Elem : Request) {

llvm/lib/ExecutionEngine/Orc/LookupAndRecordAddrs.cpp

@@ -60,8 +60,8 @@ Error lookupAndRecordAddrs(
   for (auto &KV : Pairs)
     Symbols.add(KV.first, LookupFlags);
 
-  ExecutorProcessControl::LookupRequest LR(H, Symbols);
-  auto Result = EPC.lookupSymbols(LR);
+  DylibManager::LookupRequest LR(H, Symbols);
+  auto Result = EPC.getDylibMgr().lookupSymbols(LR);
   if (!Result)
     return Result.takeError();
 

llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp

@@ -26,17 +26,17 @@ SimpleRemoteEPC::~SimpleRemoteEPC() {
 
 Expected<tpctypes::DylibHandle>
 SimpleRemoteEPC::loadDylib(const char *DylibPath) {
-  return DylibMgr->open(DylibPath, 0);
+  return EPCDylibMgr->open(DylibPath, 0);
 }
 
 /// Async helper to chain together calls to DylibMgr::lookupAsync to fulfill all
 /// all the requests.
 /// FIXME: The dylib manager should support multiple LookupRequests natively.
 static void
 lookupSymbolsAsyncHelper(EPCGenericDylibManager &DylibMgr,
-                         ArrayRef<SimpleRemoteEPC::LookupRequest> Request,
+                         ArrayRef<DylibManager::LookupRequest> Request,
                          std::vector<tpctypes::LookupResult> Result,
-                         SimpleRemoteEPC::SymbolLookupCompleteFn Complete) {
+                         DylibManager::SymbolLookupCompleteFn Complete) {
   if (Request.empty())
     return Complete(std::move(Result));
 
@@ -59,7 +59,7 @@ lookupSymbolsAsyncHelper(EPCGenericDylibManager &DylibMgr,
 
 void SimpleRemoteEPC::lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
                                          SymbolLookupCompleteFn Complete) {
-  lookupSymbolsAsyncHelper(*DylibMgr, Request, {}, std::move(Complete));
+  lookupSymbolsAsyncHelper(*EPCDylibMgr, Request, {}, std::move(Complete));
 }
 
 Expected<int32_t> SimpleRemoteEPC::runAsMain(ExecutorAddr MainFnAddr,
@@ -357,7 +357,7 @@ Error SimpleRemoteEPC::setup(Setup S) {
 
   if (auto DM =
           EPCGenericDylibManager::CreateWithDefaultBootstrapSymbols(*this))
-    DylibMgr = std::make_unique<EPCGenericDylibManager>(std::move(*DM));
+    EPCDylibMgr = std::make_unique<EPCGenericDylibManager>(std::move(*DM));
   else
     return DM.takeError();
 

llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

@@ -255,11 +255,14 @@ TEST_F(ObjectLinkingLayerTest, AddAndRemovePlugins) {
 }
 
 TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
-  class TestEPC : public UnsupportedExecutorProcessControl {
+  class TestEPC : public UnsupportedExecutorProcessControl,
+                  public DylibManager {
   public:
     TestEPC()
         : UnsupportedExecutorProcessControl(nullptr, nullptr,
-                                            "x86_64-apple-darwin") {}
+                                            "x86_64-apple-darwin") {
+      this->DylibMgr = this;
+    }
 
     Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override {
       return ExecutorAddr::fromPtr((void *)nullptr);