[ORC] Add a 'lookup' convenience function for finding symbols in a list of VSOs.

The lookup function takes a list of VSOs, a set of symbol names (or just one
symbol name) and a materialization function object. It returns an
Expected<SymbolMap> (if given a set of names) or an Expected<JITEvaluatedSymbol>
(if given just one name). The lookup method constructs an
AsynchronousSymbolQuery for the given names, applies that query to each VSO in
the list in turn, and then blocks waiting for the query to complete. If
threading is enabled then the materialization function object can be used to
execute the materialization on different threads. If threading is disabled the
MaterializeOnCurrentThread utility must be used.

llvm-svn: 327474

3 files changed, 241 insertions, 0 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 6ed2135b251..bddef0e83ba 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -314,6 +314,8 @@ private:
 /// @brief An ExecutionSession represents a running JIT program.
 class ExecutionSession {
 public:
+  using ErrorReporter = std::function<void(Error)>;
+
   /// @brief Construct an ExecutionEngine.
   ///
   /// SymbolStringPools may be shared between ExecutionSessions.
@@ -322,6 +324,16 @@ public:
   /// @brief Returns the SymbolStringPool for this ExecutionSession.
   SymbolStringPool &getSymbolStringPool() const { return SSP; }
 
+  /// @brief Set the error reporter function.
+  void setErrorReporter(ErrorReporter ReportError) {
+    this->ReportError = std::move(ReportError);
+  }
+
+  /// @brief Report a error for this execution session.
+  ///
+  /// Unhandled errors can be sent here to log them.
+  void reportError(Error Err) { ReportError(std::move(Err)); }
+
   /// @brief Allocate a module key for a new module to add to the JIT.
   VModuleKey allocateVModule();
 
@@ -331,10 +343,44 @@ public:
   void releaseVModule(VModuleKey Key);
 
 public:
+  static void logErrorsToStdErr(Error Err);
+
   SymbolStringPool &SSP;
   VModuleKey LastKey = 0;
+  ErrorReporter ReportError = logErrorsToStdErr;
+};
+
+/// Runs SymbolSource materializations on the current thread and reports errors
+/// to the given ExecutionSession.
+class MaterializeOnCurrentThread {
+public:
+  MaterializeOnCurrentThread(ExecutionSession &ES) : ES(ES) {}
+
+  void operator()(VSO &V, std::shared_ptr<SymbolSource> Source,
+                  SymbolNameSet Names) {
+    if (auto Err = Source->materialize(V, std::move(Names)))
+      ES.reportError(std::move(Err));
+  }
+
+private:
+  ExecutionSession &ES;
 };
 
