[ORC] Replace lookupFlags in JITSymbolResolver with getResponsibilitySet.

The new method name/behavior more closely models the way it was being used.
It also fixes an assertion that can occur when using the new ORC Core APIs,
where flags alone don't necessarily provide enough context to decide whether
the caller is responsible for materializing a given symbol (which was always
the reason this API existed).

The default implementation of getResponsibilitySet uses lookupFlags to determine
responsibility as before, so existing JITSymbolResolvers should continue to
work.

llvm-svn: 340874

3 files changed, 88 insertions, 83 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 97b1ab5aee8..24c5378e2d1 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -531,6 +531,57 @@ TEST_F(CoreAPIsStandardTest, AddAndMaterializeLazySymbol) {
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run";
 }
 
+TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) {
+  // Test that weak symbols are materialized correctly when we look them up.
+  BarSym.setFlags(static_cast<JITSymbolFlags::FlagNames>(BarSym.getFlags() |
+                                                         JITSymbolFlags::Weak));
+
+  bool BarMaterialized = false;
+  auto MU1 = llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Foo, FooSym.getFlags()}, {Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) {
+        R.resolve(SymbolMap({{Foo, FooSym}, {Bar, BarSym}})), R.emit();
+        BarMaterialized = true;
+      });
+
+  bool DuplicateBarDiscarded = false;
+  auto MU2 = llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) {
+        ADD_FAILURE() << "Attempt to materialize Bar from the wrong unit";
+        R.failMaterialization();
+      },
+      [&](const JITDylib &JD, SymbolStringPtr Name) {
+        EXPECT_EQ(Name, Bar) << "Expected \"Bar\" to be discarded";
+        DuplicateBarDiscarded = true;
+      });
+
+  cantFail(JD.define(MU1));
+  cantFail(JD.define(MU2));
+
+  bool OnResolvedRun = false;
+  bool OnReadyRun = false;
+
+  auto OnResolution = [&](Expected<SymbolMap> Result) {
+    cantFail(std::move(Result));
+    OnResolvedRun = true;
+  };
+
+  auto OnReady = [&](Error Err) {
+    cantFail(std::move(Err));
+    OnReadyRun = true;
+  };
+
+  ES.lookup({&JD}, {Bar}, std::move(OnResolution), std::move(OnReady),
+            NoDependenciesToRegister);
+
+  EXPECT_TRUE(OnResolvedRun) << "OnResolved not run";
+  EXPECT_TRUE(OnReadyRun) << "OnReady not run";
+  EXPECT_TRUE(BarMaterialized) << "Bar was not materialized at all";
+  EXPECT_TRUE(DuplicateBarDiscarded)
+      << "Duplicate bar definition not discarded";
+}
+
 TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
   bool ExpectNoMoreMaterialization = false;
   ES.setDispatchMaterialization(
diff --git a/llvm/unittests/ExecutionEngine/Orc/LegacyAPIInteropTest.cpp b/llvm/unittests/ExecutionEngine/Orc/LegacyAPIInteropTest.cpp
index 05fe921ee96..ec8ef5641f7 100644
--- a/llvm/unittests/ExecutionEngine/Orc/LegacyAPIInteropTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/LegacyAPIInteropTest.cpp
@@ -19,26 +19,31 @@ class LegacyAPIsStandardTest : public CoreAPIsBasedStandardTest {};
 namespace {
 
 TEST_F(LegacyAPIsStandardTest, TestLambdaSymbolResolver) {
+  BarSym.setFlags(static_cast<JITSymbolFlags::FlagNames>(BarSym.getFlags() |
+                                                         JITSymbolFlags::Weak));
+
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}, {Bar, BarSym}})));
 
   auto Resolver = createSymbolResolver(
-      [&](const SymbolNameSet &Symbols) { return JD.lookupFlags(Symbols); },
+      [&](const SymbolNameSet &Symbols) {
+        auto FlagsMap = JD.lookupFlags(Symbols);
+        llvm::dbgs() << "FlagsMap is " << FlagsMap << "\n";
+        SymbolNameSet Result;
+        for (auto &KV : FlagsMap)
+          if (!KV.second.isStrong())
+            Result.insert(KV.first);
+        return Result;
+      },
       [&](std::shared_ptr<AsynchronousSymbolQuery> Q, SymbolNameSet Symbols) {
         return JD.legacyLookup(std::move(Q), Symbols);
       });
 
-  SymbolNameSet Symbols({Foo, Bar, Baz});
-
-  SymbolFlagsMap SymbolFlags = Resolver->lookupFlags(Symbols);
+  auto RS = Resolver->getResponsibilitySet(SymbolNameSet({Bar, Baz}));
 
