[ORC] Update lazyReexports to support aliases with different symbol names.

lhames · lhames · commit d2fabd70065e · 2020-01-15T08:02:53.000-08:00
A bug in the existing implementation meant that lazyReexports would not work if
the aliased name differed from the alias's name, i.e. all lazy reexports had to
be of the form (lib1, name) -&gt; (lib2, name). This patch fixes the issue by
capturing the alias's name in the NotifyResolved callback. To simplify this
capture, and the LazyCallThroughManager code in general, the NotifyResolved
callback is updated to use llvm::unique_function rather than a custom class.

No test case yet: This can only be tested at runtime, and the only in-tree
client (lli) always uses aliases with matching names. I will add a new LLJIT
example shortly that will directly test the lazyReexports API and the
non-trivial alias use case.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
@@ -16,6 +16,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_LAZYREEXPORTS_H
 #define LLVM_EXECUTIONENGINE_ORC_LAZYREEXPORTS_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/Speculation.h"
@@ -36,50 +37,14 @@ namespace orc {
 /// function.
 class LazyCallThroughManager {
 public:
-  /// Clients will want to take some action on first resolution, e.g. updating
-  /// a stub pointer. Instances of this class can be used to implement this.
-  class NotifyResolvedFunction {
-  public:
-    virtual ~NotifyResolvedFunction() {}
-
-    /// Called the first time a lazy call through is executed and the target
-    /// symbol resolved.
-    virtual Error operator()(JITDylib &SourceJD,
-                             const SymbolStringPtr &SymbolName,
-                             JITTargetAddress ResolvedAddr) = 0;
-
-  private:
-    virtual void anchor();
-  };
-
-  template <typename NotifyResolvedImpl>
-  class NotifyResolvedFunctionImpl : public NotifyResolvedFunction {
-  public:
-    NotifyResolvedFunctionImpl(NotifyResolvedImpl NotifyResolved)
-        : NotifyResolved(std::move(NotifyResolved)) {}
-    Error operator()(JITDylib &SourceJD, const SymbolStringPtr &SymbolName,
-                     JITTargetAddress ResolvedAddr) {
-      return NotifyResolved(SourceJD, SymbolName, ResolvedAddr);
-    }
-
-  private:
-    NotifyResolvedImpl NotifyResolved;
-  };
-
-  /// Create a shared NotifyResolvedFunction from a given type that is
-  /// callable with the correct signature.
-  template <typename NotifyResolvedImpl>
-  static std::unique_ptr<NotifyResolvedFunction>
-  createNotifyResolvedFunction(NotifyResolvedImpl NotifyResolved) {
-    return std::make_unique<NotifyResolvedFunctionImpl<NotifyResolvedImpl>>(
-        std::move(NotifyResolved));
-  }
+  using NotifyResolvedFunction =
+      unique_function<Error(JITTargetAddress ResolvedAddr)>;
 
   // Return a free call-through trampoline and bind it to look up and call
   // through to the given symbol.
-  Expected<JITTargetAddress> getCallThroughTrampoline(
-      JITDylib &SourceJD, SymbolStringPtr SymbolName,
-      std::shared_ptr<NotifyResolvedFunction> NotifyResolved);
+  Expected<JITTargetAddress>
+  getCallThroughTrampoline(JITDylib &SourceJD, SymbolStringPtr SymbolName,
+                           NotifyResolvedFunction NotifyResolved);
 
 protected:
   LazyCallThroughManager(ExecutionSession &ES,
@@ -96,8 +61,7 @@ class LazyCallThroughManager {
   using ReexportsMap =
       std::map<JITTargetAddress, std::pair<JITDylib *, SymbolStringPtr>>;
 
-  using NotifiersMap =
-      std::map<JITTargetAddress, std::shared_ptr<NotifyResolvedFunction>>;
+  using NotifiersMap = std::map<JITTargetAddress, NotifyResolvedFunction>;
 
   std::mutex LCTMMutex;
   ExecutionSession &ES;
@@ -173,8 +137,6 @@ class LazyReexportsMaterializationUnit : public MaterializationUnit {
   IndirectStubsManager &ISManager;
   JITDylib &SourceJD;
   SymbolAliasMap CallableAliases;
-  std::shared_ptr<LazyCallThroughManager::NotifyResolvedFunction>
-      NotifyResolved;
   ImplSymbolMap *AliaseeTable;
 };
 
diff --git a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
@@ -16,16 +16,14 @@
 namespace llvm {
 namespace orc {
 
-void LazyCallThroughManager::NotifyResolvedFunction::anchor() {}
-
 LazyCallThroughManager::LazyCallThroughManager(
     ExecutionSession &ES, JITTargetAddress ErrorHandlerAddr,
     std::unique_ptr<TrampolinePool> TP)
     : ES(ES), ErrorHandlerAddr(ErrorHandlerAddr), TP(std::move(TP)) {}
 
 Expected<JITTargetAddress> LazyCallThroughManager::getCallThroughTrampoline(
     JITDylib &SourceJD, SymbolStringPtr SymbolName,
-    std::shared_ptr<NotifyResolvedFunction> NotifyResolved) {
+    NotifyResolvedFunction NotifyResolved) {
   std::lock_guard<std::mutex> Lock(LCTMMutex);
   auto Trampoline = TP->getTrampoline();
 
@@ -62,18 +60,18 @@ LazyCallThroughManager::callThroughToSymbol(JITTargetAddress TrampolineAddr) {
 
   auto ResolvedAddr = LookupResult->getAddress();
 
-  std::shared_ptr<NotifyResolvedFunction> NotifyResolved = nullptr;
+  NotifyResolvedFunction NotifyResolved;
   {
     std::lock_guard<std::mutex> Lock(LCTMMutex);
     auto I = Notifiers.find(TrampolineAddr);
     if (I != Notifiers.end()) {
-      NotifyResolved = I->second;
+      NotifyResolved = std::move(I->second);
       Notifiers.erase(I);
     }
   }
 
   if (NotifyResolved) {
-    if (auto Err = (*NotifyResolved)(*SourceJD, SymbolName, ResolvedAddr)) {
+    if (auto Err = NotifyResolved(ResolvedAddr)) {
       ES.reportError(std::move(Err));
       return ErrorHandlerAddr;
     }
@@ -128,11 +126,6 @@ LazyReexportsMaterializationUnit::LazyReexportsMaterializationUnit(
     : MaterializationUnit(extractFlags(CallableAliases), std::move(K)),
       LCTManager(LCTManager), ISManager(ISManager), SourceJD(SourceJD),
       CallableAliases(std::move(CallableAliases)),
-      NotifyResolved(LazyCallThroughManager::createNotifyResolvedFunction(
-          [&ISManager](JITDylib &JD, const SymbolStringPtr &SymbolName,
-                       JITTargetAddress ResolvedAddr) {
-            return ISManager.updatePointer(*SymbolName, ResolvedAddr);
-          })),
       AliaseeTable(SrcJDLoc) {}
 
 StringRef LazyReexportsMaterializationUnit::getName() const {
@@ -159,7 +152,11 @@ void LazyReexportsMaterializationUnit::materialize(
   for (auto &Alias : RequestedAliases) {
 
     auto CallThroughTrampoline = LCTManager.getCallThroughTrampoline(
-        SourceJD, Alias.second.Aliasee, NotifyResolved);
+        SourceJD, Alias.second.Aliasee,
+        [&ISManager = this->ISManager,
+         StubSym = Alias.first](JITTargetAddress ResolvedAddr) -> Error {
+          return ISManager.updatePointer(*StubSym, ResolvedAddr);
+        });
 
     if (!CallThroughTrampoline) {
       SourceJD.getExecutionSession().reportError(
diff --git a/llvm/unittests/ExecutionEngine/Orc/LazyCallThroughAndReexportsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/LazyCallThroughAndReexportsTest.cpp
@@ -51,12 +51,10 @@ TEST_F(LazyReexportsTest, BasicLocalCallThroughManagerOperation) {
       })));
 
   unsigned NotifyResolvedCount = 0;
-  auto NotifyResolved = LazyCallThroughManager::createNotifyResolvedFunction(
-      [&](JITDylib &JD, const SymbolStringPtr &SymbolName,
-          JITTargetAddress ResolvedAddr) {
-        ++NotifyResolvedCount;
-        return Error::success();
-      });
+  auto NotifyResolved = [&](JITTargetAddress ResolvedAddr) {
+    ++NotifyResolvedCount;
+    return Error::success();
+  };
 
   auto CallThroughTrampoline = cantFail((*LCTM)->getCallThroughTrampoline(
       JD, DummyTarget, std::move(NotifyResolved)));
