path: root/libs/kstd/include/kstd/bits/shared_ptr.hpp

#ifndef KSTD_BITS_SHARED_PTR_HPP
#define KSTD_BITS_SHARED_PTR_HPP

#include <atomic>
#include <cstddef>
#include <type_traits>
#include <utility>

// IWYU pragma: private, include <kstd/memory>

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
   * shared_ptr owning the object is assigned another pointer via operator= or reset(). A
   * shared_ptr can share ownership of an object while storing a pointer to another object. This feature can be used
   * to point to member objects while owning the object they belong to. The stored pointer is the one accessed by get(),
   * the dereference and the comparison operators. The managed pointer is the one passed to the deleter when use count
   * reaches zero.
   *
   * @tparam T The type of the managed object.
   */
  template<typename T>
  struct shared_ptr
  {
    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)
        , control(nullptr)
    {}

    /**
     * @brief Constructor.
     *
     * @param pointer A pointer to an object to manage (default is nullptr).
     */
    template<typename U>
      requires(std::is_convertible_v<U *, T *>)
    explicit shared_ptr(U * pointer = nullptr)
        : pointer(pointer)
        , 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)