[Orc] Make partition identification in the CompileOnDemand layer lazy.

This also breaks out the logical dylib symbol resolution logic.



git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@239561 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -15,9 +15,9 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
 
-//#include "CloneSubModule.h"
 #include "IndirectionUtils.h"
 #include "LambdaResolver.h"
+#include "LogicalDylib.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/Transforms/Utils/Cloning.h"
@@ -36,7 +36,9 @@ namespace orc {
 /// added to the layer below. When a stub is called it triggers the extraction
 /// of the function body from the original module. The extracted body is then
 /// compiled and executed.
-template <typename BaseLayerT, typename CompileCallbackMgrT>
+template <typename BaseLayerT, typename CompileCallbackMgrT,
+          typename PartitioningFtor =
+            std::function<std::set<Function*>(Function&)>>
 class CompileOnDemandLayer {
 private:
 
@@ -58,311 +60,23 @@ private:
   };
 
   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
-  class UncompiledPartition;
 
-  // Logical module.
+  struct LogicalModuleResources {
-  //
+    std::shared_ptr<Module> SourceModule;
-  //   This struct contains the handles for the global values and stubs (which
-  // cover the external symbols of the original module), plus the handes for
-  // each of the extracted partitions. These handleds are used for lookup (only
-  // the globals/stubs module is searched) and memory management. The actual
-  // searching and resource management are handled by the LogicalDylib that owns
-  // the LogicalModule.
-  struct LogicalModule {
-    LogicalModule() {}
-
-    LogicalModule(LogicalModule &&Other)
-        : SrcM(std::move(Other.SrcM)),
-          GVsAndStubsHandle(std::move(Other.GVsAndStubsHandle)),
-          ImplHandles(std::move(Other.ImplHandles)) {}
-
-    std::unique_ptr<Module> SrcM;
-    BaseLayerModuleSetHandleT GVsAndStubsHandle;
-    std::vector<BaseLayerModuleSetHandleT> ImplHandles;
   };
 
-  // Logical dylib.
+  struct LogicalDylibResources {
-  //
+    typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
-  //   This class handles symbol resolution and resource management for a set of
+      SymbolResolverFtor;
-  // modules that were added together as a logical dylib.
+    SymbolResolverFtor ExternalSymbolResolver;
-  //
+    PartitioningFtor Partitioner;
-  //   A logical dylib contains one-or-more LogicalModules plus a set of
-  // UncompiledPartitions. LogicalModules support symbol resolution and resource
-  // management for for code that has already been emitted. UncompiledPartitions
-  // represent code that has not yet been compiled.
-  class LogicalDylib {
-  private:
-    friend class UncompiledPartition;
-    typedef std::list<LogicalModule> LogicalModuleList;
-  public:
-
-    typedef unsigned UncompiledPartitionID;
-    typedef typename LogicalModuleList::iterator LMHandle;
-
-    // Construct a logical dylib.
-    LogicalDylib(CompileOnDemandLayer &CODLayer) : CODLayer(CODLayer) { }
-
-    // Delete this logical dylib, release logical module resources.
-    virtual ~LogicalDylib() {
-      releaseLogicalModuleResources();
-    }
-
-    // Get a reference to the containing layer.
-    CompileOnDemandLayer& getCODLayer() { return CODLayer; }
-
-    // Get a reference to the base layer.
-    BaseLayerT& getBaseLayer() { return CODLayer.BaseLayer; }
-
-    // Start a new context for a single logical module.
-    LMHandle createLogicalModule() {
-      LogicalModules.push_back(LogicalModule());
-      return std::prev(LogicalModules.end());
-    }
-
-    // Set the global-values-and-stubs module handle for this logical module.
-    void setGVsAndStubsHandle(LMHandle LMH, BaseLayerModuleSetHandleT H) {
-      LMH->GVsAndStubsHandle = H;
-    }
-
-    // Return the global-values-and-stubs module handle for this logical module.
-    BaseLayerModuleSetHandleT getGVsAndStubsHandle(LMHandle LMH) {
-      return LMH->GVsAndStubsHandle;
-    }
-
-    //   Add a handle to a module containing lazy function bodies to the given
-    // logical module.
-    void addToLogicalModule(LMHandle LMH, BaseLayerModuleSetHandleT H) {
-      LMH->ImplHandles.push_back(H);
-    }
-
-    // Create an UncompiledPartition attached to this LogicalDylib.
-    UncompiledPartition& createUncompiledPartition(LMHandle LMH,
-                                                   std::shared_ptr<Module> SrcM);
-
-    // Take ownership of the given UncompiledPartition from the logical dylib.
-    std::unique_ptr<UncompiledPartition>
-    takeUPOwnership(UncompiledPartitionID ID);
-
-    // Look up a symbol in this context.
-    JITSymbol findSymbolInternally(LMHandle LMH, const std::string &Name) {
-      if (auto Symbol = getBaseLayer().findSymbolIn(LMH->GVsAndStubsHandle,
-                                                    Name, false))
-        return Symbol;
-
-      for (auto I = LogicalModules.begin(), E = LogicalModules.end(); I != E;
-           ++I)
-        if (I != LMH)
-          if (auto Symbol = getBaseLayer().findSymbolIn(I->GVsAndStubsHandle,
-                                                        Name, false))
-            return Symbol;
-
-      return nullptr;
-    }
-
-    JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
-      for (auto &LM : LogicalModules)
-        if (auto Symbol = getBaseLayer().findSymbolIn(LM.GVsAndStubsHandle,
-                                                      Name,
-                                                      ExportedSymbolsOnly))
-          return Symbol;
-      return nullptr;
-    }
-
-    // Find an external symbol (via the user supplied SymbolResolver).
-    virtual RuntimeDyld::SymbolInfo
-    findSymbolExternally(const std::string &Name) const = 0;
-
-  private:
-
-    void releaseLogicalModuleResources() {
-      for (auto I = LogicalModules.begin(), E = LogicalModules.end(); I != E;
-           ++I) {
-        getBaseLayer().removeModuleSet(I->GVsAndStubsHandle);
-        for (auto H : I->ImplHandles)
-          getBaseLayer().removeModuleSet(H);
-      }
-    }
-
-    CompileOnDemandLayer &CODLayer;
-    LogicalModuleList LogicalModules;
-    std::vector<std::unique_ptr<UncompiledPartition>> UncompiledPartitions;
   };
 
-  template <typename ResolverPtrT>
+  typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
-  class LogicalDylibImpl : public LogicalDylib  {
+                       LogicalDylibResources> CODLogicalDylib;
-  public:
-    LogicalDylibImpl(CompileOnDemandLayer &CODLayer, ResolverPtrT Resolver)
-      : LogicalDylib(CODLayer), Resolver(std::move(Resolver)) {}
 
-    RuntimeDyld::SymbolInfo
+  typedef typename CODLogicalDylib::LogicalModuleHandle LogicalModuleHandle;
-    findSymbolExternally(const std::string &Name) const override {
+  typedef std::list<CODLogicalDylib> LogicalDylibList;
-      return Resolver->findSymbol(Name);
-    }
-
-  private:
-    ResolverPtrT Resolver;
-  };
-
-  template <typename ResolverPtrT>
-  static std::unique_ptr<LogicalDylib>
-  createLogicalDylib(CompileOnDemandLayer &CODLayer,
-                     ResolverPtrT Resolver) {
-    typedef LogicalDylibImpl<ResolverPtrT> Impl;
-    return llvm::make_unique<Impl>(CODLayer, std::move(Resolver));
-  }
-
-  // Uncompiled partition.
-  //
-  // Represents one as-yet uncompiled portion of a module.
-  class UncompiledPartition {
-  public:
-
-    struct PartitionEntry {
-      PartitionEntry(Function *F, TargetAddress CallbackID)
-          : F(F), CallbackID(CallbackID) {}
-      Function *F;
-      TargetAddress CallbackID;
-    };
-
-    typedef std::vector<PartitionEntry> PartitionEntryList;
-
-    // Creates an uncompiled partition with the list of functions that make up
-    // this partition.
-    UncompiledPartition(LogicalDylib &LD, typename LogicalDylib::LMHandle LMH,
-                        std::shared_ptr<Module> SrcM)
-        : LD(LD), LMH(LMH), SrcM(std::move(SrcM)), ID(~0U) {}
-
-    ~UncompiledPartition() {
-      // FIXME: When we want to support threaded lazy compilation we'll need to
-      //        lock the callback manager here.
-      auto &CCMgr = LD.getCODLayer().CompileCallbackMgr;
-      for (auto PEntry : PartitionEntries)
-        CCMgr.releaseCompileCallback(PEntry.CallbackID);
-    }
-
-    // Set the ID for this partition.
-    void setID(typename LogicalDylib::UncompiledPartitionID ID) {
-      this->ID = ID;
-    }
-
-    // Set the function set and callbacks for this partition.
-    void setPartitionEntries(PartitionEntryList PartitionEntries) {
-      this->PartitionEntries = std::move(PartitionEntries);
-    }
-
-    // Handle a compile callback for the function at index FnIdx.
-    TargetAddress compile(unsigned FnIdx) {
-      // Take ownership of self. This will ensure we delete the partition and
-      // free all its resources once we're done compiling.
-      std::unique_ptr<UncompiledPartition> This = LD.takeUPOwnership(ID);
-
-      // Release all other compile callbacks for this partition.
-      // We skip the callback for this function because that's the one that
-      // called us, and the callback manager will already have removed it.
-      auto &CCMgr = LD.getCODLayer().CompileCallbackMgr;
-      for (unsigned I = 0; I < PartitionEntries.size(); ++I)
-        if (I != FnIdx)
-          CCMgr.releaseCompileCallback(PartitionEntries[I].CallbackID);
-
-      // Grab the name of the function being called here.
-      Function *F = PartitionEntries[FnIdx].F;
-      std::string CalledFnName = Mangle(F->getName(), SrcM->getDataLayout());
-
-      // Extract the function and add it to the base layer.
-      auto PartitionImplH = emitPartition();
-      LD.addToLogicalModule(LMH, PartitionImplH);
-
-      // Update body pointers.
-      // FIXME: When we start supporting remote lazy jitting this will need to
-      //        be replaced with a user-supplied callback for updating the
-      //        remote pointers.
-      TargetAddress CalledAddr = 0;
-      for (unsigned I = 0; I < PartitionEntries.size(); ++I) {
-        auto F = PartitionEntries[I].F;
-        std::string FName(F->getName());
-        auto FnBodySym =
-          LD.getBaseLayer().findSymbolIn(PartitionImplH,
-                                         Mangle(FName, SrcM->getDataLayout()),
-                                         false);
-        auto FnPtrSym =
-          LD.getBaseLayer().findSymbolIn(LD.getGVsAndStubsHandle(LMH),
-                                         Mangle(FName + "$orc_addr",
-                                                SrcM->getDataLayout()),
-                                         false);
-        assert(FnBodySym && "Couldn't find function body.");
-        assert(FnPtrSym && "Couldn't find function body pointer.");
-
-        auto FnBodyAddr = FnBodySym.getAddress();
-        void *FnPtrAddr = reinterpret_cast<void*>(
-                            static_cast<uintptr_t>(FnPtrSym.getAddress()));
-
-        // If this is the function we're calling record the address so we can
-        // return it from this function.
-        if (I == FnIdx)
-          CalledAddr = FnBodyAddr;
-
-        memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
-      }
-
-      // Finally, clear the partition structure so we don't try to
-      // double-release the callbacks in the UncompiledPartition destructor.
-      PartitionEntries.clear();
-
-      return CalledAddr;
-    }
-
-  private:
-
-    BaseLayerModuleSetHandleT emitPartition() {
-      // Create the module.
-      std::string NewName(SrcM->getName());
-      for (auto &PEntry : PartitionEntries) {
-        NewName += ".";
-        NewName += PEntry.F->getName();
-      }
-      auto PM = llvm::make_unique<Module>(NewName, SrcM->getContext());
-      PM->setDataLayout(SrcM->getDataLayout());
-      ValueToValueMapTy VMap;
-      GlobalDeclMaterializer GDM(*PM);
-
-      // Create decls in the new module.
-      for (auto &PEntry : PartitionEntries)
-        cloneFunctionDecl(*PM, *PEntry.F, &VMap);
-
-      // Move the function bodies.
-      for (auto &PEntry : PartitionEntries)
-        moveFunctionBody(*PEntry.F, VMap);
-
-      // Create memory manager and symbol resolver.
-      auto MemMgr = llvm::make_unique<SectionMemoryManager>();
-      auto Resolver = createLambdaResolver(
-          [this](const std::string &Name) {
-            if (auto Symbol = LD.findSymbolInternally(LMH, Name))
-              return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
-                                             Symbol.getFlags());
-            return LD.findSymbolExternally(Name);
-          },
-          [this](const std::string &Name) {
-            if (auto Symbol = LD.findSymbolInternally(LMH, Name))
-              return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
-                                             Symbol.getFlags());
-            return RuntimeDyld::SymbolInfo(nullptr);
-          });
-      std::vector<std::unique_ptr<Module>> PartMSet;
-      PartMSet.push_back(std::move(PM));
-      return LD.getBaseLayer().addModuleSet(std::move(PartMSet),
-                                            std::move(MemMgr),
-                                            std::move(Resolver));
-    }
-
-    LogicalDylib &LD;
-    typename LogicalDylib::LMHandle LMH;
-    std::shared_ptr<Module> SrcM;
-    typename LogicalDylib::UncompiledPartitionID ID;
-    PartitionEntryList PartitionEntries;
-  };
-
-  typedef std::list<std::unique_ptr<LogicalDylib>> LogicalDylibList;
 
 public:
   /// @brief Handle to a set of loaded modules.
