[ORC] Add a TargetProcessControl-based dynamic library search generator.

TPCDynamicLibrarySearchGenerator uses a TargetProcessControl instance to
load libraries and search for symbol addresses in a target process. It
can be used in place of a DynamicLibrarySearchGenerator to enable
target-process agnostic lookup.

llvm/examples/OrcV2Examples/LLJITWithTargetProcessControl/LLJITWithTargetProcessControl.cpp

@@ -25,6 +25,7 @@
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
+#include "llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h"
 #include "llvm/ExecutionEngine/Orc/TPCIndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
 #include "llvm/Support/InitLLVM.h"
@@ -51,17 +52,23 @@ ExitOnError ExitOnErr;
 
 const llvm::StringRef FooMod =
     R"(
+  declare i32 @return1()
+
   define i32 @foo_body() {
   entry:
-    ret i32 1
+    %0 = call i32 @return1()
+    ret i32 %0
   }
 )";
 
 const llvm::StringRef BarMod =
     R"(
+  declare i32 @return2()
+
   define i32 @bar_body() {
   entry:
-    ret i32 2
+    %0 = call i32 @return2()
+    ret i32 %0
   }
 )";
 
@@ -91,6 +98,9 @@ const llvm::StringRef MainMod =
   declare i32 @bar()
 )";
 
+extern "C" int32_t return1() { return 1; }
+extern "C" int32_t return2() { return 2; }
+
 static void *reenter(void *Ctx, void *TrampolineAddr) {
   std::promise<void *> LandingAddressP;
   auto LandingAddressF = LandingAddressP.get_future();
@@ -105,6 +115,11 @@ static void *reenter(void *Ctx, void *TrampolineAddr) {
   return LandingAddressF.get();
 }
 
+static void reportErrorAndExit() {
+  errs() << "Unable to lazily compile function. Exiting.\n";
+  exit(1);
+}
+
 cl::list<std::string> InputArgv(cl::Positional,
                                 cl::desc("<program arguments>..."));
 
@@ -135,8 +150,11 @@ int main(int argc, char *argv[]) {
   auto TPCIU = ExitOnErr(TPCIndirectionUtils::Create(*TPC));
   ExitOnErr(TPCIU->writeResolverBlock(pointerToJITTargetAddress(&reenter),
                                       pointerToJITTargetAddress(TPCIU.get())));
-  TPCIU->createLazyCallThroughManager(J->getExecutionSession(), 0);
+  TPCIU->createLazyCallThroughManager(
+      J->getExecutionSession(), pointerToJITTargetAddress(&reportErrorAndExit));
   auto ISM = TPCIU->createIndirectStubsManager();
+  J->getMainJITDylib().addGenerator(
+      ExitOnErr(TPCDynamicLibrarySearchGenerator::GetForTargetProcess(*TPC)));
 
   // (4) Add modules.
   ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(FooMod, "foo-mod"))));

llvm/include/llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h

