implement simple weak_ptr and enable_shared_from_this

1 files changed, 286 insertions, 34 deletions

diff --git a/libs/kstd/include/kstd/bits/shared_ptr.hpp b/libs/kstd/include/kstd/bits/shared_ptr.hpp
index 6bce83f..674807d 100644
--- a/libs/kstd/include/kstd/bits/shared_ptr.hpp
+++ b/libs/kstd/include/kstd/bits/shared_ptr.hpp
@@ -10,10 +10,200 @@
 
 namespace kstd
 {
+  /**
+   * @brief Control block for shared_ptr and weak_ptr. This control block contains the reference counts for shared
+   * ownership and weak ownership. The shared_count tracks the number of shared_ptr instances that own the managed
+   * object, while the weak_count tracks the number of weak_ptr instances that reference the managed object.
+   * The weak_count is needed to determine when it is safe to delete the shared_control_block itself
+   */
+  struct shared_control_block
+  {
+    std::atomic<std::size_t> shared_count;
+    std::atomic<std::size_t> weak_count;
+
+    explicit shared_control_block(std::size_t shared = 1, std::size_t weak = 0)
+        : shared_count(shared)
+        , weak_count(weak)
+    {}
+  };
+
   template<typename T>
   struct shared_ptr;
 
   /**
+   * @brief weak_ptr is a smart pointer that holds a non-owning weak reference to an object that is managed by
+   * shared_ptr. It must be converted to shared_ptr in to be able to access the referenced object. A weak_ptr is created
+   * as a copy of a shared_ptr, or as a copy of another weak_ptr. A weak_ptr is typically used to break circular
+   * references between shared_ptr to avoid memory leaks. A weak_ptr does not contribute to the reference count of the
+   * object, and it does not manage the lifetime of the object. If the object managed by shared_ptr is destroyed, the
+   * weak_ptr becomes expired and cannot be used to access the object anymore.
+   */
+  template<typename T>
+  struct weak_ptr
+  {
+    template<typename U>
+    friend struct shared_ptr;
+
+    /**
+     * @brief Constructs a null weak_ptr.
+     */
+    weak_ptr() noexcept
+        : pointer(nullptr)
+        , control(nullptr)
+    {}
+
+    template<typename U>
+      requires(std::is_convertible_v<U *, T *>)
+    weak_ptr(shared_ptr<U> const & other)
+        : pointer(other.pointer)
+        , control(other.control)
+    {
+      if (control != nullptr)
+      {
+        ++(control->weak_count);
+      }
+    }
+
+    /**
+     * @brief Copy constructor. Constructs a weak_ptr which shares ownership of the object managed by other.
+     */
+    weak_ptr(weak_ptr const & other)
+        : pointer(other.pointer)
+        , control(other.control)
+    {
+      if (control != nullptr)
+      {
+        ++(control->weak_count);
+      }
+    }
+
+    /**
+     * @brief Move constructor. Constructs a weak_ptr which takes ownership of the object managed by other.
+     */
+    weak_ptr(weak_ptr && other) noexcept
+        : pointer(other.pointer)
+        , control(other.control)
+    {
+      other.pointer = nullptr;
+      other.control = nullptr;
+    }
+
+    /**
+     * @brief Assignment operator. Assigns the weak_ptr to another weak_ptr.
+     */
+    auto operator=(weak_ptr const & other) -> weak_ptr &
+    {
+      if (this != &other)
+      {
+        cleanup();
+        pointer = other.pointer;
+        control = other.control;
+        if (control != nullptr)
+        {
+          ++(control->weak_count);
+        }
+      }
+
+      return *this;
+    }
+
+    /**
+     * @brief Move assignment operator. Move-assigns a weak_ptr from other.
+     */
+    auto operator=(weak_ptr && other) noexcept -> weak_ptr &
+    {
+      if (this != &other)
+      {
+        cleanup();
+        pointer = other.pointer;
+        control = other.control;
+        other.pointer = nullptr;
+        other.control = nullptr;
+      }
+
+      return *this;
+    }
+
+    /**
+     * @brief Destructor. Cleans up resources if necessary.
+     */
+    ~weak_ptr()
+    {
+      cleanup();
+    }
+
+    /**
+     * @brief Returns a shared_ptr that shares ownership of the managed object if the managed object still exists, or an
+     * empty shared_ptr otherwise.
+     */
+    [[nodiscard]] auto lock() const -> shared_ptr<T>
+    {
+      return shared_ptr<T>(*this);
+    }
+
+  private:
+    auto cleanup() -> void
+    {
+      if (control != nullptr)
+      {
+        if (--(control->weak_count) == 0 && control->shared_count == 0)
+        {
+          delete control;
+        }
+      }
+    }
+
+    T * pointer;
+    shared_control_block * control;
+  };
+
+  /**
+   * @brief enable_shared_from_this is a base class that allows an object that is currently managed by a shared_ptr to
+   * create additional shared_ptr instances that share ownership of the same object. This is usefl when you want to
+   * create shared_ptr instances in a member function of the object.
+   *
+   * @tparam T The type of the managed object.
+   */
+  template<typename T>
+  struct enable_shared_from_this
+  {
+    template<typename U>
+    friend struct shared_ptr;
+
+    friend T;
+
+  public:
+    /**
+     * @brief Returns a shared_ptr that shares ownership of *this.
+     */
+    auto shared_from_this() -> shared_ptr<T>
+    {
+      return shared_ptr<T>(weak_this);
+    }
+
+    /**
+     * @brief Returns a shared_ptr that shares ownership of *this.
+     */
+    auto shared_from_this() const -> shared_ptr<T const>
+    {
+      return shared_ptr<T const>(weak_this);
+    }
+
+  private:
+    enable_shared_from_this() = default;
+    enable_shared_from_this(enable_shared_from_this const &) = default;
+    auto operator=(enable_shared_from_this const &) -> enable_shared_from_this & = default;
+    ~enable_shared_from_this() = default;
+
+    void internal_assign_ptr(shared_ptr<T> const & ptr) const
+    {
+      weak_this = ptr;
+    }
+
+    mutable weak_ptr<T> weak_this{};  ///< Weak pointer to the object, used for shared_from_this functionality.
+  };
+
+  /**
    * @brief Shared_pointer is a smart pointer that retains shared ownership of an object through a pointer. Several
    * shared_ptr objects may own the same object. The object is destroyed and its memory deallocated when either of
    * the following happens: the last remaining shared_ptr owning the object is destroyed; the last remaining
@@ -31,12 +221,23 @@ namespace kstd
     template<typename U>
     friend struct shared_ptr;
 
+    template<typename U>
+    friend struct weak_ptr;
+
+    /**
+     * @brief Construct an empty shared_ptr.
+     */
+    shared_ptr() noexcept
+        : pointer(nullptr)
+        , control(nullptr)
+    {}
+
     /**
      * @brief Construct an empty shared_ptr from nullptr.
      */
     shared_ptr(std::nullptr_t) noexcept
         : pointer(nullptr)
-        , ref_count(nullptr)
+        , control(nullptr)
     {}
 
     /**
@@ -44,11 +245,33 @@ namespace kstd
      *
      * @param pointer A pointer to an object to manage (default is nullptr).
      */
-    explicit shared_ptr(T * pointer = nullptr)
+    template<typename U>
+      requires(std::is_convertible_v<U *, T *>)
+    explicit shared_ptr(U * pointer = nullptr)
         : pointer(pointer)
-        , ref_count(pointer != nullptr ? new std::atomic<std::size_t>(1) : nullptr)
+        , control(pointer != nullptr ? new shared_control_block() : nullptr)
+    {
+      assign_enable_shared_from_this(pointer);
+    }
+
+    /**
+     * @brief Constructor a shared_ptr from a weak_ptr. If other is not expired, constructs a shared_ptr which shares
+     * ownership of the object managed by other. Otherwise, constructs an empty shared_ptr.
+     *
+     * @param other The weak_ptr to construct from.
+     */
+    template<typename U>
+      requires(std::is_convertible_v<U *, T *>)
+    explicit shared_ptr(weak_ptr<U> const & other)
+        : pointer(nullptr)
+        , control(nullptr)
     {
-      // Nothing to do.
+      if (other.control != nullptr && other.control->shared_count != 0)
+      {
+        pointer = other.pointer;
+        control = other.control;
+        ++(control->shared_count);
+      }
     }
 
     /**
@@ -58,11 +281,11 @@ namespace kstd
      */
     shared_ptr(shared_ptr const & other)
         : pointer(other.pointer)
-        , ref_count(other.ref_count)
+        , control(other.control)
     {
-      if (ref_count != nullptr)
+      if (control != nullptr)
       {
-        ++(*ref_count);
+        ++(control->shared_count);
       }
     }
 
@@ -76,11 +299,11 @@ namespace kstd
       requires(std::is_convertible_v<U *, T *>)
     shared_ptr(shared_ptr<U> const & other)
         : pointer(other.pointer)
-        , ref_count(other.ref_count)
+        , control(other.control)
     {
-      if (ref_count != nullptr)
+      if (control != nullptr)
       {
-        ++(*ref_count);
+        ++(control->shared_count);
       }
     }
 
@@ -91,10 +314,10 @@ namespace kstd
      */
     shared_ptr(shared_ptr && other) noexcept
         : pointer(other.pointer)
-        , ref_count(other.ref_count)
+        , control(other.control)
     {
       other.pointer = nullptr;
-      other.ref_count = nullptr;
+      other.control = nullptr;
     }
 
     /**
@@ -107,10 +330,10 @@ namespace kstd
       requires(std::is_convertible_v<U *, T *>)
     shared_ptr(shared_ptr<U> && other) noexcept
         : pointer(other.pointer)
-        , ref_count(other.ref_count)
+        , control(other.control)
     {
       other.pointer = nullptr;
-      other.ref_count = nullptr;
+      other.control = nullptr;
     }
 
     /**
@@ -127,11 +350,11 @@ namespace kstd
       {
         cleanup();
         pointer = other.pointer;
-        ref_count = other.ref_count;
+        control = other.control;
 
-        if (ref_count != nullptr)
+        if (control != nullptr)
         {
-          ++(*ref_count);
+          ++(control->shared_count);
         }
       }
 
@@ -151,11 +374,11 @@ namespace kstd
     {
       cleanup();
       pointer = other.pointer;
-      ref_count = other.ref_count;
+      control = other.control;
 
-      if (ref_count != nullptr)
+      if (control != nullptr)
       {
-        ++(*ref_count);
+        ++(control->shared_count);
       }
 
       return *this;
@@ -174,9 +397,9 @@ namespace kstd
       {
         cleanup();
         pointer = other.pointer;
-        ref_count = other.ref_count;
+        control = other.control;
         other.pointer = nullptr;
-        other.ref_count = nullptr;
+        other.control = nullptr;
       }
 
       return *this;
@@ -195,9 +418,9 @@ namespace kstd
     {
       cleanup();
       pointer = other.pointer;
-      ref_count = other.ref_count;
+      control = other.control;
       other.pointer = nullptr;
-      other.ref_count = nullptr;
+      other.control = nullptr;
 
       return *this;
     }
@@ -209,7 +432,7 @@ namespace kstd
     {
       cleanup();
       pointer = nullptr;
-      ref_count = nullptr;
+      control = nullptr;
       return *this;
     }
 
@@ -230,7 +453,8 @@ namespace kstd
     {
       cleanup();
       pointer = ptr;
-      ref_count = ptr != nullptr ? new std::atomic<std::size_t>(1) : nullptr;
+      control = ptr != nullptr ? new shared_control_block() : nullptr;
+      assign_enable_shared_from_this(ptr);
     }
 
     /**
@@ -242,7 +466,7 @@ namespace kstd
     void swap(shared_ptr & other)
     {
       std::swap(pointer, other.pointer);
-      std::swap(ref_count, other.ref_count);
+      std::swap(control, other.control);
     }
 
     /**
@@ -288,9 +512,9 @@ namespace kstd
      */
     [[nodiscard]] auto use_count() const -> std::size_t
     {
-      if (ref_count != nullptr)
+      if (control != nullptr)
       {
-        return *ref_count;
+        return control->shared_count;
       }
 
       return 0;
@@ -329,19 +553,47 @@ namespace kstd
 
   private:
     /**
+     * @brief If the candidate type inherits from enable_shared_from_this, assigns the internal weak pointer to this
+     * shared_ptr. This weak_ptr is used to implement shared_from_this functionality for the candidate type. If the
+     * candidate type does not inherit from enable_shared_from_this, this function does nothing.
+     *
+     * @tparam U The candidate type to check for enable_shared_from_this inheritance.
+     * @param candidate The candidate object to assign the internal weak pointer for.
+     */
+    template<typename U>
+    auto assign_enable_shared_from_this(U * candidate) -> void
+    {
+      if constexpr (requires(U * p, shared_ptr<T> const & sp) { p->internal_assign_ptr(sp); })
+      {
+        if (candidate != nullptr)
+        {
+          candidate->internal_assign_ptr(*this);
+        }
+      }
+    }
+
+    /**
      * @brief Releases ownership and deletes the object if this was the last reference to the owned managed object.
      */
     auto cleanup() -> void
     {
-      if (ref_count != nullptr && --(*ref_count) == 0)
+      if (control != nullptr)
       {
-        delete pointer;
-        delete ref_count;
+        if (--(control->shared_count) == 0)
+        {
+          delete pointer;
+          pointer = nullptr;
+
+          if (control->weak_count == 0)
+          {
+            delete control;
+          }
+        }
       }
     }
 
-    T * pointer;                           ///< The managed object.
-    std::atomic<std::size_t> * ref_count;  ///< Reference count.
+    T * pointer;                     ///< The managed object.
+    shared_control_block * control;  ///< Shared control block.
   };
 
   /**