[ORC] Update Symbol Lookup / DefinitionGenerator system.

This patch moves definition generation out from the session lock, instead
running it under a per-dylib generator lock. It also makes the
DefinitionGenerator::tryToGenerate method optionally asynchronous: Generators
are handed an opaque LookupState object which can be captured to stop/restart
the lookup process.

The new scheme provides the following benefits and guarantees:

(1) Queries that do not need to attempt definition generation (because all
    requested symbols matched against existing definitions in the JITDylib)
    can proceed without being blocked by any running definition generators.

(2) Definition generators can capture the LookupState to continue their work
    asynchronously. This allows generators to run for an arbitrary amount of
    time without blocking a thread. Definition generators that do not need to
    run asynchronously can return without capturing the LookupState to eliminate
    unnecessary recursion and improve lookup performance.

(3) Definition generators still do not need to worry about concurrency or
    re-entrance: Since they are still run under a (per-dylib) lock, generators
    will never be re-entered concurrently, or given overlapping symbol sets to
    generate.

Finally, the new system distinguishes between symbols that are candidates for
generation (generation candidates) and symbols that failed to match for a query
(due to symbol visibility). This fixes a bug where an unresolved symbol could
trigger generation of a duplicate definition for an existing hidden symbol.

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -36,6 +36,7 @@ class MaterializationUnit;
 class MaterializationResponsibility;
 class JITDylib;
 class ResourceTracker;
+class InProgressLookupState;
 
 enum class SymbolState : uint8_t;
 
@@ -215,9 +216,19 @@ public:
 
   /// Add an element to the set. The client is responsible for checking that
   /// duplicates are not added.