-  EXPECT_EQ(SymbolFlags.size(), 2U)
-      << "lookupFlags returned the wrong number of results";
-  EXPECT_EQ(SymbolFlags.count(Foo), 1U) << "Missing lookupFlags result for foo";
-  EXPECT_EQ(SymbolFlags.count(Bar), 1U) << "Missing lookupFlags result for bar";
-  EXPECT_EQ(SymbolFlags[Foo], FooSym.getFlags())
-      << "Incorrect lookupFlags result for Foo";
-  EXPECT_EQ(SymbolFlags[Bar], BarSym.getFlags())
-      << "Incorrect lookupFlags result for Bar";
+  EXPECT_EQ(RS.size(), 1U)
+      << "getResponsibilitySet returned the wrong number of results";
+  EXPECT_EQ(RS.count(Bar), 1U)
+      << "getResponsibilitySet result incorrect. Should be {'bar'}";
 
   bool OnResolvedRun = false;
 
@@ -59,68 +64,22 @@ TEST_F(LegacyAPIsStandardTest, TestLambdaSymbolResolver) {
 
   auto Q = std::make_shared<AsynchronousSymbolQuery>(SymbolNameSet({Foo, Bar}),
                                                      OnResolved, OnReady);
-  auto Unresolved = Resolver->lookup(std::move(Q), Symbols);
+  auto Unresolved =
+      Resolver->lookup(std::move(Q), SymbolNameSet({Foo, Bar, Baz}));
 
   EXPECT_EQ(Unresolved.size(), 1U) << "Expected one unresolved symbol";
   EXPECT_EQ(Unresolved.count(Baz), 1U) << "Expected baz to not be resolved";
   EXPECT_TRUE(OnResolvedRun) << "OnResolved was never run";
 }
 
