From c977251ef6fee3f7045039832331e8aab1e0a1b2 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Mon, 24 Oct 2022 13:46:10 -0700
Subject: [PATCH] [ORC] Allow EPCEHFrameRegistrar clients to specify
 registration function dylib.

Previously, EPCEHFrameRegistrar always used the
ExecutorProcessControl::loadDylib(nullptr) method to obtain a handle for the
process, but this doesn't work if the registration functions aren't visible in
a standard search of the process (e.g. if the JIT is in a plugin that is loaded
with RTLD_LOCAL).

This patch retains the old behavior by default, but allows clients to supply
their own handle for the library containing the registration functions if they
need to (e.g. to work around limitations like RDLD_LOCAL above, which aren't
expressible within the existing loadDylib / DynamicLibrary APIs).
---
 .../llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h       |  7 ++++++-
 llvm/lib/ExecutionEngine/Orc/EPCEHFrameRegistrar.cpp     | 16 +++++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h b/llvm/include/llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h
index 3de55a1..bf45683 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h
@@ -28,8 +28,13 @@ public:
   /// Create from a ExecutorProcessControl instance alone. This will use
   /// the EPC's lookupSymbols method to find the registration/deregistration
   /// funciton addresses by name.
+  ///
+  /// If RegistrationFunctionsDylib is non-None then it will be searched to
+  /// find the registration functions. If it is None then the process dylib
+  /// will be loaded to find the registration functions.
   static Expected<std::unique_ptr<EPCEHFrameRegistrar>>
-  Create(ExecutionSession &ES);
+  Create(ExecutionSession &ES,
+         Optional<ExecutorAddr> RegistrationFunctionsDylib = None);
 
   /// Create a EPCEHFrameRegistrar with the given ExecutorProcessControl
   /// object and registration/deregistration function addresses.
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCEHFrameRegistrar.cpp b/llvm/lib/ExecutionEngine/Orc/EPCEHFrameRegistrar.cpp
index 256ce94..37dea5c 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCEHFrameRegistrar.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCEHFrameRegistrar.cpp
@@ -17,16 +17,21 @@ namespace llvm {
 namespace orc {
 
 Expected<std::unique_ptr<EPCEHFrameRegistrar>>
-EPCEHFrameRegistrar::Create(ExecutionSession &ES) {
+EPCEHFrameRegistrar::Create(ExecutionSession &ES,
+                            Optional<ExecutorAddr> RegistrationFunctionsDylib) {
   // FIXME: Proper mangling here -- we really need to decouple linker mangling
   // from DataLayout.
 
   // Find the addresses of the registration/deregistration functions in the
   // executor process.
   auto &EPC = ES.getExecutorProcessControl();
-  auto ProcessHandle = EPC.loadDylib(nullptr);
-  if (!ProcessHandle)
-    return ProcessHandle.takeError();
+
+  if (!RegistrationFunctionsDylib) {
+    if (auto D = EPC.loadDylib(nullptr))
+      RegistrationFunctionsDylib = *D;
+    else
+      return D.takeError();
+  }
 
   std::string RegisterWrapperName, DeregisterWrapperName;
   if (EPC.getTargetTriple().isOSBinFormatMachO()) {
@@ -40,7 +45,8 @@ EPCEHFrameRegistrar::Create(ExecutionSession &ES) {
   RegistrationSymbols.add(EPC.intern(RegisterWrapperName));
   RegistrationSymbols.add(EPC.intern(DeregisterWrapperName));
 
-  auto Result = EPC.lookupSymbols({{*ProcessHandle, RegistrationSymbols}});
+  auto Result =
+      EPC.lookupSymbols({{*RegistrationFunctionsDylib, RegistrationSymbols}});
   if (!Result)
     return Result.takeError();
 
-- 
2.7.4