[ORC] lazyReexports: Swap IndirectStubsManager for RedirectableSymbolsManager.

RedirectableSymbolsManager is a native SymbolStringPtr API (requires fewer
string operations) and has a narrower interface that permits a wider range of
implementations.

IndirectStubsManager is updated to make it a RedirectableSymbolsManager so that
existing uses continue to work.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h

@@ -18,6 +18,7 @@
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/Process.h"
@@ -278,7 +279,7 @@ class LocalJITCompileCallbackManager : public JITCompileCallbackManager {
 };
 
 /// Base class for managing collections of named indirect stubs.
-class IndirectStubsManager {
+class IndirectStubsManager : public RedirectableSymbolManager {
 public:
   /// Map type for initializing the manager. See init.
   using StubInitsMap = StringMap<std::pair<ExecutorAddr, JITSymbolFlags>>;
@@ -305,8 +306,15 @@ class IndirectStubsManager {
   /// Change the value of the implementation pointer for the stub.
   virtual Error updatePointer(StringRef Name, ExecutorAddr NewAddr) = 0;
 
+  /// --- RedirectableSymbolManager implementation ---
+  Error redirect(JITDylib &JD, const SymbolMap &NewDests) override;
+
+  void
+  emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> MR,
+                          SymbolMap InitialDests) override;
+
 private:
-  virtual void anchor();
+  void anchor() override;
 };
 
 template <typename ORCABI> class LocalIndirectStubsInfo {

llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h

@@ -19,6 +19,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
 #include "llvm/ExecutionEngine/Orc/Speculation.h"
 
 namespace llvm {
@@ -140,7 +141,7 @@ createLocalLazyCallThroughManager(const Triple &T, ExecutionSession &ES,
 class LazyReexportsMaterializationUnit : public MaterializationUnit {
 public:
   LazyReexportsMaterializationUnit(LazyCallThroughManager &LCTManager,
-                                   IndirectStubsManager &ISManager,
+                                   RedirectableSymbolManager &RSManager,
                                    JITDylib &SourceJD,
                                    SymbolAliasMap CallableAliases,
                                    ImplSymbolMap *SrcJDLoc);
@@ -154,7 +155,7 @@ class LazyReexportsMaterializationUnit : public MaterializationUnit {
   extractFlags(const SymbolAliasMap &Aliases);
 
   LazyCallThroughManager &LCTManager;
-  IndirectStubsManager &ISManager;
+  RedirectableSymbolManager &RSManager;
   JITDylib &SourceJD;
   SymbolAliasMap CallableAliases;
   ImplSymbolMap *AliaseeTable;
@@ -165,11 +166,11 @@ class LazyReexportsMaterializationUnit : public MaterializationUnit {
 /// first call. All subsequent calls will go directly to the aliasee.
 inline std::unique_ptr<LazyReexportsMaterializationUnit>
 lazyReexports(LazyCallThroughManager &LCTManager,
-              IndirectStubsManager &ISManager, JITDylib &SourceJD,
+              RedirectableSymbolManager &RSManager, JITDylib &SourceJD,
               SymbolAliasMap CallableAliases,
               ImplSymbolMap *SrcJDLoc = nullptr) {
   return std::make_unique<LazyReexportsMaterializationUnit>(
-      LCTManager, ISManager, SourceJD, std::move(CallableAliases), SrcJDLoc);
+      LCTManager, RSManager, SourceJD, std::move(CallableAliases), SrcJDLoc);
 }
 
 } // End namespace orc

llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp

@@ -115,6 +115,38 @@ JITCompileCallbackManager::executeCompileCallback(ExecutorAddr TrampolineAddr) {
   }
 }
 
+Error IndirectStubsManager::redirect(JITDylib &JD, const SymbolMap &NewDests) {
+  for (auto &[Name, Dest] : NewDests)
+    if (auto Err = updatePointer(*Name, Dest.getAddress()))
+      return Err;
+  return Error::success();
+}
+
+void IndirectStubsManager::emitRedirectableSymbols(
+    std::unique_ptr<MaterializationResponsibility> MR, SymbolMap InitialDests) {
+  StubInitsMap StubInits;
+  for (auto &[Name, Dest] : InitialDests)
+    StubInits[*Name] = {Dest.getAddress(), Dest.getFlags()};
+  if (auto Err = createStubs(StubInits)) {
+    MR->getExecutionSession().reportError(std::move(Err));
+    return MR->failMaterialization();
+  }
+  SymbolMap Stubs;
+  for (auto &[Name, Dest] : InitialDests) {
+    auto StubSym = findStub(*Name, false);
+    assert(StubSym.getAddress() && "Stub symbol should be present");
+    Stubs[Name] = StubSym;
+  }
+  if (auto Err = MR->notifyResolved(Stubs)) {
+    MR->getExecutionSession().reportError(std::move(Err));
+    return MR->failMaterialization();
+  }
+  if (auto Err = MR->notifyEmitted({})) {
+    MR->getExecutionSession().reportError(std::move(Err));
+    return MR->failMaterialization();
+  }
+}
+
 Expected<std::unique_ptr<JITCompileCallbackManager>>
 createLocalCompileCallbackManager(const Triple &T, ExecutionSession &ES,
                                   ExecutorAddr ErrorHandlerAddress) {

llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

@@ -151,10 +151,10 @@ createLocalLazyCallThroughManager(const Triple &T, ExecutionSession &ES,
 }
 
 LazyReexportsMaterializationUnit::LazyReexportsMaterializationUnit(
-    LazyCallThroughManager &LCTManager, IndirectStubsManager &ISManager,
+    LazyCallThroughManager &LCTManager, RedirectableSymbolManager &RSManager,
     JITDylib &SourceJD, SymbolAliasMap CallableAliases, ImplSymbolMap *SrcJDLoc)
     : MaterializationUnit(extractFlags(CallableAliases)),
-      LCTManager(LCTManager), ISManager(ISManager), SourceJD(SourceJD),
+      LCTManager(LCTManager), RSManager(RSManager), SourceJD(SourceJD),
       CallableAliases(std::move(CallableAliases)), AliaseeTable(SrcJDLoc) {}
 
 StringRef LazyReexportsMaterializationUnit::getName() const {
@@ -174,51 +174,41 @@ void LazyReexportsMaterializationUnit::materialize(
   }
 
   if (!CallableAliases.empty())
-    if (auto Err = R->replace(lazyReexports(LCTManager, ISManager, SourceJD,
+    if (auto Err = R->replace(lazyReexports(LCTManager, RSManager, SourceJD,
                                             std::move(CallableAliases),
                                             AliaseeTable))) {
       R->getExecutionSession().reportError(std::move(Err));
       R->failMaterialization();
       return;
     }
 
-  IndirectStubsManager::StubInitsMap StubInits;
+  SymbolMap Inits;
   for (auto &Alias : RequestedAliases) {
-
     auto CallThroughTrampoline = LCTManager.getCallThroughTrampoline(
         SourceJD, Alias.second.Aliasee,
-        [&ISManager = this->ISManager,
+        [&TargetJD = R->getTargetJITDylib(), &RSManager = this->RSManager,
          StubSym = Alias.first](ExecutorAddr ResolvedAddr) -> Error {
-          return ISManager.updatePointer(*StubSym, ResolvedAddr);
+          return RSManager.redirect(TargetJD, StubSym,
+                                    ExecutorSymbolDef(ResolvedAddr, {}));
         });
 
     if (!CallThroughTrampoline) {
-      SourceJD.getExecutionSession().reportError(
-          CallThroughTrampoline.takeError());
+      R->getExecutionSession().reportError(CallThroughTrampoline.takeError());
       R->failMaterialization();
       return;
     }
 
-    StubInits[*Alias.first] =
-        std::make_pair(*CallThroughTrampoline, Alias.second.AliasFlags);
+    Inits[Alias.first] = {*CallThroughTrampoline, Alias.second.AliasFlags};
   }
 
   if (AliaseeTable != nullptr && !RequestedAliases.empty())
     AliaseeTable->trackImpls(RequestedAliases, &SourceJD);
 
-  if (auto Err = ISManager.createStubs(StubInits)) {
-    SourceJD.getExecutionSession().reportError(std::move(Err));
-    R->failMaterialization();
-    return;
+  if (auto Err = R->replace(std::make_unique<RedirectableMaterializationUnit>(
+          RSManager, std::move(Inits)))) {
+    R->getExecutionSession().reportError(std::move(Err));
+    return R->failMaterialization();
   }
-
-  SymbolMap Stubs;
-  for (auto &Alias : RequestedAliases)
-    Stubs[Alias.first] = ISManager.findStub(*Alias.first, false);
-
-  // No registered dependencies, so these calls cannot fail.
-  cantFail(R->notifyResolved(Stubs));
-  cantFail(R->notifyEmitted({}));
 }
 
 void LazyReexportsMaterializationUnit::discard(const JITDylib &JD,