1 files changed, 149 insertions, 110 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index d53c4558e0c..c0afbc6be06 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -59,57 +59,51 @@ private:
   DestructorFunction Destructor;
 };
 
-TEST_F(CoreAPIsStandardTest, BasicSuccessfulLookup) {
+
+TEST_F(CoreAPIsStandardTest, AsynchronousSymbolQuerySuccessfulResolutionOnly) {
   bool OnResolutionRun = false;
   bool OnReadyRun = false;
-
-  auto OnResolution = [&](Expected<SymbolMap> Result) {
-    EXPECT_TRUE(!!Result) << "Resolution unexpectedly returned error";
-    auto &Resolved = *Result;
-    auto I = Resolved.find(Foo);
-    EXPECT_NE(I, Resolved.end()) << "Could not find symbol definition";
-    EXPECT_EQ(I->second.getAddress(), FooAddr)
-        << "Resolution returned incorrect result";
-    OnResolutionRun = true;
-  };
+  auto OnResolution =
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> Result) {
+        EXPECT_TRUE(!!Result) << "Resolution unexpectedly returned error";
+        auto &Resolved = Result->Symbols;
+        auto I = Resolved.find(Foo);
+        EXPECT_NE(I, Resolved.end()) << "Could not find symbol definition";
+        EXPECT_EQ(I->second.getAddress(), FooAddr)
+            << "Resolution returned incorrect result";
+        OnResolutionRun = true;
+      };
   auto OnReady = [&](Error Err) {
     cantFail(std::move(Err));
     OnReadyRun = true;
   };
 
-  std::shared_ptr<MaterializationResponsibility> FooMR;
-
-  cantFail(V.define(llvm::make_unique<SimpleMaterializationUnit>(
-      SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
-      [&](MaterializationResponsibility R) {
-        FooMR = std::make_shared<MaterializationResponsibility>(std::move(R));
-      })));
-
-  ES.lookup({&V}, {Foo}, OnResolution, OnReady, NoDependenciesToRegister);
+  AsynchronousSymbolQuery Q(SymbolNameSet({Foo}), OnResolution, OnReady);
 
-  EXPECT_FALSE(OnResolutionRun) << "Should not have been resolved yet";
-  EXPECT_FALSE(OnReadyRun) << "Should not have been marked ready yet";
+  Q.resolve(Foo, FooSym);
 
-  FooMR->resolve({{Foo, FooSym}});
+  EXPECT_TRUE(Q.isFullyResolved()) << "Expected query to be fully resolved";
 
-  EXPECT_TRUE(OnResolutionRun) << "Should have been resolved";
-  EXPECT_FALSE(OnReadyRun) << "Should not have been marked ready yet";
+  if (!Q.isFullyResolved())
+    return;
 
-  FooMR->finalize();
+  Q.handleFullyResolved();
 
-  EXPECT_TRUE(OnReadyRun) << "Should have been marked ready";
+  EXPECT_TRUE(OnResolutionRun) << "OnResolutionCallback was not run";
+  EXPECT_FALSE(OnReadyRun) << "OnReady unexpectedly run";
 }
 
 TEST_F(CoreAPIsStandardTest, ExecutionSessionFailQuery) {
   bool OnResolutionRun = false;
   bool OnReadyRun = false;
 
-  auto OnResolution = [&](Expected<SymbolMap> Result) {
-    EXPECT_FALSE(!!Result) << "Resolution unexpectedly returned success";
-    auto Msg = toString(Result.takeError());
-    EXPECT_EQ(Msg, "xyz") << "Resolution returned incorrect result";
-    OnResolutionRun = true;
-  };
+  auto OnResolution =
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> Result) {
+        EXPECT_FALSE(!!Result) << "Resolution unexpectedly returned success";
+        auto Msg = toString(Result.takeError());
+        EXPECT_EQ(Msg, "xyz") << "Resolution returned incorrect result";
+        OnResolutionRun = true;
+      };
   auto OnReady = [&](Error Err) {
     cantFail(std::move(Err));
     OnReadyRun = true;
@@ -117,28 +111,62 @@ TEST_F(CoreAPIsStandardTest, ExecutionSessionFailQuery) {
 
   AsynchronousSymbolQuery Q(SymbolNameSet({Foo}), OnResolution, OnReady);
 
-  ES.legacyFailQuery(Q,
-                     make_error<StringError>("xyz", inconvertibleErrorCode()));
+  ES.failQuery(Q, make_error<StringError>("xyz", inconvertibleErrorCode()));
 
   EXPECT_TRUE(OnResolutionRun) << "OnResolutionCallback was not run";
   EXPECT_FALSE(OnReadyRun) << "OnReady unexpectedly run";
 }
 
-TEST_F(CoreAPIsStandardTest, EmptyLookup) {
-  bool OnResolvedRun = false;
+TEST_F(CoreAPIsStandardTest, SimpleAsynchronousSymbolQueryAgainstVSO) {
+  bool OnResolutionRun = false;
   bool OnReadyRun = false;
 
-  auto OnResolution = [&](Expected<SymbolMap> Result) {
-    cantFail(std::move(Result));
-    OnResolvedRun = true;
-  };
+  auto OnResolution =
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> Result) {
+        EXPECT_TRUE(!!Result) << "Query unexpectedly returned error";
+        auto &Resolved = Result->Symbols;
+        auto I = Resolved.find(Foo);
+        EXPECT_NE(I, Resolved.end()) << "Could not find symbol definition";
+        EXPECT_EQ(I->second.getAddress(), FooSym.getAddress())
+            << "Resolution returned incorrect result";
+        OnResolutionRun = true;
+      };
 
   auto OnReady = [&](Error Err) {
     cantFail(std::move(Err));
     OnReadyRun = true;
   };
 
-  ES.lookup({&V}, {}, OnResolution, OnReady, NoDependenciesToRegister);
+  SymbolNameSet Names({Foo});
+
+  auto Q =
+      std::make_shared<AsynchronousSymbolQuery>(Names, OnResolution, OnReady);
+
+  auto Defs = absoluteSymbols({{Foo, FooSym}});
+  cantFail(V.define(Defs));
+  assert(Defs == nullptr && "Defs should have been accepted");
+  V.lookup(Q, Names);
+
+  EXPECT_TRUE(OnResolutionRun) << "OnResolutionCallback was not run";
+  EXPECT_TRUE(OnReadyRun) << "OnReady was not run";
+}
+
+TEST_F(CoreAPIsStandardTest, EmptyVSOAndQueryLookup) {
+  bool OnResolvedRun = false;
+  bool OnReadyRun = false;
+
+  auto Q = std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet(),
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+        cantFail(std::move(RR));
+        OnResolvedRun = true;
+      },
+      [&](Error Err) {
+        cantFail(std::move(Err));
+        OnReadyRun = true;
+      });
+
+  V.lookup(std::move(Q), {});
 
   EXPECT_TRUE(OnResolvedRun) << "OnResolved was not run for empty query";
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run for empty query";
@@ -154,8 +182,8 @@ TEST_F(CoreAPIsStandardTest, ChainedVSOLookup) {
 
   auto Q = std::make_shared<AsynchronousSymbolQuery>(
       SymbolNameSet({Foo}),
-      [&](Expected<SymbolMap> Result) {
-        cantFail(std::move(Result));
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+        cantFail(std::move(RR));
         OnResolvedRun = true;
       },
       [&](Error Err) {
@@ -163,7 +191,7 @@ TEST_F(CoreAPIsStandardTest, ChainedVSOLookup) {
         OnReadyRun = true;
       });
 
-  V2.legacyLookup(Q, V.legacyLookup(Q, {Foo}));
+  V2.lookup(Q, V.lookup(Q, {Foo}));
 
   EXPECT_TRUE(OnResolvedRun) << "OnResolved was not run for empty query";
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run for empty query";
@@ -240,15 +268,20 @@ TEST_F(CoreAPIsStandardTest, TestTrivialCircularDependency) {
   cantFail(V.define(FooMU));
 
   bool FooReady = false;
-  auto OnResolution = [](Expected<SymbolMap> R) { cantFail(std::move(R)); };
-  auto OnReady = [&](Error Err) {
-    cantFail(std::move(Err));
-    FooReady = true;
-  };
+  auto Q =
+    std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet({ Foo }),
+      [](Expected<AsynchronousSymbolQuery::ResolutionResult> R) {
+        cantFail(std::move(R));
+      },
+      [&](Error Err) {
+        cantFail(std::move(Err));
+        FooReady = true;
+      });
 
-  ES.lookup({&V}, {Foo}, std::move(OnResolution), std::move(OnReady),
-            NoDependenciesToRegister);
+  V.lookup(std::move(Q), { Foo });
 
+  FooR->addDependencies({{&V, {Foo}}});
   FooR->resolve({{Foo, FooSym}});
   FooR->finalize();
 
@@ -290,52 +323,54 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneVSO) {
   // Query each of the symbols to trigger materialization.
   bool FooResolved = false;
   bool FooReady = false;
-
-  auto OnFooResolution = [&](Expected<SymbolMap> Result) {
-    cantFail(std::move(Result));
-    FooResolved = true;
-  };
-
-  auto OnFooReady = [&](Error Err) {
-    cantFail(std::move(Err));
-    FooReady = true;
-  };
-
-  // Issue a lookup for Foo. Use NoDependenciesToRegister: We're going to add
-  // the dependencies manually below.
-  ES.lookup({&V}, {Foo}, std::move(OnFooResolution), std::move(OnFooReady),
-            NoDependenciesToRegister);
+  auto FooQ = std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet({Foo}),
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+        cantFail(std::move(RR));
+        FooResolved = true;
+      },
+      [&](Error Err) {
+        cantFail(std::move(Err));
+        FooReady = true;
+      });
+  {
+    auto Unresolved = V.lookup(FooQ, {Foo});
+    EXPECT_TRUE(Unresolved.empty()) << "Failed to resolve \"Foo\"";
+  }
 
   bool BarResolved = false;
   bool BarReady = false;
