[ORC] Replace SymbolResolvers in the new ORC layers with search orders on VSOs.

A search order is a list of VSOs to be searched linearly to find symbols. Each
VSO now has a search order that will be used when fixing up definitions in that
VSO. Each VSO's search order defaults to just that VSO itself.

This is a first step towards removing symbol resolvers from ORC altogether. In
practice symbol resolvers tended to be used to implement a search order anyway,
sometimes with additional programatic generation of symbols. Now that VSOs
support programmatic generation of definitions via fallback generators, search
orders provide a cleaner way to achieve the desired effect (while removing a lot
of boilerplate).

llvm-svn: 337593

5 files changed, 148 insertions, 97 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index ebd8d43e710..d42e7b05ba6 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -130,21 +130,18 @@ namespace orc {
 
 class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
 public:
-  ExtractingIRMaterializationUnit(
-      ExecutionSession &ES, CompileOnDemandLayer2 &Parent,
-      std::unique_ptr<Module> M,
-      std::shared_ptr<SymbolResolver> BackingResolver)
-      : IRMaterializationUnit(ES, std::move(M)), Parent(Parent),
-        BackingResolver(std::move(BackingResolver)) {}
-
-  ExtractingIRMaterializationUnit(
-      std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
-      SymbolNameToDefinitionMap SymbolToDefinition,
-      CompileOnDemandLayer2 &Parent,
-      std::shared_ptr<SymbolResolver> BackingResolver)
+  ExtractingIRMaterializationUnit(ExecutionSession &ES,
+                                  CompileOnDemandLayer2 &Parent,
+                                  std::unique_ptr<Module> M)
+      : IRMaterializationUnit(ES, std::move(M)), Parent(Parent) {}
+
+  ExtractingIRMaterializationUnit(std::unique_ptr<Module> M,
+                                  SymbolFlagsMap SymbolFlags,
+                                  SymbolNameToDefinitionMap SymbolToDefinition,
+                                  CompileOnDemandLayer2 &Parent)
       : IRMaterializationUnit(std::move(M), std::move(SymbolFlags),
                               std::move(SymbolToDefinition)),
-        Parent(Parent), BackingResolver(std::move(BackingResolver)) {}
+        Parent(Parent) {}
 
 private:
   void materialize(MaterializationResponsibility R) override {
@@ -153,10 +150,6 @@ private:
     //        together. This could be used, for example, to provide a specific
     //        memory manager instance to the linking layer.
 
-    // FIXME: The derived constructor should *only* look for the names of
-    //        original function definitions in the target VSO. All other
-    //        symbols should be looked up in the backing resolver.
-
     auto RequestedSymbols = R.getRequestedSymbols();
 
     // Extract the requested functions into a new module.
@@ -201,12 +194,12 @@ private:
              "of entries");
       R.replace(llvm::make_unique<ExtractingIRMaterializationUnit>(
           std::move(M), std::move(DelegatedSymbolFlags),
-          std::move(DelegatedSymbolToDefinition), Parent, BackingResolver));
+          std::move(DelegatedSymbolToDefinition), Parent));
     }
 
     if (ExtractedFunctionsModule)
-      Parent.emitExtractedFunctionsModule(
-          std::move(R), std::move(ExtractedFunctionsModule), BackingResolver);
+      Parent.emitExtractedFunctionsModule(std::move(R),
+                                          std::move(ExtractedFunctionsModule));
   }
 
   void discard(const VSO &V, SymbolStringPtr Name) override {
@@ -218,19 +211,14 @@ private:
 
   mutable std::mutex SourceModuleMutex;
   CompileOnDemandLayer2 &Parent;
-  std::shared_ptr<SymbolResolver> BackingResolver;
 };
 
 CompileOnDemandLayer2::CompileOnDemandLayer2(
     ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
     IndirectStubsManagerBuilder BuildIndirectStubsManager,
-    GetSymbolResolverFunction GetSymbolResolver,
-    SetSymbolResolverFunction SetSymbolResolver,
     GetAvailableContextFunction GetAvailableContext)
     : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)),
