[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

58

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

108

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/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>>;
+  using ReexportsMap = std::map<JITTargetAddress, ReexportsEntry>;
 
   using NotifiersMap = std::map<JITTargetAddress, NotifyResolvedFunction>;
 
   std::mutex LCTMMutex;
   ExecutionSession &ES;
   JITTargetAddress ErrorHandlerAddr;
   std::unique_ptr<TrampolinePool> TP;
   ReexportsMap Reexports;
template <typename ORCABI> Error init() {
 
     if (!TP)
       return TP.takeError();
 
     setTrampolinePool(std::move(*TP));
     return Error::success();
   }
 
+  JITTargetAddress callThroughToSymbol(JITTargetAddress 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;
+  }
+
 public:
   /// Create a LocalLazyCallThroughManager using the given ABI. See
   /// createLocalLazyCallThroughManager.
   template <typename ORCABI>
   static Expected<std::unique_ptr<LocalLazyCallThroughManager>>
   Create(ExecutionSession &ES, JITTargetAddress ErrorHandlerAddr) {
     auto LLCTM = std::unique_ptr<LocalLazyCallThroughManager>(
         new LocalLazyCallThroughManager(ES, ErrorHandlerAddr));

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;
 
-  auto LookupResult = ES.lookup(
-      makeJITDylibSearchOrder(SourceJD, JITDylibLookupFlags::MatchAllSymbols),
-      SymbolName, SymbolState::Ready);
+Expected<JITTargetAddress>
+LazyCallThroughManager::resolveSymbol(const ReexportsEntry &RE) {
+  auto LookupResult =
+      ES.lookup(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>(