-  void add(SymbolStringPtr Name,
+  SymbolLookupSet &
-           SymbolLookupFlags Flags = SymbolLookupFlags::RequiredSymbol) {
+  add(SymbolStringPtr Name,
+      SymbolLookupFlags Flags = SymbolLookupFlags::RequiredSymbol) {
     Symbols.push_back(std::make_pair(std::move(Name), Flags));
+    return *this;
+  }
+
+  /// Quickly append one lookup set to another.
+  SymbolLookupSet &append(SymbolLookupSet Other) {
+    Symbols.reserve(Symbols.size() + Other.size());
+    for (auto &KV : Other)
+      Symbols.push_back(std::move(KV));
+    return *this;
   }
 
   bool empty() const { return Symbols.empty(); }
@@ -783,6 +794,7 @@ enum class SymbolState : uint8_t {
 /// makes a callback when all symbols are available.
 class AsynchronousSymbolQuery {
   friend class ExecutionSession;
+  friend class InProgressFullLookupState;
   friend class JITDylib;
   friend class JITSymbolResolverAdapter;
   friend class MaterializationResponsibility;
@@ -829,6 +841,22 @@ private:
   SymbolState RequiredState;
 };
 
+/// Wraps state for a lookup-in-progress.
+/// DefinitionGenerators can optionally take ownership of a LookupState object
+/// to suspend a lookup-in-progress while they search for definitions.
+class LookupState {
+  friend class ExecutionSession;
+
+public:
+  /// Continue the lookup. This can be called by DefinitionGenerators
+  /// to re-start a captured query-application operation.
+  void continueLookup(Error Err);
+
+private:
+  LookupState(std::unique_ptr<InProgressLookupState> IPLS);
+  std::unique_ptr<InProgressLookupState> IPLS;
+};
+
 /// Definition generators can be attached to JITDylibs to generate new
 /// definitions for otherwise unresolved symbols during lookup.
 class DefinitionGenerator {
@@ -841,7 +869,7 @@ public:
   /// JDLookupFlags specifies whether the search should match against
   /// hidden symbols. Finally, Symbols describes the set of unresolved
   /// symbols and their associated lookup flags.
-  virtual Error tryToGenerate(LookupKind K, JITDylib &JD,
+  virtual Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                               JITDylibLookupFlags JDLookupFlags,
                               const SymbolLookupSet &LookupSet) = 0;
 };
@@ -979,13 +1007,6 @@ public:
   /// left unmodified (no symbols are removed).
   Error remove(const SymbolNameSet &Names);
 
-  /// Search the given JITDylib for the symbols in Symbols. If found, store
-  /// the flags for each symbol in Flags. If any required symbols are not found
-  /// then an error will be returned.
-  Expected<SymbolFlagsMap> lookupFlags(LookupKind K,
-                                       JITDylibLookupFlags JDLookupFlags,
-                                       SymbolLookupSet LookupSet);
-
   /// Dump current JITDylib state to OS.
   void dump(raw_ostream &OS);
 
@@ -1104,19 +1125,19 @@ private:
   void installMaterializationUnit(std::unique_ptr<MaterializationUnit> MU,
                                   ResourceTracker &RT);
 
-  void lookupFlagsImpl(SymbolFlagsMap &Result, LookupKind K,
+  // void lookupFlagsImpl(SymbolFlagsMap &Result, LookupKind K,
-                       JITDylibLookupFlags JDLookupFlags,
+  //                      JITDylibLookupFlags JDLookupFlags,
-                       SymbolLookupSet &Unresolved);
+  //                      SymbolLookupSet &Unresolved);
 
-  Error lodgeQuery(UnmaterializedInfosList &UMIs,
+  // Error lodgeQuery(UnmaterializedInfosList &UMIs,
-                   std::shared_ptr<AsynchronousSymbolQuery> &Q, LookupKind K,
+  //                  std::shared_ptr<AsynchronousSymbolQuery> &Q, LookupKind K,
-                   JITDylibLookupFlags JDLookupFlags,
+  //                  JITDylibLookupFlags JDLookupFlags,
-                   SymbolLookupSet &Unresolved);
+  //                  SymbolLookupSet &Unresolved);
 
-  Error lodgeQueryImpl(UnmaterializedInfosList &UMIs,
+  // Error lodgeQueryImpl(UnmaterializedInfosList &UMIs,
-                       std::shared_ptr<AsynchronousSymbolQuery> &Q,
+  //                      std::shared_ptr<AsynchronousSymbolQuery> &Q,
-                       LookupKind K, JITDylibLookupFlags JDLookupFlags,
+  //                      LookupKind K, JITDylibLookupFlags JDLookupFlags,
-                       SymbolLookupSet &Unresolved);
+  //                      SymbolLookupSet &Unresolved);
 
   void detachQueryHelper(AsynchronousSymbolQuery &Q,
                          const SymbolNameSet &QuerySymbols);
@@ -1154,11 +1175,12 @@ private:
 
   ExecutionSession &ES;
   std::string JITDylibName;
+  std::mutex GeneratorsMutex;
   bool Open = true;
   SymbolTable Symbols;
   UnmaterializedInfosMap UnmaterializedInfos;
   MaterializingInfosMap MaterializingInfos;
-  std::vector<std::unique_ptr<DefinitionGenerator>> DefGenerators;
+  std::vector<std::shared_ptr<DefinitionGenerator>> DefGenerators;
   JITDylibSearchOrder LinkOrder;
   ResourceTrackerSP DefaultTracker;
 
@@ -1199,7 +1221,10 @@ public:
 
 /// An ExecutionSession represents a running JIT program.
 class ExecutionSession {
+  friend class InProgressLookupFlagsState;
+  friend class InProgressFullLookupState;
   friend class JITDylib;
+  friend class LookupState;
   friend class MaterializationResponsibility;
   friend class ResourceTracker;
 
@@ -1290,7 +1315,18 @@ public:
     return *this;
   }
 
-  /// Search the given JITDylib list for the given symbols.
+  /// Search the given JITDylibs to find the flags associated with each of the
+  /// given symbols.
+  void lookupFlags(LookupKind K, JITDylibSearchOrder SearchOrder,
+                   SymbolLookupSet Symbols,
+                   unique_function<void(Expected<SymbolFlagsMap>)> OnComplete);
+
+  /// Blocking version of lookupFlags.
+  Expected<SymbolFlagsMap> lookupFlags(LookupKind K,
+                                       JITDylibSearchOrder SearchOrder,
+                                       SymbolLookupSet Symbols);
+
+  /// Search the given JITDylibs for the given symbols.
   ///
   /// SearchOrder lists the JITDylibs to search. For each dylib, the associated
   /// boolean indicates whether the search should match against non-exported
@@ -1372,7 +1408,7 @@ private:
     MU->materialize(std::move(MR));
   }
 
-  void runOutstandingMUs();
+  void dispatchOutstandingMUs();
 
   static std::unique_ptr<MaterializationResponsibility>
   createMaterializationResponsibility(ResourceTracker &RT,
@@ -1390,8 +1426,36 @@ private:
   void transferResourceTracker(ResourceTracker &DstRT, ResourceTracker &SrcRT);
   void destroyResourceTracker(ResourceTracker &RT);
 
+  // State machine functions for query application..
 
-  /// State machine functions for MaterializationResponsibility.
+  /// IL_updateCandidatesFor is called to remove already-defined symbols that
+  /// match a given query from the set of candidate symbols to generate
+  /// definitions for (no need to generate a definition if one already exists).
+  Error IL_updateCandidatesFor(JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+                               SymbolLookupSet &Candidates,
+                               SymbolLookupSet *NonCandidates);
+
+  /// OL_applyQueryPhase1 is an optionally re-startable loop for triggering
+  /// definition generation. It is called when a lookup is performed, and again
+  /// each time that LookupState::continueLookup is called.
+  void OL_applyQueryPhase1(std::unique_ptr<InProgressLookupState> IPLS,
+                           Error Err);
+
+  /// OL_completeLookup is run once phase 1 successfully completes for a lookup
+  /// call. It attempts to attach the symbol to all symbol table entries and
+  /// collect all MaterializationUnits to dispatch. If this method fails then
+  /// all MaterializationUnits will be left un-materialized.
+  void OL_completeLookup(std::unique_ptr<InProgressLookupState> IPLS,
+                         std::shared_ptr<AsynchronousSymbolQuery> Q,
+                         RegisterDependenciesFunction RegisterDependencies);
+
+  /// OL_completeLookupFlags is run once phase 1 successfully completes for a
+  /// lookupFlags call.
+  void OL_completeLookupFlags(
+      std::unique_ptr<InProgressLookupState> IPLS,
+      unique_function<void(Expected<SymbolFlagsMap>)> OnComplete);
+
+  // State machine functions for MaterializationResponsibility.
   void OL_destroyMaterializationResponsibility(
       MaterializationResponsibility &MR);
   SymbolNameSet OL_getRequestedSymbols(const MaterializationResponsibility &MR);
@@ -1454,8 +1518,8 @@ Error MaterializationResponsibility::withResourceKeyDo(Func &&F) const {
 template <typename GeneratorT>
 GeneratorT &JITDylib::addGenerator(std::unique_ptr<GeneratorT> DefGenerator) {
   auto &G = *DefGenerator;
-  ES.runSessionLocked(
+  std::lock_guard<std::mutex> Lock(GeneratorsMutex);
-      [&]() { DefGenerators.push_back(std::move(DefGenerator)); });
+  DefGenerators.push_back(std::move(DefGenerator));
   return G;
 }
 
@@ -1546,7 +1610,7 @@ public:
                      JITDylibLookupFlags SourceJDLookupFlags,
                      SymbolPredicate Allow = SymbolPredicate());
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &LookupSet) override;
 

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -254,7 +254,7 @@ public:
     return Load(nullptr, GlobalPrefix, std::move(Allow));
   }
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 
@@ -292,7 +292,7 @@ public:
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Create(ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer);
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 

