Ref-counted ModuleSlotTracker (#614)

* Make the ModuleSlotTracker ref-counted

* minor

Author: fabianbs96

.gitignore

@@ -6,6 +6,7 @@ bin/*
 # build directories for cmake
 build/
 build_*/
+build-*/
 
 # LLVM project
 llvm-project/*

lib/PhasarLLVM/Utils/LLVMShorthands.cpp

@@ -44,6 +44,7 @@
 #include <charconv>
 #include <cstdlib>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <system_error>
 
@@ -504,32 +505,44 @@ llvm::StringRef getVarAnnotationIntrinsicName(const llvm::CallInst *CallInst) {
   return Data->getAsCString();
 }
 
+struct PhasarModuleSlotTrackerWrapper {
+  PhasarModuleSlotTrackerWrapper(const llvm::Module *M) : MST(M) {}
+
+  llvm::ModuleSlotTracker MST;
+  size_t RefCount = 0;
+};
+
 static llvm::SmallDenseMap<const llvm::Module *,
-                           std::unique_ptr<llvm::ModuleSlotTracker>, 2>
+                           std::unique_ptr<PhasarModuleSlotTrackerWrapper>, 2>
     MToST{};
 
+static std::mutex MSTMx;
+
 llvm::ModuleSlotTracker &
 ModulesToSlotTracker::getSlotTrackerForModule(const llvm::Module *M) {
+  std::lock_guard Lck(MSTMx);
+
   auto &Ret = MToST[M];
   if (M == nullptr && Ret == nullptr) {
-    Ret = std::make_unique<llvm::ModuleSlotTracker>(M);
+    Ret = std::make_unique<PhasarModuleSlotTrackerWrapper>(M);
+    Ret->RefCount++;
   }
   assert(Ret != nullptr && "no ModuleSlotTracker instance for module cached");
-  return *Ret;
+  return Ret->MST;
 }
 
 void ModulesToSlotTracker::setMSTForModule(const llvm::Module *M) {
+  std::lock_guard Lck(MSTMx);
+
   auto [It, Inserted] = MToST.try_emplace(M, nullptr);
-  if (!Inserted) {
-    llvm::report_fatal_error(
-        "Cannot register the same module twice in the ModulesToSlotTracker! "
-        "Probably you have managed the same LLVM Module with multiple "
-        "ProjectIRDB instances at the same time. Don't do that!");
+  if (Inserted) {
+    It->second = std::make_unique<PhasarModuleSlotTrackerWrapper>(M);
   }
-  It->second = std::make_unique<llvm::ModuleSlotTracker>(M);
+  It->second->RefCount++;
 }
 
 void ModulesToSlotTracker::updateMSTForModule(const llvm::Module *Module) {
+  std::lock_guard Lck(MSTMx);
   auto It = MToST.find(Module);
   if (It == MToST.end()) {
     llvm::report_fatal_error(
@@ -541,7 +554,16 @@ void ModulesToSlotTracker::updateMSTForModule(const llvm::Module *Module) {
 }
 
 void ModulesToSlotTracker::deleteMSTForModule(const llvm::Module *M) {
-  MToST.erase(M);
+  std::lock_guard Lck(MSTMx);
+
+  auto It = MToST.find(M);
+  if (It == MToST.end()) {
+    return;
+  }
+
+  if (--It->second->RefCount == 0) {
+    MToST.erase(It);
+  }
 }
 
 } // namespace psr

unittests/Utils/LLVMShorthandsTest.cpp

@@ -4,6 +4,7 @@
 #include "phasar/PhasarLLVM/DB/LLVMProjectIRDB.h"
 #include "phasar/Utils/Utilities.h"
 
+#include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 
@@ -43,6 +44,24 @@ TEST(LLVMGetterTest, HandlesLLVMTermInstruction) {
   ASSERT_EQ(getNthTermInstruction(F, 5), nullptr);
 }
 
+TEST(SlotTrackerTest, HandleTwoReferences) {
+  LLVMProjectIRDB IRDB(unittest::PathToLLTestFiles +
+                       "control_flow/global_stmt_cpp.ll");
+
+  const auto *F = IRDB.getFunctionDefinition("main");
+
+  ASSERT_NE(F, nullptr);
+  const auto *Inst = getNthInstruction(F, 6);
+  llvm::StringRef InstStr = "%0 = load i32, i32* @i, align 4 | ID: 6";
+  {
+    LLVMProjectIRDB IRDB2(IRDB.getModule());
+
+    EXPECT_EQ(llvmIRToStableString(Inst), InstStr);
+  }
+
+  EXPECT_EQ(llvmIRToStableString(Inst), InstStr);
+}
+
 int main(int Argc, char **Argv) {
   ::testing::InitGoogleTest(&Argc, Argv);
   return RUN_ALL_TESTS();