[ORC] Add ThreadSafeModule::takingModuleDo for consuming operations on TSMs.

ThreadSafeModule::takingModuleDo passes ownership of the contained llvm::Module
into the callback, allowing ThreadSafeModules to be used with consuming
operations like Linker::link.

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
index c0df69f..bc34642 100644 (file)
--- a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
@@ -139,10 +139,19 @@ public:
   /// Locks the associated ThreadSafeContext and calls the given function
   /// on the contained Module.
   template <typename Func> decltype(auto) withModuleDo(Func &&F) const {
+    assert(M && "Can not call on null module");
     auto Lock = TSCtx.getLock();
     return F(*M);
   }
 
+  /// Locks the associated ThreadSafeContext and calls the given function,
+  /// passing the contained std::unique_ptr<Module>. The given function should
+  /// consume the Module.
+  template <typename Func> decltype(auto) takingModuleDo(Func &&F) {
+    auto Lock = TSCtx.getLock();
+    return F(std::move(M));
+  }
+
   /// Get a raw pointer to the contained module without locking the context.
   Module *getModuleUnlocked() { return M.get(); }
 

diff --git a/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp
index 1ffb06d..1cda029 100644 (file)
--- a/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp
@@ -91,4 +91,28 @@ TEST(ThreadSafeModuleTest, ContextLockPreservesContext) {
   TSCtx = ThreadSafeContext();
 }
 
+TEST(ThreadSafeModuleTest, WithModuleDo) {
+  // Test non-const version of withModuleDo.
+  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  ThreadSafeModule TSM(std::make_unique<Module>("M", *TSCtx.getContext()),
+                       TSCtx);
+  TSM.withModuleDo([](Module &M) {});
+}
+
+TEST(ThreadSafeModuleTest, WithModuleDoConst) {
+  // Test const version of withModuleDo.
+  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  const ThreadSafeModule TSM(std::make_unique<Module>("M", *TSCtx.getContext()),
+                             TSCtx);
+  TSM.withModuleDo([](const Module &M) {});
+}
+
+TEST(ThreadSafeModuleTest, TakingModuleDo) {
+  // Test takingModuleDo.
+  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  ThreadSafeModule TSM(std::make_unique<Module>("M", *TSCtx.getContext()),
+                       TSCtx);
+  TSM.takingModuleDo([](std::unique_ptr<Module> M) {});
+}
+
 } // end anonymous namespace