[ORC] Replace ThreadSafeContext::getContext with withContextDo. (#146819)

This removes ThreadSafeContext::Lock, ThreadSafeContext::getLock, and
ThreadSafeContext::getContext, and replaces them with a
ThreadSafeContext::withContextDo method (and const override).

The new method can be used to access an existing
ThreadSafeContext-wrapped LLVMContext in a safe way:

ThreadSafeContext TSCtx = ... ;
TSCtx.withContextDo([](LLVMContext *Ctx) {
  // this closure has exclusive access to Ctx.
});

The new API enforces correct locking, whereas the old APIs relied on
manual locking (which almost no in-tree code preformed, relying instead
on incidental exclusive access to the ThreadSafeContext).

clang/lib/Interpreter/Interpreter.cpp

@@ -373,8 +373,11 @@ Interpreter::Interpreter(std::unique_ptr<CompilerInstance> Instance,
   auto LLVMCtx = std::make_unique<llvm::LLVMContext>();
   TSCtx = std::make_unique<llvm::orc::ThreadSafeContext>(std::move(LLVMCtx));
 
-  Act = std::make_unique<IncrementalAction>(*CI, *TSCtx->getContext(), ErrOut,
-                                            *this, std::move(Consumer));
+  Act = TSCtx->withContextDo([&](llvm::LLVMContext *Ctx) {
+    return std::make_unique<IncrementalAction>(*CI, *Ctx, ErrOut, *this,
+                                               std::move(Consumer));
+  });
+
   if (ErrOut)
     return;
   CI->ExecuteAction(*Act);
@@ -495,10 +498,10 @@ Interpreter::createWithCUDA(std::unique_ptr<CompilerInstance> CI,
   std::unique_ptr<Interpreter> Interp = std::move(*InterpOrErr);
 
   llvm::Error Err = llvm::Error::success();
-  llvm::LLVMContext &LLVMCtx = *Interp->TSCtx->getContext();
 
-  auto DeviceAct =
-      std::make_unique<IncrementalAction>(*DCI, LLVMCtx, Err, *Interp);
+  auto DeviceAct = Interp->TSCtx->withContextDo([&](llvm::LLVMContext *Ctx) {
+    return std::make_unique<IncrementalAction>(*DCI, *Ctx, Err, *Interp);
+  });
 
   if (Err)
     return std::move(Err);

llvm/examples/OrcV2Examples/LLJITWithThinLTOSummaries/LLJITWithThinLTOSummaries.cpp

@@ -169,7 +169,9 @@ Expected<ThreadSafeModule> loadModule(StringRef Path,
 
   MemoryBufferRef BitcodeBufferRef = (**BitcodeBuffer).getMemBufferRef();
   Expected<std::unique_ptr<Module>> M =
-      parseBitcodeFile(BitcodeBufferRef, *TSCtx.getContext());
+      TSCtx.withContextDo([&](LLVMContext *Ctx) {
+        return parseBitcodeFile(BitcodeBufferRef, *Ctx);
+      });
   if (!M)
     return M.takeError();
 

llvm/examples/OrcV2Examples/OrcV2CBindingsBasicUsage/OrcV2CBindingsBasicUsage.c

@@ -22,11 +22,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -57,6 +54,9 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Free the builder.
   LLVMDisposeBuilder(Builder);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);

llvm/examples/OrcV2Examples/OrcV2CBindingsDumpObjects/OrcV2CBindingsDumpObjects.c

@@ -31,8 +31,7 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  LLVMContextRef Ctx = LLVMContextCreate();
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
   LLVMTypeRef ParamTypes[] = {LLVMInt32Type(), LLVMInt32Type()};
   LLVMTypeRef SumFunctionType =
@@ -45,6 +44,8 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   LLVMValueRef SumArg1 = LLVMGetParam(SumFunction, 1);
   LLVMValueRef Result = LLVMBuildAdd(Builder, SumArg0, SumArg1, "result");
   LLVMBuildRet(Builder, Result);
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
   LLVMOrcDisposeThreadSafeContext(TSCtx);
   return TSM;

llvm/examples/OrcV2Examples/OrcV2CBindingsIRTransforms/OrcV2CBindingsIRTransforms.c

@@ -32,8 +32,7 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  LLVMContextRef Ctx = LLVMContextCreate();
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
   LLVMTypeRef ParamTypes[] = {LLVMInt32Type(), LLVMInt32Type()};
   LLVMTypeRef SumFunctionType =
@@ -47,6 +46,7 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   LLVMValueRef Result = LLVMBuildAdd(Builder, SumArg0, SumArg1, "result");
   LLVMBuildRet(Builder, Result);
   LLVMDisposeBuilder(Builder);
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
   LLVMOrcDisposeThreadSafeContext(TSCtx);
   return TSM;

llvm/examples/OrcV2Examples/OrcV2CBindingsLazy/OrcV2CBindingsLazy.c

@@ -67,11 +67,9 @@ const char MainMod[] =
 LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
                                 const char *Name,
                                 LLVMOrcThreadSafeModuleRef *TSM) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
 
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext for the Module.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Wrap Source in a MemoryBuffer
   LLVMMemoryBufferRef MB =
@@ -85,6 +83,10 @@ LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
     // TODO: LLVMDisposeMessage(ErrMsg);
   }
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   *TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);

