[llvm] 0bf85d7 - [ORC] Add ThreadSafeModule::takingModuleDo for consuming operations on TSMs.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2023-01-02T17:38:16-08:00
New Revision: 0bf85d7d8129339ac66ac8d94a7a4d636508fe1a

URL: https://github.com/llvm/llvm-project/commit/0bf85d7d8129339ac66ac8d94a7a4d636508fe1a
DIFF: https://github.com/llvm/llvm-project/commit/0bf85d7d8129339ac66ac8d94a7a4d636508fe1a.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
    llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
index c0df69ffcd882..bc3464250ddc9 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
@@ -139,10 +139,19 @@ class ThreadSafeModule {
   /// 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 1ffb06db65975..1cda029a6764e 100644
--- 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