llvm/include/llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h

@@ -50,7 +50,7 @@ public:
     return Load(TPC, nullptr);
   }
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -261,8 +261,8 @@ DynamicLibrarySearchGenerator::Load(const char *FileName, char GlobalPrefix,
 }
 
 Error DynamicLibrarySearchGenerator::tryToGenerate(
-    LookupKind K, JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+    LookupState &LS, LookupKind K, JITDylib &JD,
-    const SymbolLookupSet &Symbols) {
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
   orc::SymbolMap NewSymbols;
 
   bool HasGlobalPrefix = (GlobalPrefix != '\0');
@@ -365,8 +365,8 @@ StaticLibraryDefinitionGenerator::Create(
 }
 
 Error StaticLibraryDefinitionGenerator::tryToGenerate(
-    LookupKind K, JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+    LookupState &LS, LookupKind K, JITDylib &JD,
-    const SymbolLookupSet &Symbols) {
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
 
   // Don't materialize symbols from static archives unless this is a static
   // lookup.

llvm/lib/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.cpp

@@ -24,8 +24,8 @@ TPCDynamicLibrarySearchGenerator::Load(TargetProcessControl &TPC,
 }
 
 Error TPCDynamicLibrarySearchGenerator::tryToGenerate(
-    LookupKind K, JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+    LookupState &LS, LookupKind K, JITDylib &JD,
-    const SymbolLookupSet &Symbols) {
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
 
   if (Symbols.empty())
     return Error::success();

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -563,7 +563,7 @@ Error LLVMJITLinkObjectLinkingLayer::add(ResourceTrackerSP RT,
 
 class PhonyExternalsGenerator : public DefinitionGenerator {
 public:
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &LookupSet) override {
     SymbolMap PhonySymbols;

llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -272,9 +272,11 @@ TEST_F(CoreAPIsStandardTest, LookupFlagsTest) {
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
   cantFail(JD.define(std::move(MU)));
 
-  auto SymbolFlags = cantFail(JD.lookupFlags(
+  auto SymbolFlags = cantFail(ES.lookupFlags(
-      LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
+      LookupKind::Static,
-      SymbolLookupSet({Foo, Bar, Baz})));
+      {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
+      SymbolLookupSet({Foo, Bar, Baz},
+                      SymbolLookupFlags::WeaklyReferencedSymbol)));
 
   EXPECT_EQ(SymbolFlags.size(), 2U)
       << "Returned symbol flags contains unexpected results";
@@ -291,8 +293,8 @@ TEST_F(CoreAPIsStandardTest, LookupWithGeneratorFailure) {
 
   class BadGenerator : public DefinitionGenerator {
   public:
-    Error tryToGenerate(LookupKind K, JITDylib &, JITDylibLookupFlags,
+    Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &,
-                        const SymbolLookupSet &) override {
+                        JITDylibLookupFlags, const SymbolLookupSet &) override {
       return make_error<StringError>("BadGenerator", inconvertibleErrorCode());
     }
   };
@@ -300,8 +302,8 @@ TEST_F(CoreAPIsStandardTest, LookupWithGeneratorFailure) {
   JD.addGenerator(std::make_unique<BadGenerator>());
 
   EXPECT_THAT_ERROR(
-      JD.lookupFlags(LookupKind::Static,
+      ES.lookupFlags(LookupKind::Static,
-                     JITDylibLookupFlags::MatchExportedSymbolsOnly,
+                     {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
                      SymbolLookupSet(Foo))
           .takeError(),
       Failed<StringError>())
@@ -399,9 +401,11 @@ TEST_F(CoreAPIsStandardTest, TestReexportsGenerator) {
   JD.addGenerator(std::make_unique<ReexportsGenerator>(
       JD2, JITDylibLookupFlags::MatchExportedSymbolsOnly, Filter));
 
-  auto Flags = cantFail(JD.lookupFlags(
+  auto Flags = cantFail(ES.lookupFlags(
-      LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
+      LookupKind::Static,
-      SymbolLookupSet({Foo, Bar, Baz})));
+      {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
+      SymbolLookupSet({Foo, Bar, Baz},
+                      SymbolLookupFlags::WeaklyReferencedSymbol)));
   EXPECT_EQ(Flags.size(), 1U) << "Unexpected number of results";
   EXPECT_EQ(Flags[Foo], FooSym.getFlags()) << "Unexpected flags for Foo";
 
@@ -1045,12 +1049,14 @@ TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
 }
 
 TEST_F(CoreAPIsStandardTest, GeneratorTest) {
-  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
+  JITEvaluatedSymbol BazHiddenSym(
+      BazSym.getAddress(), BazSym.getFlags() & ~JITSymbolFlags::Exported);
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}, {Baz, BazHiddenSym}})));
 
   class TestGenerator : public DefinitionGenerator {
   public:
     TestGenerator(SymbolMap Symbols) : Symbols(std::move(Symbols)) {}
-    Error tryToGenerate(LookupKind K, JITDylib &JD,
+    Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                         JITDylibLookupFlags JDLookupFlags,
                         const SymbolLookupSet &Names) override {
       SymbolMap NewDefs;
@@ -1069,10 +1075,13 @@ TEST_F(CoreAPIsStandardTest, GeneratorTest) {
     SymbolMap Symbols;
   };
 
-  JD.addGenerator(std::make_unique<TestGenerator>(SymbolMap({{Bar, BarSym}})));
+  JD.addGenerator(std::make_unique<TestGenerator>(
+      SymbolMap({{Bar, BarSym}, {Baz, BazSym}})));
 
   auto Result = cantFail(
-      ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));
+      ES.lookup(makeJITDylibSearchOrder(&JD),
+                SymbolLookupSet({Foo, Bar})
+                    .add(Baz, SymbolLookupFlags::WeaklyReferencedSymbol)));
 
   EXPECT_EQ(Result.count(Bar), 1U) << "Expected to find fallback def for 'bar'";
   EXPECT_EQ(Result[Bar].getAddress(), BarSym.getAddress())

llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp

@@ -136,9 +136,10 @@ TEST_F(ResourceTrackerStandardTest,
   auto RT = JD.createResourceTracker();
   cantFail(JD.define(std::move(MU), RT));
   cantFail(RT->remove());
-  auto SymFlags = cantFail(JD.lookupFlags(
+  auto SymFlags = cantFail(ES.lookupFlags(
-      LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
+      LookupKind::Static,
-      SymbolLookupSet(Foo)));
+      {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
+      SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
 
   EXPECT_EQ(SymFlags.size(), 0U)
       << "Symbols should have been removed from the symbol table";
@@ -175,9 +176,10 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {
   cantFail(JD.define(std::move(MU), RT));
   cantFail(ES.lookup({&JD}, Foo));
   cantFail(RT->remove());
-  auto SymFlags = cantFail(JD.lookupFlags(
+  auto SymFlags = cantFail(ES.lookupFlags(
-      LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
+      LookupKind::Static,
-      SymbolLookupSet(Foo)));
+      {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
+      SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
 
   EXPECT_EQ(SymFlags.size(), 0U)
       << "Symbols should have been removed from the symbol table";
@@ -217,9 +219,10 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
       NoDependenciesToRegister);
 
   cantFail(RT->remove());
-  auto SymFlags = cantFail(JD.lookupFlags(
+  auto SymFlags = cantFail(ES.lookupFlags(
-      LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
+      LookupKind::Static,
-      SymbolLookupSet(Foo)));
+      {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
+      SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
 
   EXPECT_EQ(SymFlags.size(), 0U)
       << "Symbols should have been removed from the symbol table";