1 files changed, 75 insertions, 17 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
index 166d1369c72..cc5a8a5c1ce 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
+++ b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
@@ -138,18 +138,67 @@ class OrcMCJITReplacement : public ExecutionEngine {
     std::shared_ptr<MCJITMemoryManager> ClientMM;
   };
 
-  class LinkingResolver : public LegacyJITSymbolResolver {
+  class LinkingORCResolver : public orc::SymbolResolver {
   public:
-    LinkingResolver(OrcMCJITReplacement &M) : M(M) {}
+    LinkingORCResolver(OrcMCJITReplacement &M) : M(M) {}
+
+    SymbolNameSet lookupFlags(SymbolFlagsMap &SymbolFlags,
+                              const SymbolNameSet &Symbols) override {
+      SymbolNameSet UnresolvedSymbols;
+
+      for (auto &S : Symbols) {
+        if (auto Sym = M.findMangledSymbol(*S)) {
+          SymbolFlags[S] = Sym.getFlags();
+        } else if (auto Err = Sym.takeError()) {
+          M.reportError(std::move(Err));
+          return SymbolNameSet();
+        } else {
+          if (auto Sym2 = M.ClientResolver->findSymbolInLogicalDylib(*S)) {
+            SymbolFlags[S] = Sym2.getFlags();
+          } else if (auto Err = Sym2.takeError()) {
+            M.reportError(std::move(Err));
+            return SymbolNameSet();
+          } else
+            UnresolvedSymbols.insert(S);
+        }
+      }
 
-    JITSymbol findSymbol(const std::string &Name) override {
-      return M.ClientResolver->findSymbol(Name);
+      return UnresolvedSymbols;
     }
 
-    JITSymbol findSymbolInLogicalDylib(const std::string &Name) override {
-      if (auto Sym = M.findMangledSymbol(Name))
-        return Sym;
-      return M.ClientResolver->findSymbolInLogicalDylib(Name);
+    SymbolNameSet lookup(AsynchronousSymbolQuery &Query,
+                         SymbolNameSet Symbols) override {
+      SymbolNameSet UnresolvedSymbols;
+
+      for (auto &S : Symbols) {
+        if (auto Sym = M.findMangledSymbol(*S)) {
+          if (auto Addr = Sym.getAddress())
+            Query.setDefinition(S, JITEvaluatedSymbol(*Addr, Sym.getFlags()));
+          else {
+            Query.setFailed(Addr.takeError());
+            return SymbolNameSet();
+          }
+        } else if (auto Err = Sym.takeError()) {
+          Query.setFailed(std::move(Err));
+          return SymbolNameSet();
+        } else {
+          if (auto Sym2 = M.ClientResolver->findSymbol(*S)) {
+            if (auto Addr = Sym2.getAddress())
+              Query.setDefinition(S,
+                                  JITEvaluatedSymbol(*Addr, Sym2.getFlags()));
+            else {
+              Query.setFailed(Addr.takeError());
+              return SymbolNameSet();
+            }
+          } else if (auto Err = Sym2.takeError()) {
+            Query.setFailed(std::move(Err));
+            return SymbolNameSet();
+          } else
+            UnresolvedSymbols.insert(S);
+        }
+      }
+
+      return UnresolvedSymbols;
     }
 
   private:
@@ -166,18 +215,23 @@ private:
                                    std::move(TM));
   }
 
+  void reportError(Error Err) {
+    logAllUnhandledErrors(std::move(Err), errs(), "MCJIT error: ");
+  }
+
 public:
   OrcMCJITReplacement(std::shared_ptr<MCJITMemoryManager> MemMgr,
                       std::shared_ptr<LegacyJITSymbolResolver> ClientResolver,
                       std::unique_ptr<TargetMachine> TM)
-      : ExecutionEngine(TM->createDataLayout()), TM(std::move(TM)),
+      : ExecutionEngine(TM->createDataLayout()), ES(SSP), TM(std::move(TM)),
         MemMgr(
             std::make_shared<MCJITReplacementMemMgr>(*this, std::move(MemMgr))),
-        Resolver(std::make_shared<LinkingResolver>(*this)),
+        Resolver(std::make_shared<LinkingORCResolver>(*this)),
         ClientResolver(std::move(ClientResolver)), NotifyObjectLoaded(*this),
         NotifyFinalized(*this),
-        ObjectLayer([this]() { return this->MemMgr; }, NotifyObjectLoaded,
-                    NotifyFinalized),
+        ObjectLayer(ES, [this](VModuleKey K) { return this->MemMgr; },
+                    [this](VModuleKey K) { return this->Resolver; },
+                    NotifyObjectLoaded, NotifyFinalized),
         CompileLayer(ObjectLayer, SimpleCompiler(*this->TM)),
         LazyEmitLayer(CompileLayer) {}
 
@@ -201,20 +255,21 @@ public:
         delete Mod;
     };
     LocalModules.push_back(std::shared_ptr<Module>(MPtr, std::move(Deleter)));
-    cantFail(LazyEmitLayer.addModule(LocalModules.back(), Resolver));
+    cantFail(
+        LazyEmitLayer.addModule(ES.allocateVModule(), LocalModules.back()));
   }
 
   void addObjectFile(std::unique_ptr<object::ObjectFile> O) override {
     auto Obj =
       std::make_shared<object::OwningBinary<object::ObjectFile>>(std::move(O),
                                                                  nullptr);
-    cantFail(ObjectLayer.addObject(std::move(Obj), Resolver));
+    cantFail(ObjectLayer.addObject(ES.allocateVModule(), std::move(Obj)));
   }
 
   void addObjectFile(object::OwningBinary<object::ObjectFile> O) override {
     auto Obj =
       std::make_shared<object::OwningBinary<object::ObjectFile>>(std::move(O));
-    cantFail(ObjectLayer.addObject(std::move(Obj), Resolver));
+    cantFail(ObjectLayer.addObject(ES.allocateVModule(), std::move(Obj)));
   }
 
   void addArchive(object::OwningBinary<object::Archive> A) override {
@@ -322,7 +377,7 @@ private:
           auto Obj =
             std::make_shared<object::OwningBinary<object::ObjectFile>>(
               std::move(ChildObj), nullptr);
-          cantFail(ObjectLayer.addObject(std::move(Obj), Resolver));
+          cantFail(ObjectLayer.addObject(ES.allocateVModule(), std::move(Obj)));
           if (auto Sym = ObjectLayer.findSymbol(Name, true))
             return Sym;
         }
@@ -374,9 +429,12 @@ private:
   using CompileLayerT = IRCompileLayer<ObjectLayerT, orc::SimpleCompiler>;
   using LazyEmitLayerT = LazyEmittingLayer<CompileLayerT>;
 
+  SymbolStringPool SSP;
+  ExecutionSession ES;
+
   std::unique_ptr<TargetMachine> TM;
   std::shared_ptr<MCJITReplacementMemMgr> MemMgr;
-  std::shared_ptr<LinkingResolver> Resolver;
+  std::shared_ptr<LinkingORCResolver> Resolver;
   std::shared_ptr<LegacyJITSymbolResolver> ClientResolver;
   Mangler Mang;