-TEST(LegacyAPIInteropTest, QueryAgainstJITDylib) {
-
-  ExecutionSession ES(std::make_shared<SymbolStringPool>());
-  auto Foo = ES.getSymbolStringPool().intern("foo");
-
-  auto &JD = ES.createJITDylib("JD");
-  JITEvaluatedSymbol FooSym(0xdeadbeef, JITSymbolFlags::Exported);
-  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
-
-  auto LookupFlags = [&](const SymbolNameSet &Names) {
-    return JD.lookupFlags(Names);
-  };
-
-  auto Lookup = [&](std::shared_ptr<AsynchronousSymbolQuery> Query,
-                    SymbolNameSet Symbols) {
-    return JD.legacyLookup(std::move(Query), Symbols);
-  };
-
-  auto UnderlyingResolver =
-      createSymbolResolver(std::move(LookupFlags), std::move(Lookup));
-  JITSymbolResolverAdapter Resolver(ES, *UnderlyingResolver, nullptr);
-
-  JITSymbolResolver::LookupSet Names{StringRef("foo")};
-
-  auto LFR = Resolver.lookupFlags(Names);
-  EXPECT_TRUE(!!LFR) << "lookupFlags failed";
-  EXPECT_EQ(LFR->size(), 1U)
-      << "lookupFlags returned the wrong number of results";
-  EXPECT_EQ(LFR->count(*Foo), 1U)
-      << "lookupFlags did not contain a result for 'foo'";
-  EXPECT_EQ((*LFR)[*Foo], FooSym.getFlags())
-      << "lookupFlags contained the wrong result for 'foo'";
-
-  auto LR = Resolver.lookup(Names);
-  EXPECT_TRUE(!!LR) << "lookup failed";
-  EXPECT_EQ(LR->size(), 1U) << "lookup returned the wrong number of results";
-  EXPECT_EQ(LR->count(*Foo), 1U) << "lookup did not contain a result for 'foo'";
-  EXPECT_EQ((*LR)[*Foo].getFlags(), FooSym.getFlags())
-      << "lookup returned the wrong result for flags of 'foo'";
-  EXPECT_EQ((*LR)[*Foo].getAddress(), FooSym.getAddress())
-      << "lookup returned the wrong result for address of 'foo'";
-}
-
-TEST(LegacyAPIInteropTset, LegacyLookupHelpersFn) {
-  constexpr JITTargetAddress FooAddr = 0xdeadbeef;
-  JITSymbolFlags FooFlags = JITSymbolFlags::Exported;
-
+TEST_F(LegacyAPIsStandardTest, LegacyLookupHelpersFn) {
   bool BarMaterialized = false;
-  constexpr JITTargetAddress BarAddr = 0xcafef00d;
-  JITSymbolFlags BarFlags = static_cast<JITSymbolFlags::FlagNames>(
-      JITSymbolFlags::Exported | JITSymbolFlags::Weak);
+  BarSym.setFlags(static_cast<JITSymbolFlags::FlagNames>(BarSym.getFlags() |
+                                                         JITSymbolFlags::Weak));
 
   auto LegacyLookup = [&](const std::string &Name) -> JITSymbol {
     if (Name == "foo")
-      return {FooAddr, FooFlags};
+      return FooSym;
 
     if (Name == "bar") {
       auto BarMaterializer = [&]() -> Expected<JITTargetAddress> {
@@ -128,27 +87,18 @@ TEST(LegacyAPIInteropTset, LegacyLookupHelpersFn) {
         return BarAddr;
       };
 
-      return {BarMaterializer, BarFlags};
+      return {BarMaterializer, BarSym.getFlags()};
     }
 
     return nullptr;
   };
 
-  ExecutionSession ES;
-  auto Foo = ES.getSymbolStringPool().intern("foo");
-  auto Bar = ES.getSymbolStringPool().intern("bar");
-  auto Baz = ES.getSymbolStringPool().intern("baz");
-
-  SymbolNameSet Symbols({Foo, Bar, Baz});
-
-  auto SymbolFlags = lookupFlagsWithLegacyFn(Symbols, LegacyLookup);
+  auto RS =
+      getResponsibilitySetWithLegacyFn(SymbolNameSet({Bar, Baz}), LegacyLookup);
 
-  EXPECT_TRUE(!!SymbolFlags) << "Expected lookupFlagsWithLegacyFn to succeed";
-  EXPECT_EQ(SymbolFlags->size(), 2U) << "Wrong number of flags returned";
-  EXPECT_EQ(SymbolFlags->count(Foo), 1U) << "Flags for foo missing";
-  EXPECT_EQ(SymbolFlags->count(Bar), 1U) << "Flags for foo missing";
-  EXPECT_EQ((*SymbolFlags)[Foo], FooFlags) << "Wrong flags for foo";
-  EXPECT_EQ((*SymbolFlags)[Bar], BarFlags) << "Wrong flags for foo";
+  EXPECT_TRUE(!!RS) << "Expected getResponsibilitySetWithLegacyFn to succeed";
+  EXPECT_EQ(RS->size(), 1U) << "Wrong number of symbols returned";
+  EXPECT_EQ(RS->count(Bar), 1U) << "Incorrect responsibility set returned";
   EXPECT_FALSE(BarMaterialized)
       << "lookupFlags should not have materialized bar";
 
@@ -162,9 +112,11 @@ TEST(LegacyAPIInteropTset, LegacyLookupHelpersFn) {
     EXPECT_EQ(Result->count(Foo), 1U) << "Result for foo missing";
     EXPECT_EQ(Result->count(Bar), 1U) << "Result for bar missing";
     EXPECT_EQ((*Result)[Foo].getAddress(), FooAddr) << "Wrong address for foo";
-    EXPECT_EQ((*Result)[Foo].getFlags(), FooFlags) << "Wrong flags for foo";
+    EXPECT_EQ((*Result)[Foo].getFlags(), FooSym.getFlags())
+        << "Wrong flags for foo";
     EXPECT_EQ((*Result)[Bar].getAddress(), BarAddr) << "Wrong address for bar";
-    EXPECT_EQ((*Result)[Bar].getFlags(), BarFlags) << "Wrong flags for bar";
+    EXPECT_EQ((*Result)[Bar].getFlags(), BarSym.getFlags())
+        << "Wrong flags for bar";
   };
   auto OnReady = [&](Error Err) {
     EXPECT_FALSE(!!Err) << "Finalization unexpectedly failed";
@@ -172,7 +124,8 @@ TEST(LegacyAPIInteropTset, LegacyLookupHelpersFn) {
   };
 
   AsynchronousSymbolQuery Q({Foo, Bar}, OnResolved, OnReady);
-  auto Unresolved = lookupWithLegacyFn(ES, Q, Symbols, LegacyLookup);
+  auto Unresolved =
+      lookupWithLegacyFn(ES, Q, SymbolNameSet({Foo, Bar, Baz}), LegacyLookup);
 
   EXPECT_TRUE(OnResolvedRun) << "OnResolved was not run";
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run";
diff --git a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
index 420631c36ad..94b771f2640 100644
--- a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
@@ -185,7 +185,8 @@ TEST_F(RTDyldObjectLinkingLayerExecutionTest, NoDuplicateFinalization) {
 
   Resolvers[K2] = createSymbolResolver(
       [&](const SymbolNameSet &Symbols) {
-        return cantFail(lookupFlagsWithLegacyFn(Symbols, LegacyLookup));
+        return cantFail(
+            getResponsibilitySetWithLegacyFn(Symbols, LegacyLookup));
       },
       [&](std::shared_ptr<AsynchronousSymbolQuery> Query,
           const SymbolNameSet &Symbols) {