kstd: fix shared ptr deletion logic

1 files changed, 14 insertions, 8 deletions

diff --git a/libs/kstd/kstd/bits/shared_ptr.hpp b/libs/kstd/kstd/bits/shared_ptr.hpp
index 33d52e1c..7af0a00e 100644
--- a/libs/kstd/kstd/bits/shared_ptr.hpp
+++ b/libs/kstd/kstd/bits/shared_ptr.hpp
@@ -24,14 +24,16 @@ namespace kstd
       std::atomic<std::size_t> strong_count;
       //! The number of weak  references to the managed object.
       std::atomic<std::size_t> weak_count;
+      //! The deleter for the managed object
+      void (*deleter)(void *);
 
       //! Construct a new control block with a given number of references.
       //!
-      //! @param shared The number of strong references.
-      //! @param weak The number of weak references.
-      explicit shared_control_block(std::size_t shared = 1, std::size_t weak = 1)
-          : strong_count(shared)
-          , weak_count(weak)
+      //! @param deleter The function to use to delete the managed object
+      explicit shared_control_block(void (*deleter)(void *))
+          : strong_count(1)
+          , weak_count(1)
+          , deleter(deleter)
       {}
     };
 
@@ -350,7 +352,9 @@ namespace kstd
       requires(std::is_convertible_v<U *, T *>)
     explicit shared_ptr(U * pointer = nullptr)
         : m_storage(pointer)
-        , m_control_block(pointer != nullptr ? new bits::shared_control_block() : nullptr)
+        , m_control_block(pointer != nullptr
+                              ? new bits::shared_control_block([](auto object) { delete static_cast<T *>(object); })
+                              : nullptr)
     {
       assign_enable_shared_from_this(pointer);
     }
@@ -561,7 +565,9 @@ namespace kstd
     {
       cleanup();
       m_storage = ptr;
-      m_control_block = ptr != nullptr ? new bits::shared_control_block() : nullptr;
+      m_control_block = ptr != nullptr
+                            ? new bits::shared_control_block([](auto object) { delete static_cast<T *>(object); })
+                            : nullptr;
       assign_enable_shared_from_this(ptr);
     }
 
@@ -684,7 +690,7 @@ namespace kstd
       {
         if (--(m_control_block->strong_count) == 0)
         {
-          delete m_storage;
+          m_control_block->deleter(m_storage);
           m_storage = nullptr;
 
           if (--(m_control_block->weak_count) == 0)