Implement a solution to the JIT-owned llvm::Modules.

When removing code, we need to recover Clang's CodeGen and this happens by
iterating over the clang::Decls and the llvm::Module for a given Transaction in
the TransactionUnloader.

The problem is that now OrcV2 takes the ownership of the llvm::Module from the
Transaction for itself and then it can notify us when the module was compiled on
demand. Unfortuantely, there is no good way to guarantee that the same module
will be attached to the same transaction. This approach allows the
TransactionUnloader::unloadModule logic to recover CodeGen.

In a longer term, when we switch to partial translation units from clang-repl
we might be able to recover CodeGen without needing to know details about the
generated llvm::Module.

include/cling/Interpreter/Transaction.h

@@ -164,8 +164,8 @@ namespace cling {
     /// pending immediately, which sets this raw pointer. TransactionUnloader
     /// now checks this one instead of the unique pointer above. This is not
     /// nice, but it works and keeps the current infrastrucutre intact.
-    ///
-    const llvm::Module *m_UnownedModule{nullptr};
+    /// See TransactionUnloader::unloadModule
+    const llvm::Module *m_CompiledModule{nullptr};
 
     ///\brief The Executor to use m_ExeUnload on.
     ///
@@ -199,6 +199,7 @@ namespace cling {
 
     /// TransactionPool needs direct access to m_State as setState asserts
     friend class TransactionPool;
+    friend class IncrementalJIT;
 
     void Initialize();
 
@@ -490,11 +491,7 @@ namespace cling {
     }
     void setModule(std::unique_ptr<llvm::Module> M) { m_Module = std::move(M); }
 
-    const llvm::Module* getUnownedModule() const { return m_UnownedModule; }
-    void setUnownedModule(const llvm::Module *M) {
-      assert(m_UnownedModule == nullptr && "Module are emitted only once");
-      m_UnownedModule = M;
-    }
+    const llvm::Module* getCompiledModule() const { return m_CompiledModule; }
 
     IncrementalExecutor* getExecutor() const { return m_Exe; }
 

lib/Interpreter/IncrementalExecutor.cpp

@@ -121,14 +121,9 @@ IncrementalExecutor::IncrementalExecutor(clang::DiagnosticsEngine& /*diags*/,
 
   std::unique_ptr<TargetMachine> TM(CreateHostTargetMachine(CI));
   auto &TMRef = *TM;
-  auto SetReadyForUnloading = [this](const llvm::Module* M) {
-    assert (m_PendingModules.count(M) && "Unable to find the module");
-    m_PendingModules[M]->setUnownedModule(M);
-    m_PendingModules.erase(M);
-  };
   llvm::Error Err = llvm::Error::success();
   auto EPC = llvm::cantFail(llvm::orc::SelfExecutorProcessControl::Create());
-  m_JIT.reset(new IncrementalJIT(*this, std::move(TM), std::move(EPC), SetReadyForUnloading, Err));
+  m_JIT.reset(new IncrementalJIT(*this, std::move(TM), std::move(EPC), Err));
   if (Err) {
     llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "Fatal: ");
     llvm_unreachable("Propagate this error and exit gracefully");
@@ -141,7 +136,6 @@ IncrementalExecutor::IncrementalExecutor(clang::DiagnosticsEngine& /*diags*/,
                                           *m_JIT));
 }
 
-// Keep in source: ~unique_ptr<ClingJIT> needs ClingJIT
 IncrementalExecutor::~IncrementalExecutor() {}
 
 void IncrementalExecutor::runAtExitFuncs() {
@@ -254,8 +248,7 @@ IncrementalExecutor::runStaticInitializersOnce(Transaction& T) {
   if (isPracticallyEmptyModule(m))
     return kExeSuccess;
 
-  m_PendingModules[T.getModule()] = &T;
-  emitModule(T.takeModule(), T.getCompilationOpts().OptLevel);
+  emitModule(T);
 
   // We don't care whether something was unresolved before.
   m_unresolvedSymbols.clear();

lib/Interpreter/IncrementalExecutor.h

@@ -177,9 +177,8 @@ namespace cling {
     void installLazyFunctionCreator(LazyFunctionCreatorFunc_t fp);
 
     ///\brief Unload a set of JIT symbols.
-    llvm::Expected<std::unique_ptr<llvm::Module>>
-    unloadModule(const llvm::Module* M) const {
-      return m_JIT->removeModule(M);
+    llvm::Error unloadModule(const Transaction& T) const {
+      return m_JIT->removeModule(T);
     }
 
     ///\brief Run the static initializers of all modules collected to far.
@@ -252,11 +251,12 @@ namespace cling {
     ///
     /// @param[in] module - The module to pass to the execution engine.
     /// @param[in] optLevel - The optimization level to be used.
-    void emitModule(std::unique_ptr<llvm::Module> module, int optLevel) const {
+    void emitModule(Transaction &T) const {
       if (m_BackendPasses)
-        m_BackendPasses->runOnModule(*module, optLevel);
+        m_BackendPasses->runOnModule(*T.getModule(),
+                                     T.getCompilationOpts().OptLevel);
 
-      m_JIT->addModule(std::move(module));
+      m_JIT->addModule(T);
     }
 
     ///\brief Report and empty m_unresolvedSymbols.

lib/Interpreter/IncrementalJIT.cpp

@@ -68,10 +68,8 @@ public:
 
 IncrementalJIT::IncrementalJIT(
     IncrementalExecutor& Executor, std::unique_ptr<TargetMachine> TM,
-    std::unique_ptr<llvm::orc::ExecutorProcessControl> EPC,
-    ReadyForUnloadingCallback NotifyReadyForUnloading, Error& Err)
+    std::unique_ptr<llvm::orc::ExecutorProcessControl> EPC, Error& Err)
     : SkipHostProcessLookup(false),
-      NotifyReadyForUnloading(std::move(NotifyReadyForUnloading)),
       TM(std::move(TM)),
       SingleThreadedContext(std::make_unique<LLVMContext>()) {
   ErrorAsOutParameter _(&Err);
@@ -102,16 +100,14 @@ IncrementalJIT::IncrementalJIT(
     return;
   }
 
-  Jit->getIRCompileLayer().setNotifyCompiled(
-      [this](auto &MR, ThreadSafeModule TSM) {
-        // FIXME: Don't store them mapped by raw pointers.
-        const Module *Unsafe = TSM.getModuleUnlocked();
-        assert(!CompiledModules.count(Unsafe) && "Modules are compiled once");
-        assert((!Unsafe->getName().startswith("cling-module-") ||
-                ResourceTrackers.count(Unsafe)) &&
-               "All cling-modules are tracked");
-        CompiledModules[Unsafe] = std::move(TSM);
-      });
+  // We use this callback to transfer the ownership of the ThreadSafeModule,
+  // which owns the Transaction's llvm::Module, to m_CompiledModules.
+  Jit->getIRCompileLayer().setNotifyCompiled([this](auto &MR, ThreadSafeModule TSM) {
+      // FIXME: Don't store them mapped by raw pointers.
+      const Module *Unsafe = TSM.getModuleUnlocked();
+      assert(!m_CompiledModules.count(Unsafe) && "Modules are compiled once");
+      m_CompiledModules[Unsafe] = std::move(TSM);
+    });
 
   // FIXME: Make host process symbol lookup optional on a per-query basis
   char LinkerPrefix = this->TM->createDataLayout().getGlobalPrefix();
@@ -130,27 +126,33 @@ IncrementalJIT::IncrementalJIT(
   Jit->getMainJITDylib().addGenerator(std::move(Notifier));
 }
 
-void IncrementalJIT::addModule(std::unique_ptr<Module> M) {
-  const Module *RawModulePtr = M.get();
+void IncrementalJIT::addModule(Transaction& T) {
   ResourceTrackerSP RT = Jit->getMainJITDylib().createResourceTracker();
-  ResourceTrackers[RawModulePtr] = RT;
+  m_ResourceTrackers[&T] = RT;
+
+  std::ostringstream sstr;
+  sstr << T.getModule()->getModuleIdentifier() << '-' << std::hex
+       << std::showbase << (size_t)&T;
+  ThreadSafeModule TSM(T.takeModule(), SingleThreadedContext);
+
+  const Module *Unsafe = TSM.getModuleUnlocked();
+  T.m_CompiledModule = Unsafe;
 
-  ThreadSafeModule TSM(std::move(M), SingleThreadedContext);
   if (Error Err = Jit->addIRModule(RT, std::move(TSM))) {
     logAllUnhandledErrors(std::move(Err), errs(),
                           "[IncrementalJIT] addModule() failed: ");
     return;
   }
-
-  NotifyReadyForUnloading(RawModulePtr);
 }
 
-llvm::Error IncrementalJIT::removeModule(const Module* M) {
-  ResourceTrackerSP RT = std::move(ResourceTrackers[M]);
-  ResourceTrackers.erase(M);
+llvm::Error IncrementalJIT::removeModule(const Transaction& T) {
+  ResourceTrackerSP RT = std::move(m_ResourceTrackers[&T]);
+  if (!RT)
+    return llvm::Error::success();
+
+  m_ResourceTrackers.erase(&T);
   if (Error Err = RT->remove())
     return std::move(Err);
-  CompiledModules.erase(M);
   return llvm::Error::success();
 }
 

lib/Interpreter/IncrementalJIT.h

@@ -28,9 +28,7 @@
 namespace cling {
 
 class IncrementalExecutor;
-
-using ReadyForUnloadingCallback =
-    llvm::unique_function<void(const llvm::Module* M)>;
+class Transaction;
 
 class SharedAtomicFlag {
 public:
@@ -54,15 +52,14 @@ public:
   IncrementalJIT(IncrementalExecutor& Executor,
                  std::unique_ptr<llvm::TargetMachine> TM,
                  std::unique_ptr<llvm::orc::ExecutorProcessControl> EPC,
-                 ReadyForUnloadingCallback NotifyReadyForUnloading,
                  llvm::Error &Err);
 
   // FIXME: Accept a LLVMContext as well, e.g. the one that was used for the
   // particular module in Interpreter, CIFactory or BackendPasses (would be
   // more efficient)
-  void addModule(std::unique_ptr<llvm::Module> M);
+  void addModule(Transaction& T);
 
-  llvm::Error removeModule(const llvm::Module* M);
+  llvm::Error removeModule(const Transaction& T);
 
   /// Get the address of a symbol based on its IR name (as coming from clang's
   /// mangler). The ExcludeHostSymbols parameter controls whether the lookup
@@ -84,11 +81,8 @@ private:
   /// mapping via module pointers here is unnecessary. The transaction should
   /// store the resource tracker directly and pass it to `remove()` for
   /// unloading.
-  std::map<const llvm::Module *, llvm::orc::ResourceTrackerSP> ResourceTrackers;
-  std::map<const llvm::Module *, llvm::orc::ThreadSafeModule> CompiledModules;
-
-  // FIXME: Remove this once we simplified code unloading with ORCv2
-  ReadyForUnloadingCallback NotifyReadyForUnloading;
+  std::map<const Transaction*, llvm::orc::ResourceTrackerSP> m_ResourceTrackers;
+  std::map<const llvm::Module *, llvm::orc::ThreadSafeModule> m_CompiledModules;
 
   // FIXME: Move TargetMachine ownership to BackendPasses
   std::unique_ptr<llvm::TargetMachine> TM;

lib/Interpreter/TransactionUnloader.cpp

@@ -106,21 +106,33 @@ namespace cling {
     return Successful;
   }
 
+  static bool isPracticallyEmptyModule(const llvm::Module* M) {
+    return M->empty() && M->global_empty() && M->alias_empty();
+  }
+
   bool TransactionUnloader::RevertTransaction(Transaction* T) {
 
     bool Successful = true;
-    const llvm::Module *UnownedModule = T->getUnownedModule();
-    if (getExecutor() && UnownedModule) {
-      Successful = getExecutor()->unloadModule(T->getModule()) && Successful;
+    if (getExecutor()) {
+      if (llvm::Error Err = getExecutor()->unloadModule(*T)) {
+        llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "Unload: ");
+        Successful = false;
+      }
+
+      if (T->getState() == Transaction::kCommitted && !T->isNestedTransaction()) {
+        if (const llvm::Module *CompiledM = T->getCompiledModule())
+          Successful = unloadModule(const_cast<llvm::Module*>(CompiledM)) &&
+            Successful;
+        else
+          assert((!T->getModule() || isPracticallyEmptyModule(T->getModule()))
+                 && "Must have already compiled this module");
+      }
 
       // Cleanup the module from unused global values.
       // if (T->getModule()) {
       //   llvm::ModulePass* globalDCE = llvm::createGlobalDCEPass();
       //   globalDCE->runOnModule(*T->getModule());
       // }
-
-      Successful =
-          unloadModule(const_cast<llvm::Module *>(UnownedModule)) && Successful;
     }
 
     // Clean up the pending instantiations