-      GetSymbolResolver(std::move(GetSymbolResolver)),
-      SetSymbolResolver(std::move(SetSymbolResolver)),
       GetAvailableContext(std::move(GetAvailableContext)) {}
 
 Error CompileOnDemandLayer2::add(VSO &V, VModuleKey K,
@@ -306,19 +294,14 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
     StubInits[*KV.first] = KV.second;
 
   // Build the function-body-extracting materialization unit.
-  auto SR = GetSymbolResolver(K);
   if (auto Err = R.getTargetVSO().define(
-          llvm::make_unique<ExtractingIRMaterializationUnit>(
-              ES, *this, std::move(M), SR))) {
+          llvm::make_unique<ExtractingIRMaterializationUnit>(ES, *this,
+                                                             std::move(M)))) {
     ES.reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
-  // Replace the fallback symbol resolver: We will re-use M's VModuleKey for
-  // the GlobalsModule.
-  SetSymbolResolver(K, SR);
-
   // Build the stubs.
   // FIXME: Remove function bodies materialization unit if stub creation fails.
   auto &StubsMgr = getStubsManager(TargetVSO);
@@ -351,22 +334,8 @@ IndirectStubsManager &CompileOnDemandLayer2::getStubsManager(const VSO &V) {
 }
 
 void CompileOnDemandLayer2::emitExtractedFunctionsModule(
-    MaterializationResponsibility R, std::unique_ptr<Module> M,
-    std::shared_ptr<SymbolResolver> Resolver) {
-  auto &TargetVSO = R.getTargetVSO();
+    MaterializationResponsibility R, std::unique_ptr<Module> M) {
   auto K = getExecutionSession().allocateVModule();
-
-  auto ExtractedFunctionsResolver = createSymbolResolver(
-      [=](SymbolFlagsMap &Flags, const SymbolNameSet &Symbols) {
-        return Resolver->lookupFlags(Flags, Symbols);
-      },
-      [=, &TargetVSO](std::shared_ptr<AsynchronousSymbolQuery> Query,
-                      SymbolNameSet Symbols) {
-        auto RemainingSymbols = TargetVSO.lookup(Query, std::move(Symbols));
-        return Resolver->lookup(std::move(Query), std::move(RemainingSymbols));
-      });
-
-  SetSymbolResolver(K, std::move(ExtractedFunctionsResolver));
   BaseLayer.emit(std::move(R), std::move(K), std::move(M));
 }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 38f2213a85f..9275bb84d01 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -24,7 +24,6 @@ char FailedToMaterialize::ID = 0;
 char SymbolsNotFound::ID = 0;
 
 void MaterializationUnit::anchor() {}
-void SymbolResolver::anchor() {}
 
 raw_ostream &operator<<(raw_ostream &OS, const JITSymbolFlags &Flags) {
   if (Flags.isWeak())
@@ -99,6 +98,20 @@ raw_ostream &operator<<(raw_ostream &OS, const SymbolDependenceMap &Deps) {
   return OS;
 }
 
+raw_ostream &operator<<(raw_ostream &OS, const VSOList &VSOs) {
+  OS << "[";
+  if (!VSOs.empty()) {
+    assert(VSOs.front() && "VSOList entries must not be null");
+    OS << " " << VSOs.front()->getName();
+    for (auto *V : make_range(std::next(VSOs.begin()), VSOs.end())) {
+      assert(V && "VSOList entries must not be null");
+      OS << ", " << V->getName();
+    }
+  }
+  OS << " ]";
+  return OS;
+}
+
 FailedToMaterialize::FailedToMaterialize(SymbolNameSet Symbols)
     : Symbols(std::move(Symbols)) {
   assert(!this->Symbols.empty() && "Can not fail to resolve an empty set");
@@ -859,6 +872,34 @@ void VSO::runOutstandingMUs() {
   }
 }
 
+void VSO::setSearchOrder(VSOList NewSearchOrder, bool SearchThisVSOFirst) {
+  if (SearchThisVSOFirst && NewSearchOrder.front() != this)
+    NewSearchOrder.insert(NewSearchOrder.begin(), this);
+
+  ES.runSessionLocked([&]() { SearchOrder = std::move(NewSearchOrder); });
+}
+
+void VSO::addToSearchOrder(VSO &V) {
+  ES.runSessionLocked([&]() { SearchOrder.push_back(&V); });
+}
+
+void VSO::replaceInSearchOrder(VSO &OldV, VSO &NewV) {
+  ES.runSessionLocked([&]() {
+    auto I = std::find(SearchOrder.begin(), SearchOrder.end(), &OldV);
+
+    if (I != SearchOrder.end())
+      *I = &NewV;
+  });
+}
+
+void VSO::removeFromSearchOrder(VSO &V) {
+  ES.runSessionLocked([&]() {
+    auto I = std::find(SearchOrder.begin(), SearchOrder.end(), &V);
+    if (I != SearchOrder.end())
+      SearchOrder.erase(I);
+  });
+}
+
 SymbolNameSet VSO::lookupFlags(SymbolFlagsMap &Flags,
                                const SymbolNameSet &Names) {
   return ES.runSessionLocked([&, this]() {
@@ -1066,6 +1107,11 @@ void VSO::dump(raw_ostream &OS) {
   });
 }
 
+VSO::VSO(ExecutionSessionBase &ES, std::string Name)
+    : ES(ES), VSOName(std::move(Name)) {
+  SearchOrder.push_back(this);
+}
+
 Error VSO::defineImpl(MaterializationUnit &MU) {
   SymbolNameSet Duplicates;
   SymbolNameSet MUDefsOverridden;
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 059fba23d57..52ff4efe56b 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -35,7 +35,6 @@ Error LLJIT::addIRModule(VSO &V, std::unique_ptr<Module> M) {
     return Err;
 
   auto K = ES->allocateVModule();
-  Resolvers[K] = createResolverFor(V);
   return CompileLayer.add(V, K, std::move(M));
 }
 
@@ -49,23 +48,13 @@ LLJIT::LLJIT(std::unique_ptr<ExecutionSession> ES,
     : ES(std::move(ES)), Main(this->ES->createVSO("main")), TM(std::move(TM)),
       DL(std::move(DL)),
       ObjLinkingLayer(*this->ES,
-                      [this](VModuleKey K) { return getRTDyldResources(K); }),
+                      [this](VModuleKey K) { return getMemoryManager(K); }),
       CompileLayer(*this->ES, ObjLinkingLayer, SimpleCompiler(*this->TM)),
-      CtorRunner(Main), DtorRunner(Main) {
-  VSOLookupOrder[&Main] = VSOList({&Main});
-}
-
-std::shared_ptr<SymbolResolver> LLJIT::takeSymbolResolver(VModuleKey K) {
-  auto ResolverI = Resolvers.find(K);
-  assert(ResolverI != Resolvers.end() && "Missing resolver");
-  auto Resolver = std::move(ResolverI->second);
-  Resolvers.erase(ResolverI);
-  return Resolver;
-}
+      CtorRunner(Main), DtorRunner(Main) {}
 
-RTDyldObjectLinkingLayer2::Resources LLJIT::getRTDyldResources(VModuleKey K) {
-  return orc::RTDyldObjectLinkingLayer2::Resources(
-      {llvm::make_unique<SectionMemoryManager>(), takeSymbolResolver(K)});
+std::shared_ptr<RuntimeDyld::MemoryManager>
+LLJIT::getMemoryManager(VModuleKey K) {
+  return llvm::make_unique<SectionMemoryManager>();
 }
 
 std::string LLJIT::mangle(StringRef UnmangledName) {
@@ -77,21 +66,6 @@ std::string LLJIT::mangle(StringRef UnmangledName) {
   return MangledName;
 }
 
-std::unique_ptr<SymbolResolver> LLJIT::createResolverFor(VSO &V) {
-  return createSymbolResolver(
-      [&](SymbolFlagsMap &Flags, const SymbolNameSet &Symbols) {
-        return V.lookupFlags(Flags, Symbols);
-      },
-      [&, this](std::shared_ptr<AsynchronousSymbolQuery> Q,
-                SymbolNameSet Symbols) {
-        assert(VSOLookupOrder.count(&V) && "No VSO lookup order for V");
-        SymbolNameSet Unresolved = std::move(Symbols);
-        for (auto *LV : VSOLookupOrder[&V])
-          Unresolved = LV->lookup(Q, std::move(Unresolved));
-        return Unresolved;
-      });
-}
-
 Error LLJIT::applyDataLayout(Module &M) {
   if (M.getDataLayout().isDefault())
     M.setDataLayout(DL);
@@ -143,7 +117,6 @@ Error LLLazyJIT::addLazyIRModule(VSO &V, std::unique_ptr<Module> M) {
   recordCtorDtors(*M);
 
   auto K = ES->allocateVModule();
-  setSymbolResolver(K, createResolverFor(V));
   return CODLayer.add(V, K, std::move(M));
 }
 
@@ -155,17 +128,7 @@ LLLazyJIT::LLLazyJIT(
     : LLJIT(std::move(ES), std::move(TM), std::move(DL)),
       CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
       CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder),
-               [this](VModuleKey K) { return takeSymbolResolver(K); },
-               [this](VModuleKey K, std::shared_ptr<SymbolResolver> R) {
-                 setSymbolResolver(K, std::move(R));
-               },
                [&]() -> LLVMContext & { return Ctx; }) {}
 
-void LLLazyJIT::setSymbolResolver(VModuleKey K,
-                                  std::shared_ptr<SymbolResolver> R) {
-  assert(!Resolvers.count(K) && "Resolver already present for VModule K");
-  Resolvers[K] = std::move(R);
-}
-
 } // End namespace orc.
 } // End namespace llvm.
diff --git a/llvm/lib/ExecutionEngine/Orc/Legacy.cpp b/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
index 795ddd582d8..79525baf92e 100644
--- a/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Legacy.cpp
@@ -12,6 +12,8 @@
 namespace llvm {
 namespace orc {
 
+void SymbolResolver::anchor() {}
+
 JITSymbolResolverAdapter::JITSymbolResolverAdapter(
     ExecutionSession &ES, SymbolResolver &R, MaterializationResponsibility *MR)
     : ES(ES), R(R), MR(MR) {}
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 61e8ee8ff1a..8c53b4f58de 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -9,13 +9,85 @@
 
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 
+namespace {
+
+using namespace llvm;
+using namespace llvm::orc;
+
+class VSOSearchOrderResolver : public JITSymbolResolver {
+public:
+  VSOSearchOrderResolver(ExecutionSession &ES,
+                         MaterializationResponsibility &MR)
+      : ES(ES), MR(MR) {}
+
+  Expected<LookupResult> lookup(const LookupSet &Symbols) {
+    SymbolNameSet InternedSymbols;
+
+    for (auto &S : Symbols)
+      InternedSymbols.insert(ES.getSymbolStringPool().intern(S));
+
+    auto AsyncLookup = [&](std::shared_ptr<AsynchronousSymbolQuery> Q,
+                           SymbolNameSet Names) {
+      SymbolNameSet Unresolved = std::move(Names);
+      MR.getTargetVSO().withSearchOrderDo([&](const VSOList &SearchOrder) {
+        for (auto *V : SearchOrder) {
+          assert(V && "VSOList entry can not be null");
+          Unresolved = V->lookup(Q, std::move(Unresolved));
+        }
+      });
+      return Unresolved;
+    };
+
+    auto InternedResult = blockingLookup(
+        ES, std::move(AsyncLookup), std::move(InternedSymbols), false, &MR);
+
+    if (!InternedResult)
+      return InternedResult.takeError();
+
+    LookupResult Result;
+    for (auto &KV : *InternedResult)
+      Result[*KV.first] = std::move(KV.second);
+
+    return Result;
+  }
+
+  Expected<LookupFlagsResult> lookupFlags(const LookupSet &Symbols) {
+    SymbolNameSet InternedSymbols;
+
+    for (auto &S : Symbols)
+      InternedSymbols.insert(ES.getSymbolStringPool().intern(S));
+
+    SymbolFlagsMap InternedResult;
+    MR.getTargetVSO().withSearchOrderDo([&](const VSOList &VSOs) {
+      // An empty search order is pathalogical, but allowed.
+      if (VSOs.empty())
+        return;
+
+      assert(VSOs.front() && "VSOList entry can not be null");
+      VSOs.front()->lookupFlags(InternedResult, InternedSymbols);
+    });
+
+    LookupFlagsResult Result;
+    for (auto &KV : InternedResult)
+      Result[*KV.first] = std::move(KV.second);
+
+    return Result;
+  }
+
+private:
+  ExecutionSession &ES;
+  MaterializationResponsibility &MR;
+};
+
+} // end anonymous namespace
+
 namespace llvm {
 namespace orc {
 
 RTDyldObjectLinkingLayer2::RTDyldObjectLinkingLayer2(
-    ExecutionSession &ES, ResourcesGetterFunction GetResources,
+    ExecutionSession &ES, GetMemoryManagerFunction GetMemoryManager,
     NotifyLoadedFunction NotifyLoaded, NotifyFinalizedFunction NotifyFinalized)
-    : ObjectLayer(ES), GetResources(std::move(GetResources)),
+    : ObjectLayer(ES), GetMemoryManager(GetMemoryManager),
       NotifyLoaded(std::move(NotifyLoaded)),
       NotifyFinalized(std::move(NotifyFinalized)), ProcessAllSections(false) {}
 
@@ -32,11 +104,10 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
     R.failMaterialization();
   }
 
-  auto Resources = GetResources(K);
+  auto MemoryManager = GetMemoryManager(K);
 
-  JITSymbolResolverAdapter ResolverAdapter(ES, *Resources.Resolver, &R);
-  auto RTDyld =
-      llvm::make_unique<RuntimeDyld>(*Resources.MemMgr, ResolverAdapter);
+  VSOSearchOrderResolver Resolver(ES, R);
+  auto RTDyld = llvm::make_unique<RuntimeDyld>(*MemoryManager, Resolver);
   RTDyld->setProcessAllSections(ProcessAllSections);
 
   {
@@ -48,7 +119,7 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
 
     assert(!MemMgrs.count(K) &&
            "A memory manager already exists for this key?");
-    MemMgrs[K] = Resources.MemMgr;
+    MemMgrs[K] = std::move(MemoryManager);
   }
 
   auto Info = RTDyld->loadObject(**ObjFile);