fix(memory): ModuleProvider

CMakeLists.txt

@@ -72,6 +72,7 @@ nanobind_add_module(
   src/llvm/_types/PyOperandBundle.cpp
   src/llvm/_types/PyPassManagerBase.cpp
   src/llvm/_types/PyMemoryBuffer.cpp
+  src/llvm/_types/PyModuleProvider.cpp
 
   ${LLVM_INCLUDE_DIRS}
 )

src/llvm/_types.h

@@ -16,6 +16,7 @@
 #include "_types/PyOperandBundle.h"
 #include "_types/PyPassManagerBase.h"
 #include "_types/PyMemoryBuffer.h"
+#include "_types/PyModuleProvider.h"
 
 /*
   We don't define MoveOnly class to also give `Move` operation a default method
@@ -362,8 +363,9 @@ PY_FOR_EACH_TYPE_CLASS_RELASIONSHIP(DEFINE_DIRECT_SUB_CLASS)
 
 DEFINE_PY_WRAPPER_CLASS_COPYABLE(PyIntrinsic, unsigned)
 
-DEFINE_PY_WRAPPER_CLASS_SELF_DISPOSABLE_NONCOPYABLE
-  (PyModuleProvider, LLVMModuleProviderRef, LLVMDisposeModuleProvider)
+
+
+
 
 DEFINE_DIRECT_SUB_CLASS(PyPassManagerBase, PyPassManager);
 DEFINE_DIRECT_SUB_CLASS(PyPassManagerBase, PyFunctionPassManager);

src/llvm/_types/PyMemoryBuffer.h

@@ -3,7 +3,6 @@
 
 #include <llvm-c/Core.h>
 #include <memory>
-#include <stdexcept>
 #include <unordered_map>
 #include <mutex>
 

src/llvm/_types/PyMetadataEntries.h

@@ -3,7 +3,6 @@
 
 #include <llvm-c/Core.h>
 #include <memory>
-#include <stdexcept>
 #include <unordered_map>
 #include <mutex>
 

src/llvm/_types/PyModuleFlagEntries.h

@@ -2,7 +2,6 @@
 #define LLVMPYM__TYPES_PYMODULEFLAGENTRIES_H
 
 #include <llvm-c/Core.h>
-#include <stdexcept>
 #include <memory>
 #include <unordered_map>
 #include <mutex>

src/llvm/_types/PyModuleProvider.cpp

@@ -0,0 +1,40 @@
+#include "PyModuleProvider.h"
+
+std::unordered_map<LLVMModuleProviderRef,
+                   std::weak_ptr<LLVMOpaqueModuleProvider>> PyModuleProvider::mp_map;
+std::mutex PyModuleProvider::map_mutex;
+
+PyModuleProvider::PyModuleProvider(LLVMModuleProviderRef mp) : mp(get_shared_mp(mp)) {}
+
+LLVMModuleProviderRef PyModuleProvider::get() const {
+  return mp.get();
+}
+
+
+void PyModuleProvider::LLVMModuleProviderRefDeleter::operator()
+(LLVMModuleProviderRef mp) const {
+  if (mp) {
+    LLVMDisposeModuleProvider(mp);
+
+    std::lock_guard<std::mutex> lock(PyModuleProvider::map_mutex);
+    PyModuleProvider::mp_map.erase(mp);
+  }
+}
+
+
+std::shared_ptr<LLVMOpaqueModuleProvider> PyModuleProvider::get_shared_mp
+(LLVMModuleProviderRef mp) {
+  std::lock_guard<std::mutex> lock(PyModuleProvider::map_mutex);
+  auto it = PyModuleProvider::mp_map.find(mp);
+
+  if (it != PyModuleProvider::mp_map.end()) {
+    if (auto shared = it->second.lock()) {
+      return shared;
+    }
+  }
+
+  auto shared = std::shared_ptr<LLVMOpaqueModuleProvider>(mp, LLVMModuleProviderRefDeleter());
+  PyModuleProvider::mp_map[mp] = shared;
+  return shared;
+}
+

src/llvm/_types/PyModuleProvider.h

@@ -0,0 +1,29 @@
+#ifndef LLVMPYM__TYPES_PYMODULEPROVIDER_H
+#define LLVMPYM__TYPES_PYMODULEPROVIDER_H
+
+#include <llvm-c/Core.h>
+#include <memory>
+#include <unordered_map>
+#include <mutex>
+
+class PyModuleProvider {
+public:
+  explicit PyModuleProvider(LLVMModuleProviderRef mp);
+  LLVMModuleProviderRef get() const;
+
+private:
+  std::shared_ptr<LLVMOpaqueModuleProvider> mp;
+
+  struct LLVMModuleProviderRefDeleter {
+    void operator()(LLVMModuleProviderRef mp) const;
+  };
+
+  static std::shared_ptr<LLVMOpaqueModuleProvider> get_shared_mp(LLVMModuleProviderRef mp);
+
+  static std::unordered_map<LLVMModuleProviderRef,
+                            std::weak_ptr<LLVMOpaqueModuleProvider>> mp_map;
+  static std::mutex map_mutex;
+};
+
+
+#endif

src/llvm/_types/PyOperandBundle.h

@@ -3,7 +3,6 @@
 
 #include <llvm-c/Core.h>
 #include <memory>
-#include <stdexcept>
 #include <unordered_map>
 #include <mutex>
 

src/llvm/_types/PyPassManagerBase.h

@@ -3,7 +3,6 @@
 
 #include <llvm-c/Core.h>
 #include <memory>
-#include <stdexcept>
 #include <unordered_map>
 #include <mutex>