-  auto OnBarResolution = [&](Expected<SymbolMap> Result) {
-    cantFail(std::move(Result));
-    BarResolved = true;
-  };
-
-  auto OnBarReady = [&](Error Err) {
-    cantFail(std::move(Err));
-    BarReady = true;
-  };
-
-  ES.lookup({&V}, {Bar}, std::move(OnBarResolution), std::move(OnBarReady),
-            NoDependenciesToRegister);
+  auto BarQ = std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet({Bar}),
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+        cantFail(std::move(RR));
+        BarResolved = true;
+      },
+      [&](Error Err) {
+        cantFail(std::move(Err));
+        BarReady = true;
+      });
+  {
+    auto Unresolved = V.lookup(BarQ, {Bar});
+    EXPECT_TRUE(Unresolved.empty()) << "Failed to resolve \"Bar\"";
+  }
 
   bool BazResolved = false;
   bool BazReady = false;
-
-  auto OnBazResolution = [&](Expected<SymbolMap> Result) {
-    cantFail(std::move(Result));
-    BazResolved = true;
-  };
-
-  auto OnBazReady = [&](Error Err) {
-    cantFail(std::move(Err));
-    BazReady = true;
-  };
-
-  ES.lookup({&V}, {Baz}, std::move(OnBazResolution), std::move(OnBazReady),
-            NoDependenciesToRegister);
+  auto BazQ = std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet({Baz}),
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+        cantFail(std::move(RR));
+        BazResolved = true;
+      },
+      [&](Error Err) {
+        cantFail(std::move(Err));
+        BazReady = true;
+      });
+  {
+    auto Unresolved = V.lookup(BazQ, {Baz});
+    EXPECT_TRUE(Unresolved.empty()) << "Failed to resolve \"Baz\"";
+  }
 
   // Add a circular dependency: Foo -> Bar, Bar -> Baz, Baz -> Foo.
   FooR->addDependencies({{&V, SymbolNameSet({Bar})}});