+/// Materialization function object wrapper for the lookup method.
+using MaterializationDispatcher = std::function<void(
+    VSO &V, std::shared_ptr<SymbolSource> S, SymbolNameSet Names)>;
+
+/// @brief Look up a set of symbols by searching a list of VSOs.
+///
+/// All VSOs in the list should be non-null.
+Expected<SymbolMap> lookup(const std::vector<VSO *> &VSOs, SymbolNameSet Names,
+                           MaterializationDispatcher DispatchMaterialization);
+
+/// @brief Look up a symbol by searching a list of VSOs.
+Expected<JITEvaluatedSymbol>
+lookup(const std::vector<VSO *> VSOs, SymbolStringPtr Name,
+       MaterializationDispatcher DispatchMaterialization);
+
 } // End namespace orc
 } // End namespace llvm
 
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index b7707ade9fe..95ef62014c5 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -10,6 +10,10 @@
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 
+#if LLVM_ENABLE_THREADS
+#include <future>
+#endif
+
 namespace llvm {
 namespace orc {
 
@@ -344,6 +348,116 @@ VSO::LookupResult VSO::lookup(std::shared_ptr<AsynchronousSymbolQuery> Query,
   return {std::move(MaterializationWork), std::move(Names)};
 }
 
+Expected<SymbolMap> lookup(const std::vector<VSO *> &VSOs, SymbolNameSet Names,
+                           MaterializationDispatcher DispatchMaterialization) {
+#if LLVM_ENABLE_THREADS
+  // In the threaded case we use promises to return the results.
+  std::promise<SymbolMap> PromisedResult;
+  Error ResolutionError = Error::success();
+  std::promise<void> PromisedReady;
+  Error ReadyError = Error::success();
+  auto OnResolve = [&](Expected<SymbolMap> Result) {
+    ErrorAsOutParameter _(&ResolutionError);
+    if (Result)
+      PromisedResult.set_value(std::move(*Result));
+    else {
+      ResolutionError = Result.takeError();
+      PromisedResult.set_value({});
+    }
+  };
+  auto OnReady = [&](Error Err) {
+    ErrorAsOutParameter _(&ReadyError);
+    ReadyError = std::move(Err);
+    PromisedReady.set_value();
+  };
+#else
+  SymbolMap Result;
+  Error ResolutionError = Error::success();
+  Error ReadyError = Error::success();
+
+  auto OnResolve = [&](Expected<SymbolMap> R) {
+    ErrorAsOutParameter _(&ResolutionError);
+    if (R)
+      Result = std::move(*R);
+    else
+      ResolutionError = R.takeError();
+  };
+  auto OnReady = [&](Error Err) {
+    ErrorAsOutParameter _(&ReadyError);
+    if (Err)
+      ReadyError = std::move(Err);
+  };
+#endif
+
+  auto Query = std::make_shared<AsynchronousSymbolQuery>(
+      Names, std::move(OnResolve), std::move(OnReady));
+  SymbolNameSet UnresolvedSymbols(std::move(Names));
+
+  for (auto *VSO : VSOs) {
+
+    if (UnresolvedSymbols.empty())
+      break;
+
+    assert(VSO && "VSO pointers in VSOs list should be non-null");
+    auto LR = VSO->lookup(Query, UnresolvedSymbols);
+    UnresolvedSymbols = std::move(LR.UnresolvedSymbols);
+
+    for (auto I = LR.MaterializationWork.begin(),
+              E = LR.MaterializationWork.end();
+         I != E;) {
+      auto Tmp = I++;
+      std::shared_ptr<SymbolSource> Source = Tmp->first;
+      SymbolNameSet Names = std::move(Tmp->second);
+      LR.MaterializationWork.erase(Tmp);
+      DispatchMaterialization(*VSO, std::move(Source), std::move(Names));
+    }
+  }
+
+#if LLVM_ENABLE_THREADS
+  auto ResultFuture = PromisedResult.get_future();
+  auto Result = ResultFuture.get();
+  if (ResolutionError) {
+    // ReadyError will never be assigned. Consume the success value.
+    cantFail(std::move(ReadyError));
+    return std::move(ResolutionError);
+  }
+
+  auto ReadyFuture = PromisedReady.get_future();
+  ReadyFuture.get();
+
+  if (ReadyError)
+    return std::move(ReadyError);
+
+  return std::move(Result);
+
+#else
+  if (ResolutionError) {
+    // ReadyError will never be assigned. Consume the success value.
+    cantFail(std::move(ReadyError));
+    return std::move(ResolutionError);
+  }
+
+  if (ReadyError)
+    return std::move(ReadyError);
+
+  return Result;
+#endif
+}
+
+/// @brief Look up a symbol by searching a list of VSOs.
+Expected<JITEvaluatedSymbol>
+lookup(const std::vector<VSO *> VSOs, SymbolStringPtr Name,
+       MaterializationDispatcher DispatchMaterialization) {
+  SymbolNameSet Names({Name});
+  if (auto ResultMap =
+          lookup(VSOs, std::move(Names), std::move(DispatchMaterialization))) {
+    assert(ResultMap->size() == 1 && "Unexpected number of results");
+    assert(ResultMap->count(Name) && "Missing result for symbol");
+    return ResultMap->begin()->second;
+  } else
+    return ResultMap.takeError();
+}
+
 ExecutionSession::ExecutionSession(SymbolStringPool &SSP) : SSP(SSP) {}
 
 VModuleKey ExecutionSession::allocateVModule() { return ++LastKey; }
@@ -352,5 +466,9 @@ void ExecutionSession::releaseVModule(VModuleKey VMod) {
   // FIXME: Recycle keys.
 }
 
+void ExecutionSession::logErrorsToStdErr(Error Err) {
+  logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
+}
+
 } // End namespace orc.
 } // End namespace llvm.
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index a79f31fab0d..5a74246b1c3 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -12,6 +12,7 @@
 #include "gtest/gtest.h"
 
 #include <set>