llvm/examples/OrcV2Examples/OrcV2CBindingsMCJITLikeMemoryManager/OrcV2CBindingsMCJITLikeMemoryManager.c

@@ -150,11 +150,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -182,6 +179,10 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Build the return instruction.
   LLVMBuildRet(Builder, Result);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);

llvm/examples/OrcV2Examples/OrcV2CBindingsRemovableCode/OrcV2CBindingsRemovableCode.c

@@ -22,11 +22,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -57,6 +54,10 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Free the builder.
   LLVMDisposeBuilder(Builder);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);

llvm/examples/OrcV2Examples/OrcV2CBindingsVeryLazy/OrcV2CBindingsVeryLazy.c

@@ -74,11 +74,8 @@ LLVMErrorRef applyDataLayout(void *Ctx, LLVMModuleRef M) {
 LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
                                 const char *Name,
                                 LLVMOrcThreadSafeModuleRef *TSM) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Wrap Source in a MemoryBuffer
   LLVMMemoryBufferRef MB =
@@ -93,6 +90,9 @@ LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
     return Err;
   }
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
+
   // Our module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   *TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);

llvm/include/llvm-c/Orc.h

@@ -1062,20 +1062,32 @@ LLVMErrorRef LLVMOrcCreateStaticLibrarySearchGeneratorForPath(
     const char *FileName);
 
 /**
- * Create a ThreadSafeContext containing a new LLVMContext.
+ * Create a ThreadSafeContextRef containing a new LLVMContext.
  *
  * Ownership of the underlying ThreadSafeContext data is shared: Clients
- * can and should dispose of their ThreadSafeContext as soon as they no longer
- * need to refer to it directly. Other references (e.g. from ThreadSafeModules)
- * will keep the data alive as long as it is needed.
+ * can and should dispose of their ThreadSafeContextRef as soon as they no
+ * longer need to refer to it directly. Other references (e.g. from
+ * ThreadSafeModules) will keep the underlying data alive as long as it is
+ * needed.
  */
 LLVMOrcThreadSafeContextRef LLVMOrcCreateNewThreadSafeContext(void);
 
 /**
- * Get a reference to the wrapped LLVMContext.
+ * Create a ThreadSafeContextRef from a given LLVMContext, which must not be
+ * associated with any existing ThreadSafeContext.
+ *
+ * The underlying ThreadSafeContext will take ownership of the LLVMContext
+ * object, so clients should not free the LLVMContext passed to this
+ * function.
+ *
+ * Ownership of the underlying ThreadSafeContext data is shared: Clients
+ * can and should dispose of their ThreadSafeContextRef as soon as they no
+ * longer need to refer to it directly. Other references (e.g. from
+ * ThreadSafeModules) will keep the underlying data alive as long as it is
+ * needed.
  */