@@ -443,23 +478,28 @@ TEST_F(CoreAPIsStandardTest, AddAndMaterializeLazySymbol) {
   bool OnResolutionRun = false;
   bool OnReadyRun = false;
 
-  auto OnResolution = [&](Expected<SymbolMap> Result) {
-    EXPECT_TRUE(!!Result) << "Resolution unexpectedly returned error";
-    auto I = Result->find(Foo);
-    EXPECT_NE(I, Result->end()) << "Could not find symbol definition";
-    EXPECT_EQ(I->second.getAddress(), FooSym.getAddress())
-        << "Resolution returned incorrect result";
-    OnResolutionRun = true;
-  };
+  auto OnResolution =
+      [&](Expected<AsynchronousSymbolQuery::ResolutionResult> Result) {
+        EXPECT_TRUE(!!Result) << "Resolution unexpectedly returned error";
+        auto I = Result->Symbols.find(Foo);
+        EXPECT_NE(I, Result->Symbols.end())
+            << "Could not find symbol definition";
+        EXPECT_EQ(I->second.getAddress(), FooSym.getAddress())
+            << "Resolution returned incorrect result";
+        OnResolutionRun = true;
+      };
 
   auto OnReady = [&](Error Err) {
     cantFail(std::move(Err));
     OnReadyRun = true;
   };
 
-  ES.lookup({&V}, Names, std::move(OnResolution), std::move(OnReady),
-            NoDependenciesToRegister);
+  auto Q =
+      std::make_shared<AsynchronousSymbolQuery>(Names, OnResolution, OnReady);
 
+  auto Unresolved = V.lookup(std::move(Q), Names);
+
+  EXPECT_TRUE(Unresolved.empty()) << "Could not find Foo in dylib";
   EXPECT_TRUE(FooMaterialized) << "Foo was not materialized";
   EXPECT_TRUE(BarDiscarded) << "Bar was not discarded";
   EXPECT_TRUE(OnResolutionRun) << "OnResolutionCallback was not run";
@@ -674,14 +714,13 @@ TEST_F(CoreAPIsStandardTest, TestMaterializeWeakSymbol) {
       });
 
   cantFail(V.define(MU));
-  auto OnResolution = [](Expected<SymbolMap> Result) {
-    cantFail(std::move(Result));
-  };
-
-  auto OnReady = [](Error Err) { cantFail(std::move(Err)); };
-
-  ES.lookup({&V}, {Foo}, std::move(OnResolution), std::move(OnReady),
-            NoDependenciesToRegister);
+  auto Q = std::make_shared<AsynchronousSymbolQuery>(
+      SymbolNameSet({Foo}),
+      [](Expected<AsynchronousSymbolQuery::ResolutionResult> R) {
+        cantFail(std::move(R));
+      },
+      [](Error Err) { cantFail(std::move(Err)); });
+  V.lookup(std::move(Q), SymbolNameSet({Foo}));
 
   auto MU2 = llvm::make_unique<SimpleMaterializationUnit>(
       SymbolFlagsMap({{Foo, JITSymbolFlags::Exported}}),