[ORC] Add Executor Resolver Utility (#143654)

This patch adds the new **executor-side resolver API** as suggested by
@lhames. It introduces a `DylibSymbolResolver` that helps resolve
symbols for each loaded dylib.

Previously, we returned a `DylibHandle` to the controller. Now, we wrap
the native handle inside `DylibSymbolResolver` and return a
`ResolverHandle` instead. This makes the code cleaner and separates the
symbol resolution logic from raw handle management.

Author: SahilPatidar

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h

@@ -34,7 +34,7 @@ class EPCGenericDylibManager {
   struct SymbolAddrs {
     ExecutorAddr Instance;
     ExecutorAddr Open;
-    ExecutorAddr Lookup;
+    ExecutorAddr Resolve;
   };
 
   /// Create an EPCGenericMemoryAccess instance from a given set of
@@ -51,25 +51,25 @@ class EPCGenericDylibManager {
   LLVM_ABI Expected<tpctypes::DylibHandle> open(StringRef Path, uint64_t Mode);
 
   /// Looks up symbols within the given dylib.
-  Expected<std::vector<ExecutorSymbolDef>>
-  lookup(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup) {
-    std::promise<MSVCPExpected<std::vector<ExecutorSymbolDef>>> RP;
+  Expected<tpctypes::LookupResult> lookup(tpctypes::DylibHandle H,
+                                          const SymbolLookupSet &Lookup) {
+    std::promise<MSVCPExpected<tpctypes::LookupResult>> RP;
     auto RF = RP.get_future();
     lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
     return RF.get();
   }
 
   /// Looks up symbols within the given dylib.
-  Expected<std::vector<ExecutorSymbolDef>>
-  lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup) {
-    std::promise<MSVCPExpected<std::vector<ExecutorSymbolDef>>> RP;
+  Expected<tpctypes::LookupResult> lookup(tpctypes::DylibHandle H,
+                                          const RemoteSymbolLookupSet &Lookup) {
+    std::promise<MSVCPExpected<tpctypes::LookupResult>> RP;
     auto RF = RP.get_future();
     lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
     return RF.get();
   }
 
   using SymbolLookupCompleteFn =
-      unique_function<void(Expected<std::vector<ExecutorSymbolDef>>)>;
+      unique_function<void(Expected<tpctypes::LookupResult>)>;
 
   /// Looks up symbols within the given dylib.
   LLVM_ABI void lookupAsync(tpctypes::DylibHandle H,

llvm/include/llvm/ExecutionEngine/Orc/ExecutorResolutionGenerator.h

@@ -0,0 +1,74 @@
+//===----- ExecutorResolver.h - Resolve symbols in executor -----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Declares ExecutorResolutionGenerator for symbol resolution,
+// dynamic library loading, and lookup in an executor process via
+// ExecutorResolver.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_EXECUTORRESOLUTIONGENERATOR_H
+#define LLVM_EXECUTIONENGINE_ORC_EXECUTORRESOLUTIONGENERATOR_H
+
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+namespace llvm::orc {
+
+class ExecutorResolutionGenerator : public DefinitionGenerator {
+public:
+  using SymbolPredicate = unique_function<bool(const SymbolStringPtr &)>;
+  using AbsoluteSymbolsFn =
+      unique_function<std::unique_ptr<MaterializationUnit>(SymbolMap)>;
+
+  ExecutorResolutionGenerator(
+      ExecutionSession &ES, tpctypes::ResolverHandle H,
+      SymbolPredicate Allow = SymbolPredicate(),
+      AbsoluteSymbolsFn AbsoluteSymbols = absoluteSymbols)
+      : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)),
+        AbsoluteSymbols(std::move(AbsoluteSymbols)) {}
+
+  ExecutorResolutionGenerator(
+      ExecutionSession &ES, SymbolPredicate Allow = SymbolPredicate(),
+      AbsoluteSymbolsFn AbsoluteSymbols = absoluteSymbols)
+      : EPC(ES.getExecutorProcessControl()), Allow(std::move(Allow)),
+        AbsoluteSymbols(std::move(AbsoluteSymbols)) {}
+
+  /// Permanently loads the library at the given path and, on success, returns
+  /// an ExecutorResolutionGenerator 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<ExecutorResolutionGenerator>>
+  Load(ExecutionSession &ES, const char *LibraryPath,
+       SymbolPredicate Allow = SymbolPredicate(),
+       AbsoluteSymbolsFn AbsoluteSymbols = absoluteSymbols);
+
+  /// Creates a ExecutorResolutionGenerator that searches for symbols in
+  /// the target process.
+  static Expected<std::unique_ptr<ExecutorResolutionGenerator>>
+  GetForTargetProcess(ExecutionSession &ES,
+                      SymbolPredicate Allow = SymbolPredicate(),
+                      AbsoluteSymbolsFn AbsoluteSymbols = absoluteSymbols) {
+    return Load(ES, nullptr, std::move(Allow), std::move(AbsoluteSymbols));
+  }
+
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
+                      JITDylibLookupFlags JDLookupFlags,
+                      const SymbolLookupSet &LookupSet) override;
+
+private:
+  ExecutorProcessControl &EPC;
+  tpctypes::ResolverHandle H;
+  SymbolPredicate Allow;
+  AbsoluteSymbolsFn AbsoluteSymbols;
+};
+
+} // namespace llvm::orc
+
+#endif // LLVM_EXECUTIONENGINE_ORC_EXECUTORRESOLUTIONGENERATOR_H

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

@@ -25,7 +25,7 @@ namespace rt {
 
 LLVM_ABI extern const char *SimpleExecutorDylibManagerInstanceName;
 LLVM_ABI extern const char *SimpleExecutorDylibManagerOpenWrapperName;
-LLVM_ABI extern const char *SimpleExecutorDylibManagerLookupWrapperName;
+LLVM_ABI extern const char *SimpleExecutorDylibManagerResolveWrapperName;
 
 LLVM_ABI extern const char *SimpleExecutorMemoryManagerInstanceName;
 LLVM_ABI extern const char *SimpleExecutorMemoryManagerReserveWrapperName;
@@ -66,10 +66,9 @@ using SPSSimpleExecutorDylibManagerOpenSignature =
     shared::SPSExpected<shared::SPSExecutorAddr>(shared::SPSExecutorAddr,
                                                  shared::SPSString, uint64_t);
 
-using SPSSimpleExecutorDylibManagerLookupSignature =
-    shared::SPSExpected<shared::SPSSequence<shared::SPSExecutorSymbolDef>>(
-        shared::SPSExecutorAddr, shared::SPSExecutorAddr,
-        shared::SPSRemoteSymbolLookupSet);
+using SPSSimpleExecutorDylibManagerResolveSignature = shared::SPSExpected<
+    shared::SPSSequence<shared::SPSOptional<shared::SPSExecutorSymbolDef>>>(
+    shared::SPSExecutorAddr, shared::SPSRemoteSymbolLookupSet);
 
 using SPSSimpleExecutorMemoryManagerReserveSignature =
     shared::SPSExpected<shared::SPSExecutorAddr>(shared::SPSExecutorAddr,

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

@@ -114,7 +114,11 @@ struct PointerWrite {
 /// A handle used to represent a loaded dylib in the target process.
 using DylibHandle = ExecutorAddr;
 
-using LookupResult = std::vector<ExecutorSymbolDef>;
+/// A handle used to reference the resolver associated with a loaded
+///  dylib in the target process.
+using ResolverHandle = ExecutorAddr;
+
+using LookupResult = std::vector<std::optional<ExecutorSymbolDef>>;
 
 } // end namespace tpctypes
 

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/ExecutorResolver.h

@@ -0,0 +1,48 @@
+//===----- ExecutorResolver.h - Symbol resolver -----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Executor Symbol resolver.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORRESOLVER_H
+#define LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORRESOLVER_H
+
+#include "llvm/ADT/FunctionExtras.h"
+
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
+#include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
+#include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
+
+namespace llvm::orc {
+
+class ExecutorResolver {
+public:
+  using ResolveResult = Expected<std::vector<std::optional<ExecutorSymbolDef>>>;
+  using YieldResolveResultFn = unique_function<void(ResolveResult)>;
+
+  virtual ~ExecutorResolver() = default;
+
+  virtual void resolveAsync(const RemoteSymbolLookupSet &L,
+                            YieldResolveResultFn &&OnResolve) = 0;
+};
+
+class DylibSymbolResolver : public ExecutorResolver {
+public:
+  DylibSymbolResolver(tpctypes::DylibHandle H) : Handle(H) {}
+
+  void
+  resolveAsync(const RemoteSymbolLookupSet &L,
+               ExecutorResolver::YieldResolveResultFn &&OnResolve) override;
+
+private:
+  tpctypes::DylibHandle Handle;
+};
+
+} // end namespace llvm::orc
+#endif // LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORRESOLVER_H

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorDylibManager.h

@@ -23,6 +23,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorResolver.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/Error.h"
@@ -39,8 +40,6 @@ class LLVM_ABI SimpleExecutorDylibManager : public ExecutorBootstrapService {
   virtual ~SimpleExecutorDylibManager();
 
   Expected<tpctypes::DylibHandle> open(const std::string &Path, uint64_t Mode);
-  Expected<std::vector<ExecutorSymbolDef>>
-  lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &L);
 
   Error shutdown() override;
   void addBootstrapSymbols(StringMap<ExecutorAddr> &M) override;
@@ -52,10 +51,11 @@ class LLVM_ABI SimpleExecutorDylibManager : public ExecutorBootstrapService {
   openWrapper(const char *ArgData, size_t ArgSize);
 
   static llvm::orc::shared::CWrapperFunctionResult
-  lookupWrapper(const char *ArgData, size_t ArgSize);
+  resolveWrapper(const char *ArgData, size_t ArgSize);
 
   std::mutex M;
   DylibSet Dylibs;
+  std::vector<std::unique_ptr<ExecutorResolver>> Resolvers;
 };
 
 } // end namespace rt_bootstrap

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -24,6 +24,7 @@ add_llvm_component_library(LLVMOrcJIT
   EPCGenericRTDyldMemoryManager.cpp
   EPCIndirectionUtils.cpp
   ExecutionUtils.cpp
+  ExecutorResolutionGenerator.cpp
   ObjectFileInterface.cpp
   GetDylibInterface.cpp
   IndirectionUtils.cpp

llvm/lib/ExecutionEngine/Orc/EPCDebugObjectRegistrar.cpp

@@ -42,7 +42,12 @@ Expected<std::unique_ptr<EPCDebugObjectRegistrar>> createJITLoaderGDBRegistrar(
   assert((*Result)[0].size() == 1 &&
          "Unexpected number of addresses in result");
 
-  ExecutorAddr RegisterAddr = (*Result)[0][0].getAddress();
+  if (!(*Result)[0][0].has_value())
+    return make_error<StringError>(
+        "Expected a valid address in the lookup result",
+        inconvertibleErrorCode());
+
+  ExecutorAddr RegisterAddr = (*Result)[0][0]->getAddress();
   return std::make_unique<EPCDebugObjectRegistrar>(ES, RegisterAddr);
 }
 

llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp

@@ -79,12 +79,16 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     assert(Result->front().size() == LookupSymbols.size() &&
            "Result has incorrect number of elements");
 
+    auto SymsIt = Result->front().begin();
+    SymbolNameSet MissingSymbols;
     SymbolMap NewSymbols;
-    auto ResultI = Result->front().begin();
-    for (auto &KV : LookupSymbols) {
-      if (ResultI->getAddress())
-        NewSymbols[KV.first] = *ResultI;
-      ++ResultI;
+    for (auto &[Name, Flags] : LookupSymbols) {
+      const auto &Sym = *SymsIt++;
+      if (Sym && Sym->getAddress())
+        NewSymbols[Name] = *Sym;
+      else if (LLVM_UNLIKELY(!Sym &&
+                             Flags == SymbolLookupFlags::RequiredSymbol))
+        MissingSymbols.insert(Name);
     }
 
     LLVM_DEBUG({
@@ -96,6 +100,10 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     if (NewSymbols.empty())
       return LS.continueLookup(Error::success());
 
+    if (LLVM_UNLIKELY(!MissingSymbols.empty()))
+      return LS.continueLookup(make_error<SymbolsNotFound>(
+          this->EPC.getSymbolStringPool(), std::move(MissingSymbols)));
+
     // Define resolved symbols.
     Error Err = addAbsolutes(JD, std::move(NewSymbols));
 

llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp

@@ -66,7 +66,7 @@ EPCGenericDylibManager::CreateWithDefaultBootstrapSymbols(
   if (auto Err = EPC.getBootstrapSymbols(
           {{SAs.Instance, rt::SimpleExecutorDylibManagerInstanceName},
            {SAs.Open, rt::SimpleExecutorDylibManagerOpenWrapperName},
-           {SAs.Lookup, rt::SimpleExecutorDylibManagerLookupWrapperName}}))
+           {SAs.Resolve, rt::SimpleExecutorDylibManagerResolveWrapperName}}))
     return std::move(Err);
   return EPCGenericDylibManager(EPC, std::move(SAs));
 }
@@ -84,37 +84,39 @@ Expected<tpctypes::DylibHandle> EPCGenericDylibManager::open(StringRef Path,
 void EPCGenericDylibManager::lookupAsync(tpctypes::DylibHandle H,
                                          const SymbolLookupSet &Lookup,
                                          SymbolLookupCompleteFn Complete) {
-  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
-      SAs.Lookup,
+  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerResolveSignature>(
+      SAs.Resolve,
       [Complete = std::move(Complete)](
           Error SerializationErr,
-          Expected<std::vector<ExecutorSymbolDef>> Result) mutable {
+          Expected<std::vector<std::optional<ExecutorSymbolDef>>>
+              Result) mutable {
         if (SerializationErr) {
           cantFail(Result.takeError());
           Complete(std::move(SerializationErr));
           return;
         }
         Complete(std::move(Result));
       },
-      SAs.Instance, H, Lookup);
+      H, Lookup);
 }
 
 void EPCGenericDylibManager::lookupAsync(tpctypes::DylibHandle H,
                                          const RemoteSymbolLookupSet &Lookup,
                                          SymbolLookupCompleteFn Complete) {
-  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
-      SAs.Lookup,
+  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerResolveSignature>(
+      SAs.Resolve,
       [Complete = std::move(Complete)](
           Error SerializationErr,
-          Expected<std::vector<ExecutorSymbolDef>> Result) mutable {
+          Expected<std::vector<std::optional<ExecutorSymbolDef>>>
+              Result) mutable {
         if (SerializationErr) {
           cantFail(Result.takeError());
           Complete(std::move(SerializationErr));
           return;
         }
         Complete(std::move(Result));
       },
-      SAs.Instance, H, Lookup);
+      H, Lookup);
 }
 
 } // end namespace orc