[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

7 files changed, 69 insertions, 74 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/Legacy.cpp b/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
index 18be9a042f7..517176e5f42 100644
--- a/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
@@ -48,17 +48,17 @@ JITSymbolResolverAdapter::lookup(const LookupSet &Symbols) {
   return Result;
 }
 
-Expected<JITSymbolResolverAdapter::LookupFlagsResult>
-JITSymbolResolverAdapter::lookupFlags(const LookupSet &Symbols) {
+Expected<JITSymbolResolverAdapter::LookupSet>
+JITSymbolResolverAdapter::getResponsibilitySet(const LookupSet &Symbols) {
   SymbolNameSet InternedSymbols;
   for (auto &S : Symbols)
     InternedSymbols.insert(ES.getSymbolStringPool().intern(S));
 
-  SymbolFlagsMap SymbolFlags = R.lookupFlags(InternedSymbols);
-  LookupFlagsResult Result;
-  for (auto &KV : SymbolFlags) {
-    ResolvedStrings.insert(KV.first);
-    Result[*KV.first] = KV.second;
+  auto InternedResult = R.getResponsibilitySet(InternedSymbols);
+  LookupSet Result;
+  for (auto &S : InternedResult) {
+    ResolvedStrings.insert(S);
+    Result.insert(*S);
   }
 
   return Result;
diff --git a/llvm/lib/ExecutionEngine/Orc/NullResolver.cpp b/llvm/lib/ExecutionEngine/Orc/NullResolver.cpp
index 3796e3d37bc..922fc6f021c 100644
--- a/llvm/lib/ExecutionEngine/Orc/NullResolver.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/NullResolver.cpp
@@ -14,8 +14,8 @@
 namespace llvm {
 namespace orc {
 
-SymbolFlagsMap NullResolver::lookupFlags(const SymbolNameSet &Symbols) {
-  return SymbolFlagsMap();
+SymbolNameSet NullResolver::getResponsibilitySet(const SymbolNameSet &Symbols) {
+  return Symbols;
 }
 
 SymbolNameSet
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h b/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
index b9f8a370d2f..a67215f659a 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
+++ b/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
@@ -129,20 +129,21 @@ private:
         : Stack(Stack), ExternalResolver(std::move(ExternalResolver)),
           ExternalResolverCtx(std::move(ExternalResolverCtx)) {}
 
-    orc::SymbolFlagsMap
-    lookupFlags(const orc::SymbolNameSet &Symbols) override {
-      orc::SymbolFlagsMap SymbolFlags;
+    orc::SymbolNameSet
+    getResponsibilitySet(const orc::SymbolNameSet &Symbols) override {
+      orc::SymbolNameSet Result;
 
       for (auto &S : Symbols) {
-        if (auto Sym = findSymbol(*S))
-          SymbolFlags[S] = Sym.getFlags();
-        else if (auto Err = Sym.takeError()) {
+        if (auto Sym = findSymbol(*S)) {
+          if (!Sym.getFlags().isStrong())
+            Result.insert(S);
+        } else if (auto Err = Sym.takeError()) {
           Stack.reportError(std::move(Err));
-          return orc::SymbolFlagsMap();
+          return orc::SymbolNameSet();
         }
       }
 
-      return SymbolFlags;
+      return Result;
     }
 
     orc::SymbolNameSet
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
index abe89ce70af..6515c0da360 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
+++ b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
@@ -144,26 +144,29 @@ class OrcMCJITReplacement : public ExecutionEngine {
   public:
     LinkingORCResolver(OrcMCJITReplacement &M) : M(M) {}
 
-    SymbolFlagsMap lookupFlags(const SymbolNameSet &Symbols) override {
-      SymbolFlagsMap SymbolFlags;
+    SymbolNameSet getResponsibilitySet(const SymbolNameSet &Symbols) override {
+      SymbolNameSet Result;
 
       for (auto &S : Symbols) {
         if (auto Sym = M.findMangledSymbol(*S)) {
-          SymbolFlags[S] = Sym.getFlags();
+          if (!Sym.getFlags().isStrong())
+            Result.insert(S);
         } else if (auto Err = Sym.takeError()) {
           M.reportError(std::move(Err));
-          return SymbolFlagsMap();
+          return SymbolNameSet();
         } else {
           if (auto Sym2 = M.ClientResolver->findSymbolInLogicalDylib(*S)) {
-            SymbolFlags[S] = Sym2.getFlags();
+            if (!Sym2.getFlags().isStrong())
+              Result.insert(S);
           } else if (auto Err = Sym2.takeError()) {
             M.reportError(std::move(Err));
-            return SymbolFlagsMap();
-          }
+            return SymbolNameSet();
+          } else
+            Result.insert(S);
         }
       }
 
-      return SymbolFlags;
+      return Result;
     }
 
     SymbolNameSet lookup(std::shared_ptr<AsynchronousSymbolQuery> Query,
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 1dfa90acb76..f82f5ecfed5 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -44,27 +44,13 @@ public:
     return Result;
   }
 
-  Expected<LookupFlagsResult> lookupFlags(const LookupSet &Symbols) {
-    auto &ES = MR.getTargetJITDylib().getExecutionSession();
-
-    SymbolNameSet InternedSymbols;
-
-    for (auto &S : Symbols)
-      InternedSymbols.insert(ES.getSymbolStringPool().intern(S));
-
-    SymbolFlagsMap InternedResult;
-    MR.getTargetJITDylib().withSearchOrderDo([&](const JITDylibList &JDs) {
-      // An empty search order is pathalogical, but allowed.
-      if (JDs.empty())
-        return;
-
-      assert(JDs.front() && "VSOList entry can not be null");
-      InternedResult = JDs.front()->lookupFlags(InternedSymbols);
-    });
+  Expected<LookupSet> getResponsibilitySet(const LookupSet &Symbols) {
+    LookupSet Result;
 
-    LookupFlagsResult Result;
-    for (auto &KV : InternedResult)
-      Result[*KV.first] = std::move(KV.second);
+    for (auto &KV : MR.getSymbols()) {
+      if (Symbols.count(*KV.first))
+        Result.insert(*KV.first);
+    }
 
     return Result;
   }
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
index 88b0cd0e336..d865216cf31 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/JITSymbol.cpp
@@ -94,16 +94,24 @@ LegacyJITSymbolResolver::lookup(const LookupSet &Symbols) {
 
 /// Performs flags lookup by calling findSymbolInLogicalDylib and
 ///        returning the flags value for that symbol.
-Expected<JITSymbolResolver::LookupFlagsResult>
-LegacyJITSymbolResolver::lookupFlags(const LookupSet &Symbols) {
-  JITSymbolResolver::LookupFlagsResult Result;
+Expected<JITSymbolResolver::LookupSet>
+LegacyJITSymbolResolver::getResponsibilitySet(const LookupSet &Symbols) {
+  JITSymbolResolver::LookupSet Result;
 
   for (auto &Symbol : Symbols) {
     std::string SymName = Symbol.str();
-    if (auto Sym = findSymbolInLogicalDylib(SymName))
-      Result[Symbol] = Sym.getFlags();
-    else if (auto Err = Sym.takeError())
+    if (auto Sym = findSymbolInLogicalDylib(SymName)) {
+      // If there's an existing def but it is not strong, then the caller is
+      // responsible for it.
+      if (!Sym.getFlags().isStrong())
+        Result.insert(Symbol);
+    } else if (auto Err = Sym.takeError())
       return std::move(Err);
+    else {
+      // If there is no existing definition then the caller is responsible for
+      // it.
+      Result.insert(Symbol);
+    }
   }
 
   return std::move(Result);
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index df4bd5728a6..43c99b52c34 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -204,7 +204,7 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
 
   // First, collect all weak and common symbols. We need to know if stronger
   // definitions occur elsewhere.
-  JITSymbolResolver::LookupFlagsResult SymbolFlags;
+  JITSymbolResolver::LookupSet ResponsibilitySet;
   {
     JITSymbolResolver::LookupSet Symbols;
     for (auto &Sym : Obj.symbols()) {
@@ -218,10 +218,10 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
       }
     }
 
-    if (auto FlagsResultOrErr = Resolver.lookupFlags(Symbols))
-      SymbolFlags = std::move(*FlagsResultOrErr);
+    if (auto ResultOrErr = Resolver.getResponsibilitySet(Symbols))
+      ResponsibilitySet = std::move(*ResultOrErr);
     else
-      return FlagsResultOrErr.takeError();
+      return ResultOrErr.takeError();
   }
 
   // Parse symbols
@@ -259,29 +259,26 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
     // strong.
     if (JITSymFlags->isWeak() || JITSymFlags->isCommon()) {
       // First check whether there's already a definition in this instance.
-      // FIXME: Override existing weak definitions with strong ones.
       if (GlobalSymbolTable.count(Name))
         continue;
 
-      // Then check whether we found flags for an existing symbol during the
-      // flags lookup earlier.
-      auto FlagsI = SymbolFlags.find(Name);
-      if (FlagsI == SymbolFlags.end() ||
-          (JITSymFlags->isWeak() && !FlagsI->second.isStrong()) ||
-          (JITSymFlags->isCommon() && FlagsI->second.isCommon())) {
-        if (JITSymFlags->isWeak())
-          *JITSymFlags &= ~JITSymbolFlags::Weak;
-        if (JITSymFlags->isCommon()) {
-          *JITSymFlags &= ~JITSymbolFlags::Common;
-          uint32_t Align = I->getAlignment();
-          uint64_t Size = I->getCommonSize();
-          if (!CommonAlign)
-            CommonAlign = Align;
-          CommonSize = alignTo(CommonSize, Align) + Size;
-          CommonSymbolsToAllocate.push_back(*I);
-        }
-      } else
+      // If we're not responsible for this symbol, skip it.
+      if (!ResponsibilitySet.count(Name))
         continue;
+
+      // Otherwise update the flags on the symbol to make this definition
+      // strong.
+      if (JITSymFlags->isWeak())
+        *JITSymFlags &= ~JITSymbolFlags::Weak;
+      if (JITSymFlags->isCommon()) {
+        *JITSymFlags &= ~JITSymbolFlags::Common;
+        uint32_t Align = I->getAlignment();
+        uint64_t Size = I->getCommonSize();
+        if (!CommonAlign)
+          CommonAlign = Align;
+        CommonSize = alignTo(CommonSize, Align) + Size;
+        CommonSymbolsToAllocate.push_back(*I);
+      }
     }
 
     if (Flags & SymbolRef::SF_Absolute &&