[ORC] Add bootstrap symbols to ExecutorProcessControl.

Bootstrap symbols are symbols whose addresses may be required to bootstrap
the rest of the JIT. The bootstrap symbols map generalizes the existing
JITDispatchInfo class provide an arbitrary map of symbol names to addresses.

The JITDispatchInfo class will be replaced by bootstrap symbols with reserved
names in upcoming commits.

Author: lhames

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

@@ -158,6 +158,29 @@ class ExecutorProcessControl {
     return *MemMgr;
   }
 
+  /// Returns the bootstrap symbol map.
+  const StringMap<ExecutorAddress> &getBootstrapSymbolsMap() const {
+    return BootstrapSymbols;
+  }
+
+  /// For each (ExecutorAddress&, StringRef) pair, looks up the string in the
+  /// bootstrap symbols map and writes its address to the ExecutorAddress if
+  /// found. If any symbol is not found then the function returns an error.
+  Error getBootstrapSymbols(
+      ArrayRef<std::pair<ExecutorAddress &, StringRef>> Pairs) const {
+    for (auto &KV : Pairs) {
+      auto I = BootstrapSymbols.find(KV.second);
+      if (I == BootstrapSymbols.end())
+        return make_error<StringError>("Symbol \"" + KV.second +
+                                           "\" not found "
+                                           "in bootstrap symbols map",
+                                       inconvertibleErrorCode());
+
+      KV.first = I->second;
+    }
+    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.
@@ -252,6 +275,7 @@ class ExecutorProcessControl {
   JITDispatchInfo JDI;
   MemoryAccess *MemAccess = nullptr;
   jitlink::JITLinkMemoryManager *MemMgr = nullptr;
+  StringMap<ExecutorAddress> BootstrapSymbols;
 };
 
 /// A ExecutorProcessControl instance that asserts if any of its methods are

llvm/include/llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h

@@ -45,27 +45,6 @@ struct SimpleRemoteEPCExecutorInfo {
   std::string TargetTriple;
   uint64_t PageSize;
   StringMap<ExecutorAddress> BootstrapSymbols;
-
-  Expected<ExecutorAddress> getBootstrapSymbol(StringRef Name) const {
-    auto I = BootstrapSymbols.find(Name);
-    if (I == BootstrapSymbols.end())
-      return make_error<StringError>("Symbol \"" + Name +
-                                         "\" not found in "
-                                         "bootstrap symbols map",
-                                     inconvertibleErrorCode());
-    return I->second;
-  }
-
-  Error getBootstrapSymbols(
-      ArrayRef<std::pair<ExecutorAddress &, StringRef>> Pairs) const {
-    for (auto &KV : Pairs) {
-      if (auto A = getBootstrapSymbol(KV.second))
-        KV.first = *A;
-      else
-        return A.takeError();
-    }
-    return Error::success();
-  }
 };
 
 using SimpleRemoteEPCArgBytesVector = SmallVector<char, 128>;

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

@@ -61,14 +61,6 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   SimpleRemoteEPC &operator=(SimpleRemoteEPC &&) = delete;
   ~SimpleRemoteEPC();
 
-  /// Called at the end of the construction process to set up the instance.
-  ///
-  /// Override to set up custom memory manager and/or memory access objects.
-  /// This method must be called at the *end* of the subclass's
-  /// implementation.
-  virtual Error setup(std::unique_ptr<SimpleRemoteEPCTransport> T,
-                      const SimpleRemoteEPCExecutorInfo &EI);
-
   Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
 
   Expected<std::vector<tpctypes::LookupResult>>
@@ -91,20 +83,20 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   void handleDisconnect(Error Err) override;
 
 protected:
-  void setMemoryManager(std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr);
-  void setMemoryAccess(std::unique_ptr<MemoryAccess> MemAccess);
+  virtual Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
+  createMemoryManager();
+  virtual Expected<std::unique_ptr<MemoryAccess>> createMemoryAccess();
 
 private:
   SimpleRemoteEPC(std::shared_ptr<SymbolStringPool> SSP)
       : ExecutorProcessControl(std::move(SSP)) {}
 
-  Error setupDefaultMemoryManager(const SimpleRemoteEPCExecutorInfo &EI);
-  Error setupDefaultMemoryAccess(const SimpleRemoteEPCExecutorInfo &EI);
-
   Error handleSetup(uint64_t SeqNo, ExecutorAddress TagAddr,
                     SimpleRemoteEPCArgBytesVector ArgBytes);
   void prepareToReceiveSetupMessage(
       std::promise<MSVCPExpected<SimpleRemoteEPCExecutorInfo>> &ExecInfoP);
+  Error setup(std::unique_ptr<SimpleRemoteEPCTransport> T,
+              SimpleRemoteEPCExecutorInfo EI);
 
   Error handleResult(uint64_t SeqNo, ExecutorAddress TagAddr,
                      SimpleRemoteEPCArgBytesVector ArgBytes);

llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp

@@ -64,40 +64,6 @@ SimpleRemoteEPC::~SimpleRemoteEPC() {
   assert(Disconnected && "Destroyed without disconnection");
 }
 
-Error SimpleRemoteEPC::setup(std::unique_ptr<SimpleRemoteEPCTransport> T,
-                             const SimpleRemoteEPCExecutorInfo &EI) {
-  using namespace SimpleRemoteEPCDefaultBootstrapSymbolNames;
-  LLVM_DEBUG({
-    dbgs() << "SimpleRemoteEPC received setup message:\n"
-           << "  Triple: " << EI.TargetTriple << "\n"
-           << "  Page size: " << EI.PageSize << "\n"
-           << "  Bootstrap symbols:\n";
-    for (const auto &KV : EI.BootstrapSymbols)
-      dbgs() << "    " << KV.first() << ": "
-             << formatv("{0:x16}", KV.second.getValue()) << "\n";
-  });
-  this->T = std::move(T);
-  TargetTriple = Triple(EI.TargetTriple);
-  PageSize = EI.PageSize;
-
-  if (auto Err = EI.getBootstrapSymbols(
-          {{JDI.JITDispatchContextAddress, ExecutorSessionObjectName},
-           {JDI.JITDispatchFunctionAddress, DispatchFnName},
-           {LoadDylibAddr, "__llvm_orc_load_dylib"},
-           {LookupSymbolsAddr, "__llvm_orc_lookup_symbols"},
-           {RunAsMainAddr, "__llvm_orc_run_as_main"}}))
-    return Err;
-
-  if (!MemMgr)
-    if (auto Err = setupDefaultMemoryManager(EI))
-      return Err;
-  if (!MemAccess)
-    if (auto Err = setupDefaultMemoryAccess(EI))
-      return Err;
-
-  return Error::success();
-}
-
 Expected<tpctypes::DylibHandle>
 SimpleRemoteEPC::loadDylib(const char *DylibPath) {
   Expected<tpctypes::DylibHandle> H((tpctypes::DylibHandle()));
@@ -201,37 +167,22 @@ void SimpleRemoteEPC::handleDisconnect(Error Err) {
   }
 }
 
-void SimpleRemoteEPC::setMemoryManager(
-    std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr) {
-  OwnedMemMgr = std::move(MemMgr);
-  this->MemMgr = OwnedMemMgr.get();
-}
-
-void SimpleRemoteEPC::setMemoryAccess(std::unique_ptr<MemoryAccess> MemAccess) {
-  OwnedMemAccess = std::move(MemAccess);
-  this->MemAccess = OwnedMemAccess.get();
-}
-
-Error SimpleRemoteEPC::setupDefaultMemoryManager(
-    const SimpleRemoteEPCExecutorInfo &EI) {
-
+Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
+SimpleRemoteEPC::createMemoryManager() {
   EPCGenericJITLinkMemoryManager::FuncAddrs FAs;
-
-  if (auto Err = EI.getBootstrapSymbols(
+  if (auto Err = getBootstrapSymbols(
           {{FAs.Reserve, "__llvm_orc_memory_reserve"},
            {FAs.Finalize, "__llvm_orc_memory_finalize"},
            {FAs.Deallocate, "__llvm_orc_memory_deallocate"}}))
-    return Err;
+    return std::move(Err);
 
-  setMemoryManager(
-      std::make_unique<EPCGenericJITLinkMemoryManager>(*this, FAs));
-  return Error::success();
+  return std::make_unique<EPCGenericJITLinkMemoryManager>(*this, FAs);
 }
 
-Error SimpleRemoteEPC::setupDefaultMemoryAccess(
-    const SimpleRemoteEPCExecutorInfo &EI) {
+Expected<std::unique_ptr<ExecutorProcessControl::MemoryAccess>>
+SimpleRemoteEPC::createMemoryAccess() {
 
-  return Error::success();
+  return nullptr;
 }
 
 Error SimpleRemoteEPC::handleSetup(uint64_t SeqNo, ExecutorAddress TagAddr,
@@ -279,6 +230,46 @@ void SimpleRemoteEPC::prepareToReceiveSetupMessage(
       };
 }
 
+Error SimpleRemoteEPC::setup(std::unique_ptr<SimpleRemoteEPCTransport> T,
+                             SimpleRemoteEPCExecutorInfo EI) {
+  using namespace SimpleRemoteEPCDefaultBootstrapSymbolNames;
+  LLVM_DEBUG({
+    dbgs() << "SimpleRemoteEPC received setup message:\n"
+           << "  Triple: " << EI.TargetTriple << "\n"
+           << "  Page size: " << EI.PageSize << "\n"
+           << "  Bootstrap symbols:\n";
+    for (const auto &KV : EI.BootstrapSymbols)
+      dbgs() << "    " << KV.first() << ": "
+             << formatv("{0:x16}", KV.second.getValue()) << "\n";
+  });
+  this->T = std::move(T);
+  TargetTriple = Triple(EI.TargetTriple);
+  PageSize = EI.PageSize;
+  BootstrapSymbols = std::move(EI.BootstrapSymbols);
+
+  if (auto Err = getBootstrapSymbols(
+          {{JDI.JITDispatchContextAddress, ExecutorSessionObjectName},
+           {JDI.JITDispatchFunctionAddress, DispatchFnName},
+           {LoadDylibAddr, "__llvm_orc_load_dylib"},
+           {LookupSymbolsAddr, "__llvm_orc_lookup_symbols"},
+           {RunAsMainAddr, "__llvm_orc_run_as_main"}}))
+    return Err;
+
+  if (auto MemMgr = createMemoryManager()) {
+    OwnedMemMgr = std::move(*MemMgr);
+    this->MemMgr = OwnedMemMgr.get();
+  } else
+    return MemMgr.takeError();
+
+  if (auto MemAccess = createMemoryAccess()) {
+    OwnedMemAccess = std::move(*MemAccess);
+    this->MemAccess = OwnedMemAccess.get();
+  } else
+    return MemAccess.takeError();
+
+  return Error::success();
+}
+
 Error SimpleRemoteEPC::handleResult(uint64_t SeqNo, ExecutorAddress TagAddr,
                                     SimpleRemoteEPCArgBytesVector ArgBytes) {
   SendResultFunction SendResult;