-LLVMContextRef
-LLVMOrcThreadSafeContextGetContext(LLVMOrcThreadSafeContextRef TSCtx);
+LLVMOrcThreadSafeContextRef
+LLVMOrcCreateNewThreadSafeContextFromLLVMContext(LLVMContextRef Ctx);
 
 /**
  * Dispose of a ThreadSafeContext.

llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h

@@ -36,16 +36,6 @@ private:
   };
 
 public:
-  // RAII based lock for ThreadSafeContext.
-  class [[nodiscard]] Lock {
-  public:
-    Lock(std::shared_ptr<State> S) : S(std::move(S)), L(this->S->Mutex) {}
-
-  private:
-    std::shared_ptr<State> S;
-    std::unique_lock<std::recursive_mutex> L;
-  };
-
   /// Construct a null context.
   ThreadSafeContext() = default;
 
@@ -56,17 +46,20 @@ public:
            "Can not construct a ThreadSafeContext from a nullptr");
   }
 
-  /// Returns a pointer to the LLVMContext that was used to construct this
-  /// instance, or null if the instance was default constructed.
-  LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
+  template <typename Func> decltype(auto) withContextDo(Func &&F) {
+    if (auto TmpS = S) {
+      std::lock_guard<std::recursive_mutex> Lock(TmpS->Mutex);
+      return F(TmpS->Ctx.get());
+    } else
+      return F((LLVMContext *)nullptr);
+  }
 
-  /// Returns a pointer to the LLVMContext that was used to construct this
-  /// instance, or null if the instance was default constructed.
-  const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
-
-  Lock getLock() const {
-    assert(S && "Can not lock an empty ThreadSafeContext");
-    return Lock(S);
+  template <typename Func> decltype(auto) withContextDo(Func &&F) const {
+    if (auto TmpS = S) {
+      std::lock_guard<std::recursive_mutex> Lock(TmpS->Mutex);
+      return F(const_cast<const LLVMContext *>(TmpS->Ctx.get()));
+    } else
+      return F((const LLVMContext *)nullptr);
   }
 
 private:
@@ -89,10 +82,7 @@ public:
     // *before* the context that it depends on.
     // We also need to lock the context to make sure the module tear-down
     // does not overlap any other work on the context.
-    if (M) {
-      auto L = TSCtx.getLock();
-      M = nullptr;
-    }
+    TSCtx.withContextDo([this](LLVMContext *Ctx) { M = nullptr; });
     M = std::move(Other.M);
     TSCtx = std::move(Other.TSCtx);
     return *this;
@@ -111,45 +101,39 @@ public:
 
   ~ThreadSafeModule() {
     // We need to lock the context while we destruct the module.
-    if (M) {
-      auto L = TSCtx.getLock();
-      M = nullptr;
-    }
+    TSCtx.withContextDo([this](LLVMContext *Ctx) { M = nullptr; });
   }
 
   /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
   /// wraps a non-null module.
-  explicit operator bool() const {
-    if (M) {
-      assert(TSCtx.getContext() &&
-             "Non-null module must have non-null context");
-      return true;
-    }
-    return false;
-  }
+  explicit operator bool() const { return !!M; }
 
   /// Locks the associated ThreadSafeContext and calls the given function
   /// on the contained Module.
   template <typename Func> decltype(auto) withModuleDo(Func &&F) {
-    assert(M && "Can not call on null module");
-    auto Lock = TSCtx.getLock();
-    return F(*M);
+    return TSCtx.withContextDo([&](LLVMContext *) {
+      assert(M && "Can not call on null module");
+      return F(*M);
+    });
   }
 
   /// 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);
+    return TSCtx.withContextDo([&](const LLVMContext *) {
+      assert(M && "Can not call on null module");
+      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) consumingModuleDo(Func &&F) {
-    auto Lock = TSCtx.getLock();
-    return F(std::move(M));
+    return TSCtx.withContextDo([&](LLVMContext *) {
+      assert(M && "Can not call on null module");
+      return F(std::move(M));
+    });
   }
 
   /// Get a raw pointer to the contained module without locking the context.

llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp

@@ -729,9 +729,9 @@ LLVMOrcThreadSafeContextRef LLVMOrcCreateNewThreadSafeContext(void) {
   return wrap(new ThreadSafeContext(std::make_unique<LLVMContext>()));
 }
 
-LLVMContextRef
-LLVMOrcThreadSafeContextGetContext(LLVMOrcThreadSafeContextRef TSCtx) {
-  return wrap(unwrap(TSCtx)->getContext());
+LLVMOrcThreadSafeContextRef
+LLVMOrcCreateNewThreadSafeContextFromLLVMContext(LLVMContextRef Ctx) {
+  return wrap(new ThreadSafeContext(std::unique_ptr<LLVMContext>(unwrap(Ctx))));
 }
 
 void LLVMOrcDisposeThreadSafeContext(LLVMOrcThreadSafeContextRef TSCtx) {

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

@@ -60,8 +60,6 @@ void IRSpeculationLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
                               ThreadSafeModule TSM) {
 
   assert(TSM && "Speculation Layer received Null Module ?");
-  assert(TSM.getContext().getContext() != nullptr &&
-         "Module with null LLVMContext?");
 
   // Instrumentation of runtime calls, lock the Module
   TSM.withModuleDo([this, &R](Module &M) {

llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp

@@ -53,9 +53,11 @@ ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,
         "cloned module buffer");
     ThreadSafeContext NewTSCtx(std::make_unique<LLVMContext>());
 
-    auto ClonedModule = cantFail(
-        parseBitcodeFile(ClonedModuleBufferRef, *NewTSCtx.getContext()));
-    ClonedModule->setModuleIdentifier(M.getName());
+    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));
   });
 }

llvm/tools/lli/lli.cpp

@@ -878,7 +878,8 @@ static void exitOnLazyCallThroughFailure() { exit(1); }
 Expected<orc::ThreadSafeModule>
 loadModule(StringRef Path, orc::ThreadSafeContext TSCtx) {
   SMDiagnostic Err;
-  auto M = parseIRFile(Path, Err, *TSCtx.getContext());
+  auto M = TSCtx.withContextDo(
+      [&](LLVMContext *Ctx) { return parseIRFile(Path, Err, *Ctx); });
   if (!M) {
     std::string ErrMsg;
     {

llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp

@@ -128,29 +128,25 @@ TEST(RTDyldObjectLinkingLayerTest, TestOverrideObjectFlags) {
   };
 
   // Create a module with two void() functions: foo and bar.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
   ThreadSafeModule M;
   {
-    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
+    auto Ctx = std::make_unique<LLVMContext>();
+    ModuleBuilder MB(*Ctx, TM->getTargetTriple().str(), "dummy");
     MB.getModule()->setDataLayout(TM->createDataLayout());
 
     Function *FooImpl = MB.createFunctionDecl(
-        FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
-        "foo");
-    BasicBlock *FooEntry =
-        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
+        FunctionType::get(Type::getVoidTy(*Ctx), {}, false), "foo");
+    BasicBlock *FooEntry = BasicBlock::Create(*Ctx, "entry", FooImpl);
     IRBuilder<> B1(FooEntry);
     B1.CreateRetVoid();
 
     Function *BarImpl = MB.createFunctionDecl(
-        FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
-        "bar");
-    BasicBlock *BarEntry =
-        BasicBlock::Create(*TSCtx.getContext(), "entry", BarImpl);
+        FunctionType::get(Type::getVoidTy(*Ctx), {}, false), "bar");
+    BasicBlock *BarEntry = BasicBlock::Create(*Ctx, "entry", BarImpl);
     IRBuilder<> B2(BarEntry);
     B2.CreateRetVoid();
 
-    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
+    M = ThreadSafeModule(MB.takeModule(), std::move(Ctx));
   }
 
   // Create a simple stack and set the override flags option.
@@ -207,21 +203,19 @@ TEST(RTDyldObjectLinkingLayerTest, TestAutoClaimResponsibilityForSymbols) {
   };
 
   // Create a module with two void() functions: foo and bar.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
   ThreadSafeModule M;
   {
-    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
+    auto Ctx = std::make_unique<LLVMContext>();
+    ModuleBuilder MB(*Ctx, TM->getTargetTriple().str(), "dummy");
     MB.getModule()->setDataLayout(TM->createDataLayout());
 
     Function *FooImpl = MB.createFunctionDecl(
-        FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
-        "foo");
-    BasicBlock *FooEntry =
-        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
+        FunctionType::get(Type::getVoidTy(*Ctx), {}, false), "foo");
+    BasicBlock *FooEntry = BasicBlock::Create(*Ctx, "entry", FooImpl);
     IRBuilder<> B(FooEntry);
     B.CreateRetVoid();
 
-    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
+    M = ThreadSafeModule(MB.takeModule(), std::move(Ctx));
   }
 
   // Create a simple stack and set the override flags option.
@@ -258,21 +252,19 @@ TEST(RTDyldObjectLinkingLayerTest, TestMemoryBufferNamePropagation) {
     GTEST_SKIP();
 
   // Create a module with two void() functions: foo and bar.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
   ThreadSafeModule M;
   {
-    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
+    auto Ctx = std::make_unique<LLVMContext>();
+    ModuleBuilder MB(*Ctx, TM->getTargetTriple().str(), "dummy");
     MB.getModule()->setDataLayout(TM->createDataLayout());
 
     Function *FooImpl = MB.createFunctionDecl(
-        FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
-        "foo");
-    BasicBlock *FooEntry =
-        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
+        FunctionType::get(Type::getVoidTy(*Ctx), {}, false), "foo");
+    BasicBlock *FooEntry = BasicBlock::Create(*Ctx, "entry", FooImpl);
     IRBuilder<> B1(FooEntry);
     B1.CreateRetVoid();
 
-    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
+    M = ThreadSafeModule(MB.takeModule(), std::move(Ctx));
   }
 
   ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};

llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp

@@ -172,8 +172,8 @@ TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
       });
   EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded());
 
-  ThreadSafeContext Ctx(std::make_unique<LLVMContext>());
-  auto M = std::make_unique<Module>("<main>", *Ctx.getContext());
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = std::make_unique<Module>("<main>", *Ctx);
   M->setTargetTriple(Triple(sys::getProcessTriple()));
 
   (void)createRetFunction(M.get(), "main", 42);

llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp

@@ -21,20 +21,21 @@ namespace {
 TEST(ThreadSafeModuleTest, ContextWhollyOwnedByOneModule) {
   // Test that ownership of a context can be transferred to a single
   // ThreadSafeModule.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  auto M = std::make_unique<Module>("M", *TSCtx.getContext());
-  ThreadSafeModule TSM(std::move(M), std::move(TSCtx));
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = std::make_unique<Module>("M", *Ctx);
+  ThreadSafeModule TSM(std::move(M), std::move(Ctx));
 }
 
 TEST(ThreadSafeModuleTest, ContextOwnershipSharedByTwoModules) {
   // Test that ownership of a context can be shared between more than one
   // ThreadSafeModule.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  auto Ctx = std::make_unique<LLVMContext>();
 
-  auto M1 = std::make_unique<Module>("M1", *TSCtx.getContext());
+  auto M1 = std::make_unique<Module>("M1", *Ctx);
+  auto M2 = std::make_unique<Module>("M2", *Ctx);
+
+  ThreadSafeContext TSCtx(std::move(Ctx));
   ThreadSafeModule TSM1(std::move(M1), TSCtx);
-
-  auto M2 = std::make_unique<Module>("M2", *TSCtx.getContext());
   ThreadSafeModule TSM2(std::move(M2), std::move(TSCtx));
 }
 
@@ -45,12 +46,14 @@ TEST(ThreadSafeModuleTest, ContextOwnershipSharedWithClient) {
 
   {
     // Create and destroy a module.
-    auto M1 = std::make_unique<Module>("M1", *TSCtx.getContext());
+    auto M1 = TSCtx.withContextDo(
+        [](LLVMContext *Ctx) { return std::make_unique<Module>("M1", *Ctx); });
     ThreadSafeModule TSM1(std::move(M1), TSCtx);
   }
 
   // Verify that the context is still available for re-use.
-  auto M2 = std::make_unique<Module>("M2", *TSCtx.getContext());
+  auto M2 = TSCtx.withContextDo(
+      [](LLVMContext *Ctx) { return std::make_unique<Module>("M2", *Ctx); });
   ThreadSafeModule TSM2(std::move(M2), std::move(TSCtx));
 }
 
@@ -59,59 +62,44 @@ TEST(ThreadSafeModuleTest, ThreadSafeModuleMoveAssignment) {
   // to the field order) to ensure that overwriting with an empty
   // ThreadSafeModule does not destroy the context early.
   ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  auto M = std::make_unique<Module>("M", *TSCtx.getContext());
+  auto M = TSCtx.withContextDo(
+      [](LLVMContext *Ctx) { return std::make_unique<Module>("M", *Ctx); });
   ThreadSafeModule TSM(std::move(M), std::move(TSCtx));
   TSM = ThreadSafeModule();
 }
 
-TEST(ThreadSafeModuleTest, BasicContextLockAPI) {
-  // Test that basic lock API calls work.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  auto M = std::make_unique<Module>("M", *TSCtx.getContext());
-  ThreadSafeModule TSM(std::move(M), TSCtx);
+TEST(ThreadSafeModuleTest, WithContextDoPreservesContext) {
+  // Test that withContextDo passes through the LLVMContext that was used
+  // to create the ThreadSafeContext.
 
-  { auto L = TSCtx.getLock(); }
-
-  { auto L = TSM.getContext().getLock(); }
-}
-
-TEST(ThreadSafeModuleTest, ContextLockPreservesContext) {
-  // Test that the existence of a context lock preserves the attached
-  // context.
-  // The trick to verify this is a bit of a hack: We attach a Module
-  // (without the ThreadSafeModule wrapper) to the context, then verify
-  // that this Module destructs safely (which it will not if its context
-  // has been destroyed) even though all references to the context have
-  // been thrown away (apart from the lock).
-
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  auto L = TSCtx.getLock();
-  auto &Ctx = *TSCtx.getContext();
-  auto M = std::make_unique<Module>("M", Ctx);
-  TSCtx = ThreadSafeContext();
+  auto Ctx = std::make_unique<LLVMContext>();
+  LLVMContext *OriginalCtx = Ctx.get();
+  ThreadSafeContext TSCtx(std::move(Ctx));
+  TSCtx.withContextDo(
+      [&](LLVMContext *ClosureCtx) { EXPECT_EQ(ClosureCtx, OriginalCtx); });
 }
 
 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);
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = std::make_unique<Module>("M", *Ctx);
+  ThreadSafeModule TSM(std::move(M), std::move(Ctx));
   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);
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = std::make_unique<Module>("M", *Ctx);
+  const ThreadSafeModule TSM(std::move(M), std::move(Ctx));
   TSM.withModuleDo([](const Module &M) {});
 }
 
 TEST(ThreadSafeModuleTest, ConsumingModuleDo) {
   // Test consumingModuleDo.
-  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  ThreadSafeModule TSM(std::make_unique<Module>("M", *TSCtx.getContext()),
-                       TSCtx);
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = std::make_unique<Module>("M", *Ctx);
+  ThreadSafeModule TSM(std::move(M), std::move(Ctx));
   TSM.consumingModuleDo([](std::unique_ptr<Module> M) {});
 }