@@ -0,0 +1,60 @@
+//===------------ TPCDynamicLibrarySearchGenerator.h ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Support loading and searching of dynamic libraries in a target process via
+// the TargetProcessControl class.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_TPCDYNAMICLIBRARYSEARCHGENERATOR_H
+#define LLVM_EXECUTIONENGINE_ORC_TPCDYNAMICLIBRARYSEARCHGENERATOR_H
+
+#include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
+
+namespace llvm {
+namespace orc {
+
+class TPCDynamicLibrarySearchGenerator : public JITDylib::DefinitionGenerator {
+public:
+  /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
+  /// library with the given handle.
+  ///
+  /// 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.
+  TPCDynamicLibrarySearchGenerator(
+      TargetProcessControl &TPC,
+      TargetProcessControl::DynamicLibraryHandle DylibHandle)
+      : TPC(TPC), DylibHandle(DylibHandle) {}
+
+  /// 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<TPCDynamicLibrarySearchGenerator>>
+  Load(TargetProcessControl &TPC, const char *LibraryPath);
+
+  /// Creates a TPCDynamicLibrarySearchGenerator that searches for symbols in
+  /// the target process.
+  static Expected<std::unique_ptr<TPCDynamicLibrarySearchGenerator>>
+  GetForTargetProcess(TargetProcessControl &TPC) {
+    return Load(TPC, nullptr);
+  }
+
+  Error tryToGenerate(LookupKind K, JITDylib &JD,
+                      JITDylibLookupFlags JDLookupFlags,
+                      const SymbolLookupSet &Symbols) override;
+
+private:
+  TargetProcessControl &TPC;
+  TargetProcessControl::DynamicLibraryHandle DylibHandle;
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_TPCDYNAMICLIBRARYSEARCHGENERATOR_H

llvm/include/llvm/ExecutionEngine/Orc/TargetProcessControl.h

@@ -13,11 +13,16 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_TARGETPROCESSCONTROL_H
 #define LLVM_EXECUTIONENGINE_ORC_TARGETPROCESSCONTROL_H
 
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/MSVCErrorWorkarounds.h"
 
 #include <future>
+#include <vector>
 
 namespace llvm {
 namespace orc {
@@ -105,6 +110,22 @@ public:
     }
   };
 
+  using DynamicLibraryHandle = JITTargetAddress;
+
+  /// Request lookup within a single library.
+  /// If Library is None then the whole target process should be searched.
+  struct LookupRequestElement {
+    LookupRequestElement(DynamicLibraryHandle Handle,
+                         const SymbolLookupSet &Symbols)
+        : Handle(Handle), Symbols(Symbols) {}
+    DynamicLibraryHandle Handle;
+    const SymbolLookupSet &Symbols;
+  };
+
+  using LookupRequest = ArrayRef<LookupRequestElement>;
+
+  using LookupResult = std::vector<std::vector<JITTargetAddress>>;
+
   virtual ~TargetProcessControl();
 
   /// Return the Triple for the target process.
@@ -119,7 +140,16 @@ public:
   /// Return a MemoryAccess object for the target process.
   MemoryAccess &getMemoryAccess() const { return *MemAccess; }
 
-  /// Load the library at the given path.
+  /// Load the library at the given path. Returns a handle to the loaded
+  /// library. If LibraryPath is null this function will return the global
+  /// handle for the target process.
+  virtual Expected<DynamicLibraryHandle>
+  loadLibrary(const char *LibraryPath) = 0;
+
+  /// Search for symbols in the target process.
+  /// The result of the lookup is a 2-dimentional array of target addresses
+  /// that correspond to the lookup order.
+  virtual Expected<LookupResult> lookupSymbols(LookupRequest Request) = 0;
 
 protected:
   TargetProcessControl(Triple TT, unsigned PageSize);
@@ -138,6 +168,10 @@ public:
 
   static Expected<std::unique_ptr<SelfTargetProcessControl>> Create();
 
+  Expected<DynamicLibraryHandle> loadLibrary(const char *LibraryPath) override;
+
+  Expected<LookupResult> lookupSymbols(LookupRequest Request) override;
+
 private:
   void writeUInt8s(ArrayRef<UInt8Write> Ws,
                    WriteResultFn OnWriteComplete) override;
@@ -156,6 +190,9 @@ private:
 
   std::unique_ptr<jitlink::InProcessMemoryManager> IPMM =
       std::make_unique<jitlink::InProcessMemoryManager>();
+
+  char GlobalManglingPrefix = 0;
+  std::vector<std::unique_ptr<sys::DynamicLibrary>> DynamicLibraries;
 };
 
 } // end namespace orc

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -26,6 +26,7 @@ add_llvm_component_library(LLVMOrcJIT
   SpeculateAnalyses.cpp
   TargetProcessControl.cpp
   ThreadSafeModule.cpp
+  TPCDynamicLibrarySearchGenerator.cpp
   TPCIndirectionUtils.cpp
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc

llvm/lib/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.cpp

@@ -0,0 +1,51 @@
+//===---------------- TPCDynamicLibrarySearchGenerator.cpp ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h"
+
+namespace llvm {
+namespace orc {
+
+Expected<std::unique_ptr<TPCDynamicLibrarySearchGenerator>>
+TPCDynamicLibrarySearchGenerator::Load(TargetProcessControl &TPC,
+                                       const char *LibraryPath) {
+  auto Handle = TPC.loadLibrary(LibraryPath);
+  if (!Handle)
+    return Handle.takeError();
+
+  return std::make_unique<TPCDynamicLibrarySearchGenerator>(TPC, *Handle);
+}
+
+Error TPCDynamicLibrarySearchGenerator::tryToGenerate(
+    LookupKind K, JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+    const SymbolLookupSet &Symbols) {
+
+  if (Symbols.empty())
+    return Error::success();
+
+  SymbolMap NewSymbols;
+
+  TargetProcessControl::LookupRequestElement Request(DylibHandle, Symbols);
+  auto Result = TPC.lookupSymbols(Request);
+  if (!Result)
+    return Result.takeError();
+
+  assert(Result->size() == 1 && "Results for more than one library returned");
+  assert(Result->front().size() == Symbols.size() &&
+         "Result has incorrect number of elements");
+
+  auto ResultI = Result->front().begin();
+  for (auto &KV : Symbols)
+    NewSymbols[KV.first] =
+        JITEvaluatedSymbol(*ResultI++, JITSymbolFlags::Exported);
+
+  return JD.define(absoluteSymbols(std::move(NewSymbols)));
+}
+
+} // end namespace orc
+} // end namespace llvm

llvm/lib/ExecutionEngine/Orc/TargetProcessControl.cpp

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/Process.h"
 
@@ -26,6 +28,8 @@ SelfTargetProcessControl::SelfTargetProcessControl(Triple TT, unsigned PageSize)
     : TargetProcessControl(std::move(TT), PageSize) {
   this->MemMgr = IPMM.get();
   this->MemAccess = this;
+  if (this->TT.isOSBinFormatMachO())
+    GlobalManglingPrefix = '_';
 }
 
 Expected<std::unique_ptr<SelfTargetProcessControl>>
@@ -39,6 +43,48 @@ SelfTargetProcessControl::Create() {
   return std::make_unique<SelfTargetProcessControl>(std::move(TT), *PageSize);
 }
 
+Expected<TargetProcessControl::DynamicLibraryHandle>
+SelfTargetProcessControl::loadLibrary(const char *LibraryPath) {
+  std::string ErrMsg;
+  auto Dylib = std::make_unique<sys::DynamicLibrary>(
+      sys::DynamicLibrary::getPermanentLibrary(LibraryPath, &ErrMsg));
+  if (!Dylib->isValid())
+    return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
+  DynamicLibraries.push_back(std::move(Dylib));
+  return pointerToJITTargetAddress(DynamicLibraries.back().get());
+}
+
+Expected<TargetProcessControl::LookupResult>
+SelfTargetProcessControl::lookupSymbols(LookupRequest Request) {
+  LookupResult R;
+
+  for (auto &Elem : Request) {
+    auto *Dylib = jitTargetAddressToPointer<sys::DynamicLibrary *>(Elem.Handle);
+    assert(llvm::find_if(DynamicLibraries,
+                         [=](const std::unique_ptr<sys::DynamicLibrary> &DL) {
+                           return DL.get() == Dylib;
+                         }) != DynamicLibraries.end() &&
+           "Invalid handle");
+
+    R.push_back(std::vector<JITTargetAddress>());
+    for (auto &KV : Elem.Symbols) {
+      auto &Sym = KV.first;
+      std::string Tmp((*Sym).data() + !!GlobalManglingPrefix,
+                      (*Sym).size() - !!GlobalManglingPrefix);
+      if (void *Addr = Dylib->getAddressOfSymbol(Tmp.c_str()))
+        R.back().push_back(pointerToJITTargetAddress(Addr));
+      else if (KV.second == SymbolLookupFlags::RequiredSymbol) {
+        // FIXME: Collect all failing symbols before erroring out.
+        SymbolNameVector MissingSymbols;
+        MissingSymbols.push_back(Sym);
+        return make_error<SymbolsNotFound>(std::move(MissingSymbols));
+      }
+    }
+  }
+
+  return R;
+}
+
 void SelfTargetProcessControl::writeUInt8s(ArrayRef<UInt8Write> Ws,
                                            WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)