1 files changed, 88 insertions, 56 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index d87078f2e94..5d814bd3b9d 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -88,36 +88,30 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
                                      std::unique_ptr<MemoryBuffer> O) {
   assert(O && "Object must not be null");
 
-  auto &ES = getExecutionSession();
-
-  auto ObjFile = object::ObjectFile::createObjectFile(*O);
-  if (!ObjFile) {
-    getExecutionSession().reportError(ObjFile.takeError());
-    R.failMaterialization();
-  }
+  // This method launches an asynchronous link step that will fulfill our
+  // materialization responsibility. We need to switch R to be heap
+  // allocated before that happens so it can live as long as the asynchronous
+  // link needs it to (i.e. it must be able to outlive this method).
+  auto SharedR = std::make_shared<MaterializationResponsibility>(std::move(R));
 
-  auto MemoryManager = GetMemoryManager(K);
-
-  JITDylibSearchOrderResolver Resolver(R);
-  auto RTDyld = llvm::make_unique<RuntimeDyld>(*MemoryManager, Resolver);
-  RTDyld->setProcessAllSections(ProcessAllSections);
+  auto &ES = getExecutionSession();
 
-  {
-    std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  auto Obj = object::ObjectFile::createObjectFile(*O);
 
-    assert(!MemMgrs.count(K) &&
-           "A memory manager already exists for this key?");
-    MemMgrs[K] = std::move(MemoryManager);
+  if (!Obj) {
+    getExecutionSession().reportError(Obj.takeError());
+    SharedR->failMaterialization();
+    return;
   }
 
-  auto Info = RTDyld->loadObject(**ObjFile);
-
+  // Collect the internal symbols from the object file: We will need to
+  // filter these later.
+  auto InternalSymbols = std::make_shared<std::set<StringRef>>();
   {
-    std::set<StringRef> InternalSymbols;
-    for (auto &Sym : (*ObjFile)->symbols()) {
+    for (auto &Sym : (*Obj)->symbols()) {
       if (!(Sym.getFlags() & object::BasicSymbolRef::SF_Global)) {
         if (auto SymName = Sym.getName())
-          InternalSymbols.insert(*SymName);
+          InternalSymbols->insert(*SymName);
         else {
           ES.reportError(SymName.takeError());
           R.failMaterialization();
@@ -125,51 +119,89 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
         }
       }
     }
+  }
 
-    SymbolFlagsMap ExtraSymbolsToClaim;
-    SymbolMap Symbols;
-    for (auto &KV : RTDyld->getSymbolTable()) {
-      // Scan the symbols and add them to the Symbols map for resolution.
-
-      // We never claim internal symbols.
-      if (InternalSymbols.count(KV.first))
-        continue;
-
-      auto InternedName = ES.getSymbolStringPool().intern(KV.first);
-      auto Flags = KV.second.getFlags();
+  auto MemoryManager = GetMemoryManager(K);
+  auto &MemMgr = *MemoryManager;
+  {
+    std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
 
-      // Override object flags and claim responsibility for symbols if
-      // requested.
-      if (OverrideObjectFlags || AutoClaimObjectSymbols) {
-        auto I = R.getSymbols().find(InternedName);
+    assert(!MemMgrs.count(K) &&
+           "A memory manager already exists for this key?");
+    MemMgrs[K] = std::move(MemoryManager);
+  }
 
-        if (OverrideObjectFlags && I != R.getSymbols().end())
-          Flags = JITSymbolFlags::stripTransientFlags(I->second);
-        else if (AutoClaimObjectSymbols && I == R.getSymbols().end())
-          ExtraSymbolsToClaim[InternedName] = Flags;
-      }
+  JITDylibSearchOrderResolver Resolver(*SharedR);
+
+  /* Thoughts on proper cross-dylib weak symbol handling:
+   *
+   * Change selection of canonical defs to be a manually triggered process, and
+   * add a 'canonical' bit to symbol definitions. When canonical def selection
+   * is triggered, sweep the JITDylibs to mark defs as canonical, discard
+   * duplicate defs.
+   */
+  jitLinkForORC(
+      **Obj, std::move(O), MemMgr, Resolver, ProcessAllSections,
+      [this, K, SharedR, &Obj, InternalSymbols](
+          std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
+          std::map<StringRef, JITEvaluatedSymbol> ResolvedSymbols) {
+        return onObjLoad(K, *SharedR, **Obj, std::move(LoadedObjInfo),
+                         ResolvedSymbols, *InternalSymbols);
+      },
+      [this, K, SharedR](Error Err) {
+        onObjEmit(K, *SharedR, std::move(Err));
+      });
+}
 
-      Symbols[InternedName] = JITEvaluatedSymbol(KV.second.getAddress(), Flags);
+Error RTDyldObjectLinkingLayer2::onObjLoad(
+    VModuleKey K, MaterializationResponsibility &R, object::ObjectFile &Obj,
+    std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
+    std::map<StringRef, JITEvaluatedSymbol> Resolved,
+    std::set<StringRef> &InternalSymbols) {
+  SymbolFlagsMap ExtraSymbolsToClaim;
+  SymbolMap Symbols;
+  for (auto &KV : Resolved) {
+    // Scan the symbols and add them to the Symbols map for resolution.
+
+    // We never claim internal symbols.
+    if (InternalSymbols.count(KV.first))
+      continue;
+
+    auto InternedName =
+        getExecutionSession().getSymbolStringPool().intern(KV.first);
+    auto Flags = KV.second.getFlags();
+
+    // Override object flags and claim responsibility for symbols if
+    // requested.
+    if (OverrideObjectFlags || AutoClaimObjectSymbols) {
+      auto I = R.getSymbols().find(InternedName);
+
+      if (OverrideObjectFlags && I != R.getSymbols().end())
+        Flags = JITSymbolFlags::stripTransientFlags(I->second);
+      else if (AutoClaimObjectSymbols && I == R.getSymbols().end())
+        ExtraSymbolsToClaim[InternedName] = Flags;
     }
 
-    if (!ExtraSymbolsToClaim.empty())
-      if (auto Err = R.defineMaterializing(ExtraSymbolsToClaim)) {
-        ES.reportError(std::move(Err));
-        R.failMaterialization();
-        return;
-      }
-
-    R.resolve(Symbols);
+    Symbols[InternedName] = JITEvaluatedSymbol(KV.second.getAddress(), Flags);
   }
 
+  if (!ExtraSymbolsToClaim.empty())
+    if (auto Err = R.defineMaterializing(ExtraSymbolsToClaim))
+      return Err;
+
+  R.resolve(Symbols);
+
   if (NotifyLoaded)
-    NotifyLoaded(K, **ObjFile, *Info);
+    NotifyLoaded(K, Obj, *LoadedObjInfo);
 
-  RTDyld->finalizeWithMemoryManagerLocking();
+  return Error::success();
+}
 
-  if (RTDyld->hasError()) {
-    ES.reportError(make_error<StringError>(RTDyld->getErrorString(),
-                                           inconvertibleErrorCode()));
+void RTDyldObjectLinkingLayer2::onObjEmit(VModuleKey K,
+                                          MaterializationResponsibility &R,
+                                          Error Err) {
+  if (Err) {
+    getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
     return;
   }