@@ -382,20 +96,25 @@ public:
     assert(MemMgr == nullptr &&
            "User supplied memory managers not supported with COD yet.");
 
-    LogicalDylibs.push_back(createLogicalDylib(*this, std::move(Resolver)));
+    LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
+    auto &LDLResources = LogicalDylibs.back().getDylibResources();
+
+    LDLResources.ExternalSymbolResolver =
+      [Resolver](const std::string &Name) {
+        return Resolver->findSymbol(Name);
+      };
+
+    LDLResources.Partitioner =
+      [](Function &F) {
+        std::set<Function*> Partition;
+        Partition.insert(&F);
+        return Partition;
+      };
 
     // Process each of the modules in this module set.
-    for (auto &M : Ms) {
+    for (auto &M : Ms)
-      std::vector<std::vector<Function*>> Partitioning;
+      addLogicalModule(LogicalDylibs.back(),
-      for (auto &F : *M) {
+                       std::shared_ptr<Module>(std::move(M)));
-        if (F.isDeclaration())
-          continue;
-        Partitioning.emplace_back(1, &F);
-      }
-      addLogicalModule(*LogicalDylibs.back(),
-                       std::shared_ptr<Module>(std::move(M)),
-                       std::move(Partitioning));
-    }
 
     return std::prev(LogicalDylibs.end());
   }
