[ORC] Add cloneToContext: Clone Module to a given ThreadSafeContext (#146852)

This is a generalization of the existing cloneToNewContext operation:
rather than cloning the given module into a new ThreadSafeContext it
clones it into any given ThreadSafeContext. The given ThreadSafeContext
is locked to ensure that the cloning operation is safe.

1 files changed, 46 insertions, 34 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
index fadd53e..19c000e 100644
--- a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
@@ -14,51 +14,63 @@
 namespace llvm {
 namespace orc {
 
-ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,
-                                   GVPredicate ShouldCloneDef,
-                                   GVModifier UpdateClonedDefSource) {
+ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
+                                ThreadSafeContext TSCtx,
+                                GVPredicate ShouldCloneDef,
+                                GVModifier UpdateClonedDefSource) {
   assert(TSM && "Can not clone null module");
 
   if (!ShouldCloneDef)
     ShouldCloneDef = [](const GlobalValue &) { return true; };
 
-  return TSM.withModuleDo([&](Module &M) {
-    SmallVector<char, 1> ClonedModuleBuffer;
+  // First copy the source module into a buffer.
+  std::string ModuleName;
+  SmallVector<char, 1> ClonedModuleBuffer;
+  TSM.withModuleDo([&](Module &M) {
+    ModuleName = M.getModuleIdentifier();
+    std::set<GlobalValue *> ClonedDefsInSrc;
+    ValueToValueMapTy VMap;
+    auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
+      if (ShouldCloneDef(*GV)) {
+        ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
+        return true;
+      }
+      return false;
+    });
 
-    {
-      std::set<GlobalValue *> ClonedDefsInSrc;
-      ValueToValueMapTy VMap;
-      auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
-        if (ShouldCloneDef(*GV)) {
-          ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
-          return true;
-        }
-        return false;
-      });
+    if (UpdateClonedDefSource)
+      for (auto *GV : ClonedDefsInSrc)
+        UpdateClonedDefSource(*GV);
 
-      if (UpdateClonedDefSource)
-        for (auto *GV : ClonedDefsInSrc)
-          UpdateClonedDefSource(*GV);
+    BitcodeWriter BCWriter(ClonedModuleBuffer);
+    BCWriter.writeModule(*Tmp);
+    BCWriter.writeSymtab();
+    BCWriter.writeStrtab();
+  });
+
+  MemoryBufferRef ClonedModuleBufferRef(
+      StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
+      "cloned module buffer");
 
-      BitcodeWriter BCWriter(ClonedModuleBuffer);
+  // Then parse the buffer into the new Module.
+  auto M = TSCtx.withContextDo([&](LLVMContext *Ctx) {
+    assert(Ctx && "No LLVMContext provided");
+    auto TmpM = cantFail(parseBitcodeFile(ClonedModuleBufferRef, *Ctx));
+    TmpM->setModuleIdentifier(ModuleName);
+    return TmpM;
+  });
 
-      BCWriter.writeModule(*Tmp);
-      BCWriter.writeSymtab();
-      BCWriter.writeStrtab();
-    }
+  return ThreadSafeModule(std::move(M), std::move(TSCtx));
+}
 
-    MemoryBufferRef ClonedModuleBufferRef(
-        StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
-        "cloned module buffer");
-    ThreadSafeContext NewTSCtx(std::make_unique<LLVMContext>());
+ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,
+                                   GVPredicate ShouldCloneDef,
+                                   GVModifier UpdateClonedDefSource) {
+  assert(TSM && "Can not clone null module");
 
-    auto ClonedModule = NewTSCtx.withContextDo([&](LLVMContext *Ctx) {
-      auto TmpM = cantFail(parseBitcodeFile(ClonedModuleBufferRef, *Ctx));
-      TmpM->setModuleIdentifier(M.getName());
-      return TmpM;
-    });
-    return ThreadSafeModule(std::move(ClonedModule), std::move(NewTSCtx));
-  });
+  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  return cloneToContext(TSM, std::move(TSCtx), std::move(ShouldCloneDef),
+                        std::move(UpdateClonedDefSource));
 }
 
 } // end namespace orc