+#include <thread>
 
 using namespace llvm;
 using namespace llvm::orc;
@@ -325,4 +326,80 @@ TEST(CoreAPIsTest, TestLambdaSymbolResolver) {
   EXPECT_TRUE(OnResolvedRun) << "OnResolved was never run";
 }
 
+TEST(CoreAPIsTest, TestLookupWithUnthreadedMaterialization) {
+  constexpr JITTargetAddress FakeFooAddr = 0xdeadbeef;
+  JITEvaluatedSymbol FooSym(FakeFooAddr, JITSymbolFlags::Exported);
+
+  SymbolStringPool SSP;
+  auto Foo = SSP.intern("foo");
+
+  auto Source = std::make_shared<SimpleSource>(
+      [&](VSO &V, SymbolNameSet Symbols) -> Error {
+        V.resolve({{Foo, FooSym}});
+        V.finalize({Foo});
+        return Error::success();
+      },
+      [](VSO &V, SymbolStringPtr Name) -> Error {
+        llvm_unreachable("Not expecting finalization");
+      });
+
+  VSO V;
+
+  SymbolFlagsMap InitialSymbols({{Foo, JITSymbolFlags::Exported}});
+  cantFail(V.defineLazy(InitialSymbols, Source));
+
+  ExecutionSession ES(SSP);
+  auto FooLookupResult =
+      cantFail(lookup({&V}, Foo, MaterializeOnCurrentThread(ES)));
+
+  EXPECT_EQ(FooLookupResult.getAddress(), FooSym.getAddress())
+      << "lookup returned an incorrect address";
+  EXPECT_EQ(FooLookupResult.getFlags(), FooSym.getFlags())
+      << "lookup returned incorrect flags";
+}
+
+TEST(CoreAPIsTest, TestLookupWithThreadedMaterialization) {
+#if LLVM_ENABLE_THREADS
+  constexpr JITTargetAddress FakeFooAddr = 0xdeadbeef;
+  JITEvaluatedSymbol FooSym(FakeFooAddr, JITSymbolFlags::Exported);
+
+  SymbolStringPool SSP;
+  auto Foo = SSP.intern("foo");
+
+  auto Source = std::make_shared<SimpleSource>(
+      [&](VSO &V, SymbolNameSet Symbols) -> Error {
+        V.resolve({{Foo, FooSym}});
+        V.finalize({Foo});
+        return Error::success();
+      },
+      [](VSO &V, SymbolStringPtr Name) -> Error {
+        llvm_unreachable("Not expecting finalization");
+      });
+
+  VSO V;
+
+  SymbolFlagsMap InitialSymbols({{Foo, JITSymbolFlags::Exported}});
+  cantFail(V.defineLazy(InitialSymbols, Source));
+
+  ExecutionSession ES(SSP);
+
+  auto MaterializeOnNewThread =
+    [&ES](VSO &V, std::shared_ptr<SymbolSource> Source, SymbolNameSet Names) {
+      std::thread(
+        [&ES, &V, Source, Names]() {
+          if (auto Err = Source->materialize(V, std::move(Names)))
+            ES.reportError(std::move(Err));
+        }).detach();
+    };
+
+  auto FooLookupResult =
+    cantFail(lookup({&V}, Foo, MaterializeOnNewThread));
+
+  EXPECT_EQ(FooLookupResult.getAddress(), FooSym.getAddress())
+      << "lookup returned an incorrect address";
+  EXPECT_EQ(FooLookupResult.getFlags(), FooSym.getFlags())
+      << "lookup returned incorrect flags";
+#endif
+}
+
 } // namespace