@@ -420,13 +139,12 @@ public:
   ///        below this one.
   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
                          bool ExportedSymbolsOnly) {
-    return (*H)->findSymbol(Name, ExportedSymbolsOnly);
+    return H->findSymbol(Name, ExportedSymbolsOnly);
   }
 
 private:
 
-  void addLogicalModule(LogicalDylib &LD, std::shared_ptr<Module> SrcM,
+  void addLogicalModule(CODLogicalDylib &LD, std::shared_ptr<Module> SrcM) {
-                        std::vector<std::vector<Function*>> Partitions) {
 
     // Bump the linkage and rename any anonymous/privote members in SrcM to
     // ensure that everything will resolve properly after we partition SrcM.
@@ -434,6 +152,7 @@ private:
 
     // Create a logical module handle for SrcM within the logical dylib.
     auto LMH = LD.createLogicalModule();
+    LD.getLogicalModuleResources(LMH).SourceModule = SrcM;
 
     // Create the GVs-and-stubs module.
     auto GVsAndStubsM = llvm::make_unique<Module>(
@@ -442,31 +161,31 @@ private:
     GVsAndStubsM->setDataLayout(SrcM->getDataLayout());
     ValueToValueMapTy VMap;
 
-    // Process partitions and create stubs.
+    // Process module and create stubs.
     // We create the stubs before copying the global variables as we know the
     // stubs won't refer to any globals (they only refer to their implementation
     // pointer) so there's no ordering/value-mapping issues.
-    for (auto& Partition : Partitions) {
+    for (auto &F : *SrcM) {
-      auto &UP = LD.createUncompiledPartition(LMH, SrcM);
-      typename UncompiledPartition::PartitionEntryList PartitionEntries;
-      for (auto &F : Partition) {
-        assert(!F->isDeclaration() &&
-               "Partition should only contain definitions");
-        unsigned FnIdx = PartitionEntries.size();
-        auto CCI = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
-        PartitionEntries.push_back(
-          typename UncompiledPartition::PartitionEntry(F, CCI.getAddress()));
-        Function *StubF = cloneFunctionDecl(*GVsAndStubsM, *F, &VMap);
-        GlobalVariable *FnBodyPtr =
-          createImplPointer(*StubF->getType(), *StubF->getParent(),
-                            StubF->getName() + "$orc_addr",
-                            createIRTypedAddress(*StubF->getFunctionType(),
-                                                 CCI.getAddress()));
-        makeStub(*StubF, *FnBodyPtr);
-        CCI.setCompileAction([&UP, FnIdx]() { return UP.compile(FnIdx); });
-      }
 
-      UP.setPartitionEntries(std::move(PartitionEntries));
+      // Skip declarations.
+      if (F.isDeclaration())
+        continue;
+
+      // For each definition: create a callback, a stub, and a function body
+      // pointer. Initialize the function body pointer to point at the callback,
+      // and set the callback to compile the function body.
+      auto CCInfo = CompileCallbackMgr.getCompileCallback(SrcM->getContext());
+      Function *StubF = cloneFunctionDecl(*GVsAndStubsM, F, &VMap);
+      GlobalVariable *FnBodyPtr =
+        createImplPointer(*StubF->getType(), *StubF->getParent(),
+                          StubF->getName() + "$orc_addr",
+                          createIRTypedAddress(*StubF->getFunctionType(),
+                                               CCInfo.getAddress()));
+      makeStub(*StubF, *FnBodyPtr);
+      CCInfo.setCompileAction(
+        [this, &LD, LMH, &F]() {
+          return extractAndCompile(LD, LMH, F);
+        });
     }
 
     // Now clone the global variable declarations.
@@ -483,12 +202,9 @@ private:
     // Build a resolver for the stubs module and add it to the base layer.
     auto GVsAndStubsResolver = createLambdaResolver(
         [&LD](const std::string &Name) {
-          if (auto Symbol = LD.findSymbol(Name, false))
+          return LD.getDylibResources().ExternalSymbolResolver(Name);
-            return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
-                                           Symbol.getFlags());
-          return LD.findSymbolExternally(Name);
         },
-        [&LD](const std::string &Name) {
+        [](const std::string &Name) {
           return RuntimeDyld::SymbolInfo(nullptr);
         });
 
@@ -498,7 +214,7 @@ private:
       BaseLayer.addModuleSet(std::move(GVsAndStubsMSet),
                              llvm::make_unique<SectionMemoryManager>(),
                              std::move(GVsAndStubsResolver));
-    LD.setGVsAndStubsHandle(LMH, GVsAndStubsH);
+    LD.addToLogicalModule(LMH, GVsAndStubsH);
   }
 
   static std::string Mangle(StringRef Name, const DataLayout &DL) {
@@ -511,35 +227,101 @@ private:
     return MangledName;
   }
 
+  TargetAddress extractAndCompile(CODLogicalDylib &LD,
+                                  LogicalModuleHandle LMH,
+                                  Function &F) {
+    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+
+    // If F is a declaration we must already have compiled it.
+    if (F.isDeclaration())
+      return 0;
+
+    // Grab the name of the function being called here.
+    std::string CalledFnName = Mangle(F.getName(), SrcM.getDataLayout());
+
+    const auto &Partition = LD.getDylibResources().Partitioner(F);
+    auto PartitionH = emitPartition(LD, LMH, Partition);
+
+    TargetAddress CalledAddr = 0;
+    for (auto *SubF : Partition) {
+      std::string FName(SubF->getName());
+      auto FnBodySym =
+        BaseLayer.findSymbolIn(PartitionH, Mangle(FName, SrcM.getDataLayout()),
+                               false);
+      auto FnPtrSym =
+        BaseLayer.findSymbolIn(*LD.moduleHandlesBegin(LMH),
+                               Mangle(FName + "$orc_addr",
+                                      SrcM.getDataLayout()),
+                               false);
+      assert(FnBodySym && "Couldn't find function body.");
+      assert(FnPtrSym && "Couldn't find function body pointer.");
+
+      auto FnBodyAddr = FnBodySym.getAddress();
+      void *FnPtrAddr = reinterpret_cast<void*>(
+          static_cast<uintptr_t>(FnPtrSym.getAddress()));
+
+      // If this is the function we're calling record the address so we can
+      // return it from this function.
+      if (SubF == &F)
+        CalledAddr = FnBodyAddr;
+
+      memcpy(FnPtrAddr, &FnBodyAddr, sizeof(uintptr_t));
+    }
+
+    return CalledAddr;
+  }
+
+  BaseLayerModuleSetHandleT emitPartition(CODLogicalDylib &LD,
+                                          LogicalModuleHandle LMH,
+                                          const std::set<Function*> &Partition) {
+    Module &SrcM = *LD.getLogicalModuleResources(LMH).SourceModule;
+
+    // Create the module.
+    std::string NewName(SrcM.getName());
+    for (auto *F : Partition) {
+      NewName += ".";
+      NewName += F->getName();
+    }
+
+    auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
+    M->setDataLayout(SrcM.getDataLayout());
+    ValueToValueMapTy VMap;
+    GlobalDeclMaterializer GDM(*M);
+
+    // Create decls in the new module.
+    for (auto *F : Partition)
+      cloneFunctionDecl(*M, *F, &VMap);
+
+    // Move the function bodies.
+    for (auto *F : Partition)
+      moveFunctionBody(*F, VMap);
+
+    // Create memory manager and symbol resolver.
+    auto MemMgr = llvm::make_unique<SectionMemoryManager>();
+    auto Resolver = createLambdaResolver(
+        [this, &LD, LMH](const std::string &Name) {
+          if (auto Symbol = LD.findSymbolInternally(LMH, Name))
+            return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                           Symbol.getFlags());
+          return LD.getDylibResources().ExternalSymbolResolver(Name);
+        },
+        [this, &LD, LMH](const std::string &Name) {
+          if (auto Symbol = LD.findSymbolInternally(LMH, Name))
+            return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                           Symbol.getFlags());
+          return RuntimeDyld::SymbolInfo(nullptr);
+        });
+    std::vector<std::unique_ptr<Module>> PartMSet;
+    PartMSet.push_back(std::move(M));
+    return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
+                                  std::move(Resolver));
+  }
+
   BaseLayerT &BaseLayer;
   CompileCallbackMgrT &CompileCallbackMgr;
   LogicalDylibList LogicalDylibs;
 };
 
