[ORC] Decompose LazyCallThroughManager::callThroughToSymbol()

[sgraenitz]

Summary

Decompose callThroughToSymbol() into findReexport(), resolveSymbol(), notifyResolved() and reportCallThroughError(). This allows derived classes to reuse the functionality while adding their own code in between. ThinLtoJIT's NotifyingCallThroughManager is an example for that. It wants to send a notification as soon as it figured out the symbol name.

[sgraenitz]

llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

45

So far this error ends up in reportCallThroughError(), which reports it to the ExecutionSession and eventually returns ErrorHandlerAddr to the JITed code. Not sure this can be recovered from, but if so, it should probably get its own ErrorInfo.

[lhames]

Style suggestion aside, LGTM. Thanks very much Stefan!

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

109

Could this lambda be factored out into a callThroughToSymbol method on LocalLazyCallThroughManager for consistency with NotifyingCallThroughManager?

[sgraenitz]

Factor out callThroughToSymbol(). Drop example use case from ThinLtoJIT. Format change set.

[sgraenitz]

Remove assert(Notifiers.find(TrampolineAddr) != Notifiers.end()); from findReexport() -- NotifyResolved may have been issued from a concurrent call-through

llvm/examples/ThinLtoJIT/OnDemandLayer.cpp

public:
 
   void onCallThroughToSymbol(NotifyCallThroughToSymbolFunction F) {
     NotifyCallThroughToSymbol = std::move(F);
   }
 };
 
 JITTargetAddress NotifyingCallThroughManager::callThroughToSymbol(
     JITTargetAddress TrampolineAddr) {
-  // TODO: Allow derived classes to find the reexports entry for TrampolineAddr
+  auto Entry = findReexport(TrampolineAddr);
+  if (!Entry)
+    return reportCallThroughError(Entry.takeError());
+
+  PSTATS("[Call-through to ", *Entry->SymbolName , "] Found reexport");
 
   // Kick-off discovery from here
   if (NotifyCallThroughToSymbol) {
-    NotifyCallThroughToSymbol(SymbolStringPtr());
+    NotifyCallThroughToSymbol(Entry->SymbolName);
   }
 
-  // TODO: Allow derived classes to resolve the call-through.
-  return JITTargetAddress();
+  PSTATS("[Call-through to ", *Entry->SymbolName , "] Kicked-off discovery");
+
+  auto ResolvedAddr = resolveSymbol(*Entry);
+  if (!ResolvedAddr)
+    return reportCallThroughError(ResolvedAddr.takeError());
+
+  PSTATS("[Call-through to ", *Entry->SymbolName , "] Resolved address");
+
+  if (Error Err = notifyResolved(TrampolineAddr, *ResolvedAddr))
+    return reportCallThroughError(std::move(Err));
+
+  PSTATS("[Call-through to ", *Entry->SymbolName , "] Returning");
+  return *ResolvedAddr;
 }
 
 CompileWholeModuleOnDemandLayer::CompileWholeModuleOnDemandLayer(
     ExecutionSession &ES, IRLayer &BaseLayer, JITTargetMachineBuilder JTMB,
     ModuleIndex &GlobalIndex, char GlobalManglePrefix,
     JITTargetAddress ErrorHandlerAddr)
     : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer),
       CallThroughManager(cantFail(

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

public:
   getCallThroughTrampoline(JITDylib &SourceJD, SymbolStringPtr SymbolName,
                            NotifyResolvedFunction NotifyResolved);
 
 protected:
   LazyCallThroughManager(ExecutionSession &ES,
                          JITTargetAddress ErrorHandlerAddr,
                          std::unique_ptr<TrampolinePool> TP);
 
-  JITTargetAddress callThroughToSymbol(JITTargetAddress TrampolineAddr);
+  struct ReexportsEntry {
+    JITDylib *SourceJD;
+    SymbolStringPtr SymbolName;
+  };
+
+  Expected<ReexportsEntry> findReexport(JITTargetAddress TrampolineAddr);
+  Expected<JITTargetAddress> resolveSymbol(const ReexportsEntry &RE);
+
+  Error notifyResolved(JITTargetAddress TrampolineAddr,
+                       JITTargetAddress ResolvedAddr);
+
+  JITTargetAddress reportCallThroughError(Error Err) {
+    ES.reportError(std::move(Err));
+    return ErrorHandlerAddr;
+  }
 
   void setTrampolinePool(std::unique_ptr<TrampolinePool> TP) {
     this->TP = std::move(TP);
   }
 
 private:
   using ReexportsMap =
-      std::map<JITTargetAddress, std::pair<JITDylib *, SymbolStringPtr>>;
+      std::map<JITTargetAddress, ReexportsEntry>;
 
   using NotifiersMap = std::map<JITTargetAddress, NotifyResolvedFunction>;
 
   std::mutex LCTMMutex;
   ExecutionSession &ES;
   JITTargetAddress ErrorHandlerAddr;
   std::unique_ptr<TrampolinePool> TP;
   ReexportsMap Reexports;
   NotifiersMap Notifiers;
 };
 
 /// A lazy call-through manager that builds trampolines in the current process.
 class LocalLazyCallThroughManager : public LazyCallThroughManager {
 private:
   LocalLazyCallThroughManager(ExecutionSession &ES,
                               JITTargetAddress ErrorHandlerAddr)
       : LazyCallThroughManager(ES, ErrorHandlerAddr, nullptr) {}
 
   template <typename ORCABI> Error init() {
     auto TP = LocalTrampolinePool<ORCABI>::Create(
         [this](JITTargetAddress TrampolineAddr) {
-          return callThroughToSymbol(TrampolineAddr);
+          auto Entry = findReexport(TrampolineAddr);
+          if (!Entry)
+            return reportCallThroughError(Entry.takeError());
+
+          auto ResolvedAddr = resolveSymbol(std::move(*Entry));
+          if (!ResolvedAddr)
+            return reportCallThroughError(ResolvedAddr.takeError());
+
+          if (Error Err = notifyResolved(TrampolineAddr, *ResolvedAddr))
+            return reportCallThroughError(std::move(Err));
+
+          return *ResolvedAddr;
         });
 
     if (!TP)
       return TP.takeError();
 
     setTrampolinePool(std::move(*TP));
     return Error::success();
   }

llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

Expected<JITTargetAddress> LazyCallThroughManager::getCallThroughTrampoline(
     JITDylib &SourceJD, SymbolStringPtr SymbolName,
     NotifyResolvedFunction NotifyResolved) {
   std::lock_guard<std::mutex> Lock(LCTMMutex);
   auto Trampoline = TP->getTrampoline();
 
   if (!Trampoline)
     return Trampoline.takeError();
 
-  Reexports[*Trampoline] = std::make_pair(&SourceJD, std::move(SymbolName));
+  Reexports[*Trampoline] = ReexportsEntry{&SourceJD, std::move(SymbolName)};
   Notifiers[*Trampoline] = std::move(NotifyResolved);
   return *Trampoline;
 }
 
-JITTargetAddress
-LazyCallThroughManager::callThroughToSymbol(JITTargetAddress TrampolineAddr) {
-  JITDylib *SourceJD = nullptr;
-  SymbolStringPtr SymbolName;
-
-  {
-    std::lock_guard<std::mutex> Lock(LCTMMutex);
-    auto I = Reexports.find(TrampolineAddr);
-    if (I == Reexports.end())
-      return ErrorHandlerAddr;
-    SourceJD = I->second.first;
-    SymbolName = I->second.second;
-  }
+Expected<LazyCallThroughManager::ReexportsEntry>
+LazyCallThroughManager::findReexport(JITTargetAddress TrampolineAddr) {
+    return createStringError(inconvertibleErrorCode(),
+                             "Missing reexport for trampoline address %p",
+                             TrampolineAddr);
+
+  assert(Notifiers.find(TrampolineAddr) != Notifiers.end());
+  return I->second;
 
+Expected<JITTargetAddress>
+LazyCallThroughManager::resolveSymbol(const ReexportsEntry &RE) {
   auto LookupResult = ES.lookup(
-      makeJITDylibSearchOrder(SourceJD, JITDylibLookupFlags::MatchAllSymbols),
-      SymbolName, SymbolState::Ready);
+      makeJITDylibSearchOrder(RE.SourceJD, JITDylibLookupFlags::MatchAllSymbols),
+      RE.SymbolName, SymbolState::Ready);
 
-  if (!LookupResult) {
-    ES.reportError(LookupResult.takeError());
-    return ErrorHandlerAddr;
-  }
+  if (!LookupResult)
+    return LookupResult.takeError();
 
-  auto ResolvedAddr = LookupResult->getAddress();
+  return LookupResult->getAddress();
+}
 
+Error LazyCallThroughManager::notifyResolved(JITTargetAddress TrampolineAddr,
+                                             JITTargetAddress ResolvedAddr) {
   NotifyResolvedFunction NotifyResolved;
   {
     std::lock_guard<std::mutex> Lock(LCTMMutex);
     auto I = Notifiers.find(TrampolineAddr);
     if (I != Notifiers.end()) {
       NotifyResolved = std::move(I->second);
       Notifiers.erase(I);
     }
   }
 
-  if (NotifyResolved) {
-    if (auto Err = NotifyResolved(ResolvedAddr)) {
-      ES.reportError(std::move(Err));
-      return ErrorHandlerAddr;
-    }
-  }
-
-  return ResolvedAddr;
+  return NotifyResolved ? NotifyResolved(ResolvedAddr) : Error::success();
 }
 
 Expected<std::unique_ptr<LazyCallThroughManager>>
 createLocalLazyCallThroughManager(const Triple &T, ExecutionSession &ES,
                                   JITTargetAddress ErrorHandlerAddr) {
   switch (T.getArch()) {
   default:
     return make_error<StringError>(