-template <typename BaseLayerT, typename CompileCallbackMgrT>
-typename CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::
-           UncompiledPartition&
-CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::LogicalDylib::
-  createUncompiledPartition(LMHandle LMH, std::shared_ptr<Module> SrcM) {
-  UncompiledPartitions.push_back(
-      llvm::make_unique<UncompiledPartition>(*this, LMH, std::move(SrcM)));
-  UncompiledPartitions.back()->setID(UncompiledPartitions.size() - 1);
-  return *UncompiledPartitions.back();
-}
-
-template <typename BaseLayerT, typename CompileCallbackMgrT>
-std::unique_ptr<typename CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::
-                  UncompiledPartition>
-CompileOnDemandLayer<BaseLayerT, CompileCallbackMgrT>::LogicalDylib::
-  takeUPOwnership(UncompiledPartitionID ID) {
-
-  std::swap(UncompiledPartitions[ID], UncompiledPartitions.back());
-  UncompiledPartitions[ID]->setID(ID);
-  auto UP = std::move(UncompiledPartitions.back());
-  UncompiledPartitions.pop_back();
-  return UP;
-}
-
 } // End namespace orc.
 } // End namespace llvm.
 

include/llvm/ExecutionEngine/Orc/LogicalDylib.h

@@ -0,0 +1,118 @@
+//===--- LogicalDylib.h - Simulates dylib-style symbol lookup ---*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Simulates symbol resolution inside a dylib.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LOGICALDYLIB_H
+#define LLVM_EXECUTIONENGINE_ORC_LOGICALDYLIB_H
+
+#include "llvm/ADT/iterator.h"
+#include "llvm/ADT/Optional.h"
+
+namespace llvm {
+namespace orc {
+
+template <typename BaseLayerT,
+          typename LogicalModuleResources,
+          typename LogicalDylibResources>
+class LogicalDylib {
+public:
+  typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
+private:
+
+  typedef std::vector<BaseLayerModuleSetHandleT> BaseLayerHandleList;
+
+  struct LogicalModule {
+    LogicalModuleResources Resources;
+    BaseLayerHandleList BaseLayerHandles;
+  };
+  typedef std::vector<LogicalModule> LogicalModuleList;
+
+public:
+
+  typedef typename BaseLayerHandleList::iterator BaseLayerHandleIterator;
+  typedef typename LogicalModuleList::iterator LogicalModuleHandle;
+
+  LogicalDylib(BaseLayerT &BaseLayer) : BaseLayer(BaseLayer) {}
+
+  ~LogicalDylib() {
+    for (auto &LM : LogicalModules)
+      for (auto BLH : LM.BaseLayerHandles)
+        BaseLayer.removeModuleSet(BLH);
+  }
+
+  LogicalModuleHandle createLogicalModule() {
+    LogicalModules.push_back(LogicalModule());
+    return std::prev(LogicalModules.end());
+  }
+
+  void addToLogicalModule(LogicalModuleHandle LMH,
+                          BaseLayerModuleSetHandleT BaseLayerHandle) {
+    LMH->BaseLayerHandles.push_back(BaseLayerHandle);
+  }
+
+  LogicalModuleResources& getLogicalModuleResources(LogicalModuleHandle LMH) {
+    return LMH->Resources;
+  }
+
+  BaseLayerHandleIterator moduleHandlesBegin(LogicalModuleHandle LMH) {
+    return LMH->BaseLayerHandles.begin();
+  }
+
+  BaseLayerHandleIterator moduleHandlesEnd(LogicalModuleHandle LMH) {
+    return LMH->BaseLayerHandles.end();
+  }
+
+  JITSymbol findSymbolInLogicalModule(LogicalModuleHandle LMH,
+                                      const std::string &Name) {
+    for (auto BLH : LMH->BaseLayerHandles)
+      if (auto Symbol = BaseLayer.findSymbolIn(BLH, Name, false))
+        return Symbol;
+    return nullptr;
+  }
+
+  JITSymbol findSymbolInternally(LogicalModuleHandle LMH,
+                                 const std::string &Name) {
+    if (auto Symbol = findSymbolInLogicalModule(LMH, Name))
+      return Symbol;
+
+    for (auto LMI = LogicalModules.begin(), LME = LogicalModules.end();
+           LMI != LME; ++LMI) {
+      if (LMI != LMH)
+        if (auto Symbol = findSymbolInLogicalModule(LMI, Name))
+          return Symbol;
+    }
+
+    return nullptr;
+  }
+
+  JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
+    for (auto &LM : LogicalModules)
+      for (auto BLH : LM.BaseLayerHandles)
+        if (auto Symbol =
+            BaseLayer.findSymbolIn(BLH, Name, ExportedSymbolsOnly))
+          return Symbol;
+    return nullptr;
+  }
+
+  LogicalDylibResources& getDylibResources() { return DylibResources; }
+
+protected:
+  BaseLayerT BaseLayer;
+  LogicalModuleList LogicalModules;
+  LogicalDylibResources DylibResources;
+
+};
+
+} // End namespace orc.
+} // End namespace llvm.
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LOGICALDYLIB_H

tools/lli/OrcLazyJIT.h

@@ -89,22 +89,26 @@ public:
     //   2) Check for C++ runtime overrides.
     //   3) Search the host process (LLI)'s symbol table.
     auto Resolver =
-      orc::createLambdaResolver(
+      std::shared_ptr<RuntimeDyld::SymbolResolver>(
-        [this](const std::string &Name) {
+        orc::createLambdaResolver(
+          [this](const std::string &Name) {              
+            if (auto Sym = CODLayer.findSymbol(Name, true))
+              return RuntimeDyld::SymbolInfo(Sym.getAddress(),
+                                             Sym.getFlags());
                 
-          if (auto Sym = CODLayer.findSymbol(Name, true))
+            if (auto Sym = CXXRuntimeOverrides.searchOverrides(Name))
-            return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
+              return Sym;
 
-          if (auto Sym = CXXRuntimeOverrides.searchOverrides(Name))
+            if (auto Addr =
-            return Sym;
+                RTDyldMemoryManager::getSymbolAddressInProcess(Name))
+              return RuntimeDyld::SymbolInfo(Addr, JITSymbolFlags::Exported);
 
-          if (auto Addr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
+            return RuntimeDyld::SymbolInfo(nullptr);
-            return RuntimeDyld::SymbolInfo(Addr, JITSymbolFlags::Exported);
+          },
-
+          [](const std::string &Name) {
-          return RuntimeDyld::SymbolInfo(nullptr);
+            return RuntimeDyld::SymbolInfo(nullptr);
-        },
+          }
-        [](const std::string &Name) { return RuntimeDyld::SymbolInfo(nullptr); }
+        ));
-      );
 
     // Add the module to the JIT.
     std::vector<std::unique_ptr<Module>> S;