memory.hpp

/*
    Copied from
   https://github.com/libhal/libhal/blob/main/include/libhal/pointers.hpp
 
    [atlas] modifications
    1.) File renamed to memory
    2.) In the atlas:: namespace.
*/
#include <atomic>
#include <memory_resource>
#include <type_traits>
#include <utility>
#include <array>
#include <atomic>
#include <optional>
#include <cstdint>
 
namespace atlas::memory {
    // Forward declarations
    template<typename T>
    class strong_ptr;
 
    template<typename T>
    class weak_ptr;
 
    namespace detail {
        struct ref_info {
            using destroy_fn_t = size_t(void const*);
 
            std::pmr::polymorphic_allocator<> allocator;
            destroy_fn_t* destroy;
            std::atomic<int32_t> strong_count = 1;
            std::atomic<int32_t> weak_count = 0;
        };
 
        inline void ptr_add_ref(ref_info* p_info) {
            p_info->strong_count.fetch_add(1, std::memory_order_relaxed);
        }
 
        inline void ptr_release(ref_info* p_info) {
            if (p_info->strong_count.fetch_sub(1, std::memory_order_acq_rel) ==
                1) {
                // No more strong references, destroy the object but keep
                // control block if there are weak references
 
                // Call the destroy function which will:
                // 1. Call the destructor of the object
                // 2. Return the size of the rc for deallocation when needed
                size_t const object_size = p_info->destroy(p_info);
 
                // If there are no weak references, deallocate memory
                if (p_info->weak_count.load(std::memory_order_acquire) == 0) {
                    // Save allocator for deallocating
                    auto alloc = p_info->allocator;
 
                    // Deallocate memory
                    alloc.deallocate_bytes(p_info, object_size);
                }
            }
        }
 
        inline void ptr_add_weak(ref_info* p_info) {
            p_info->weak_count.fetch_add(1, std::memory_order_relaxed);
        }
 
        inline void ptr_release_weak(ref_info* p_info) {
            if (p_info->weak_count.fetch_sub(1, std::memory_order_acq_rel) ==
                0) {
                // No more weak references, check if we can deallocate
                if (p_info->strong_count.load(std::memory_order_acquire) == 0) {
                    // No strong references either, get the size from the
                    // destroy function
                    auto const object_size = p_info->destroy(nullptr);
 
                    // Save allocator for deallocating
                    auto alloc = p_info->allocator;
 
                    // Deallocate memory
                    alloc.deallocate_bytes(p_info, object_size);
                }
            }
        }
 
        template<typename T>
        struct rc {
            ref_info m_info;
            T m_object;
 
            // Constructor that forwards arguments to the object
            template<typename... Args>
            rc(std::pmr::polymorphic_allocator<> p_alloc, Args&&... args)
              : m_info{ .allocator = p_alloc, .destroy = &destroy_function }
              , m_object(std::forward<Args>(args)...) {}
 
            // Static function to destroy an instance and return its size
            static size_t destroy_function(void const* p_object) {
                if (p_object != nullptr) {
                    auto const* obj = static_cast<rc<T> const*>(p_object);
                    // Call destructor for the object only
                    obj->m_object.~T();
                }
                // Return size for future deallocation
                return sizeof(rc<T>);
            }
        };
        // Check if a type is an array or std::array
        template<typename T>
        struct is_array_like : std::false_type {};
 
        // NOLINTBEGIN(modernize-avoid-c-arrays)
        // Specialization for C-style arrays
        template<typename T, std::size_t N>
        struct is_array_like<T[N]> : std::true_type {};
        // NOLINTEND(modernize-avoid-c-arrays)
 
        // Specialization for std::array
        template<typename T, std::size_t N>
        struct is_array_like<std::array<T, N>> : std::true_type {};
 
        // Helper variable template
        template<typename T>
        inline constexpr bool is_array_like_v = is_array_like<T>::value;
 
        // Concept for array-like types
        template<typename T>
        concept array_like = is_array_like_v<T>;
 
        // Concept for non-array-like types
        template<typename T>
        concept non_array_like = !array_like<T>;
    } // namespace detail
 
    template<typename T>
    class strong_ptr {
    public:
        using element_type = T;
 
        strong_ptr() = delete;
 
        strong_ptr(std::nullptr_t) = delete;
 
        strong_ptr(strong_ptr const& p_other) noexcept
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(p_other.m_ptr) {
            ptr_add_ref(m_ctrl);
        }
 
        template<typename U>
        strong_ptr(strong_ptr<U> const& p_other) noexcept
            requires(std::is_convertible_v<U*, T*>)
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(p_other.m_ptr) {
            ptr_add_ref(m_ctrl);
        }
 
        strong_ptr(strong_ptr&& p_other) noexcept
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(p_other.m_ptr) {
            ptr_add_ref(m_ctrl);
        }
 
        strong_ptr& operator=(strong_ptr&& p_other) noexcept {
            if (this != &p_other) {
                release();
                m_ctrl = p_other.m_ctrl;
                m_ptr = p_other.m_ptr;
                ptr_add_ref(m_ctrl);
            }
            return *this;
        }
 
        template<typename U>
        strong_ptr(strong_ptr<U> const&, void const*) noexcept {
            // NOTE: The conditional used here is to prevent the compiler from
            // jumping-the-gun and emitting the static assert error during
            // template instantiation of the class. With this conditional, the
            // error only appears when this constructor is used.
            static_assert(
              std::is_same_v<U, void> && !std::is_same_v<U, void>,
              "Aliasing constructor only works with pointers-to-members "
              "and does not work with arbitrary pointers like std::shared_ptr "
              "allows.");
        }
 
        template<typename U, detail::non_array_like M>
        strong_ptr(strong_ptr<U> const& p_other,
                   // clang-format off
             M U::* p_member_ptr
                   // clang-format on
                   ) noexcept
 
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(&((*p_other).*p_member_ptr)) {
            ptr_add_ref(m_ctrl);
        }
 
        template<typename U, typename E, std::size_t N>
        strong_ptr(strong_ptr<U> const& p_other,
                   // clang-format off
             std::array<E, N> U::* p_array_ptr,
                   // clang-format on
                   std::size_t p_index) {
            static_assert(std::is_convertible_v<E*, T*>,
                          "Array element type must be convertible to T");
            throw_if_out_of_bounds(N, p_index);
            m_ctrl = p_other.m_ctrl;
            m_ptr = &((*p_other).*p_array_ptr)[p_index];
            ptr_add_ref(m_ctrl);
        }
 
        // NOLINTBEGIN(modernize-avoid-c-arrays)
        template<typename U, typename E, std::size_t N>
        strong_ptr(strong_ptr<U> const& p_other,
                   E (U::*p_array_ptr)[N],
                   std::size_t p_index) {
            static_assert(std::is_convertible_v<E*, T*>,
                          "Array element type must be convertible to T");
            throw_if_out_of_bounds(N, p_index);
            m_ctrl = p_other.m_ctrl;
            m_ptr = &((*p_other).*p_array_ptr)[p_index];
            ptr_add_ref(m_ctrl);
        }
        // NOLINTEND(modernize-avoid-c-arrays)
 
        ~strong_ptr() { release(); }
 
        strong_ptr& operator=(strong_ptr const& p_other) noexcept {
            if (this != &p_other) {
                release();
                m_ctrl = p_other.m_ctrl;
                m_ptr = p_other.m_ptr;
                ptr_add_ref(m_ctrl);
            }
            return *this;
        }
 
        template<typename U>
        strong_ptr& operator=(strong_ptr<U> const& p_other) noexcept
            requires(std::is_convertible_v<U*, T*>)
        {
            release();
            m_ctrl = p_other.m_ctrl;
            m_ptr = p_other.m_ptr;
            ptr_add_ref(m_ctrl);
            return *this;
        }
 
        void swap(strong_ptr& p_other) noexcept {
            std::swap(m_ctrl, p_other.m_ctrl);
            std::swap(m_ptr, p_other.m_ptr);
        }
 
        T& operator*() && = delete;
 
        T* operator->() && = delete;
 
        [[nodiscard]] T& operator*() const& noexcept { return *m_ptr; }
 
        [[nodiscard]] T* operator->() const& noexcept { return m_ptr; }
 
        [[nodiscard]] auto use_count() const noexcept {
            return m_ctrl ? m_ctrl->strong_count.load(std::memory_order_relaxed)
                          : 0;
        }
 
    private:
        template<class U, typename... Args>
        friend strong_ptr<U> make_strong_ptr(std::pmr::polymorphic_allocator<>,
                                             Args&&...);
 
        template<typename U>
        friend class strong_ptr;
 
        template<typename U>
        friend class weak_ptr;
 
        template<typename U>
        friend class optional_ptr;
 
        inline void throw_if_out_of_bounds(size_t p_size, size_t p_index) {
            if (p_index >= p_size) {
                throw std::out_of_range(
                  this, { .m_index = p_index, .m_capacity = p_size });
            }
        }
 
        // Internal constructor with control block and pointer - used by make()
        // and aliasing
        strong_ptr(detail::ref_info* p_ctrl, T* p_ptr) noexcept
          : m_ctrl(p_ctrl)
          , m_ptr(p_ptr) {}
 
        void release() {
            if (m_ctrl) {
                ptr_release(m_ctrl);
            }
        }
 
        detail::ref_info* m_ctrl = nullptr;
        T* m_ptr = nullptr;
    };
 
    template<typename T>
    class enable_strong_from_this {
    public:
        [[nodiscard]] strong_ptr<T> strong_from_this() {
            auto locked = m_weak_this.lock();
            if (!locked) {
                throw std::bad_weak_ptr(&weak_from_this);
            }
            return locked.value();
        }
 
        [[nodiscard]] strong_ptr<T const> strong_from_this() const {
            auto locked = m_weak_this.lock();
            if (!locked) {
                throw std::bad_weak_ptr(&weak_from_this);
            }
            // Cast the strong_ptr<T> to strong_ptr<T const>
            return strong_ptr<T const>(locked.value());
        }
 
        [[nodiscard]] weak_ptr<T> weak_from_this() noexcept {
            return m_weak_this;
        }
 
        [[nodiscard]] weak_ptr<T const> weak_from_this() const noexcept {
            return weak_ptr<T const>(m_weak_this);
        }
 
    protected:
        enable_strong_from_this() = default;
 
        enable_strong_from_this(enable_strong_from_this const&) noexcept {
            // Intentionally don't copy m_weak_this
        }
 
        enable_strong_from_this& operator=(
          enable_strong_from_this const&) noexcept {
            // Intentionally don't assign m_weak_this
            return *this;
        }
 
        ~enable_strong_from_this() = default;
 
    private:
        template<class U, typename... Args>
        friend strong_ptr<U> make_strong_ptr(std::pmr::polymorphic_allocator<>,
                                             Args&&...);
 
        void init_weak_this(strong_ptr<T> const& p_self) noexcept {
            m_weak_this = p_self;
        }
 
        mutable weak_ptr<T> m_weak_this;
    };
 
    template<typename T>
    class optional_ptr;
 
    template<typename T>
    class weak_ptr {
    public:
        template<typename U>
        friend class strong_ptr;
 
        template<typename U>
        friend class weak_ptr;
 
        using element_type = T;
 
        weak_ptr() noexcept = default;
 
        weak_ptr(strong_ptr<T> const& p_strong) noexcept
          : m_ctrl(p_strong.m_ctrl)
          , m_ptr(p_strong.m_ptr) {
            if (m_ctrl) {
                ptr_add_weak(m_ctrl);
            }
        }
 
        weak_ptr(weak_ptr const& p_other) noexcept
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(p_other.m_ptr) {
            if (m_ctrl) {
                ptr_add_weak(m_ctrl);
            }
        }
 
        weak_ptr(weak_ptr&& p_other) noexcept
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(p_other.m_ptr) {
            p_other.m_ctrl = nullptr;
            p_other.m_ptr = nullptr;
        }
 
        template<typename U>
        weak_ptr(weak_ptr<U> const& p_other) noexcept
            requires(std::is_convertible_v<U*, T*>)
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(static_cast<T*>(p_other.m_ptr)) {
            if (m_ctrl) {
                ptr_add_weak(m_ctrl);
            }
        }
 
        template<typename U>
        weak_ptr(weak_ptr<U>&& p_other) noexcept
            requires(std::is_convertible_v<U*, T*>)
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(static_cast<T*>(p_other.m_ptr)) {
            p_other.m_ctrl = nullptr;
            p_other.m_ptr = nullptr;
        }
 
        template<typename U>
        weak_ptr(strong_ptr<U> const& p_other) noexcept
            requires(std::is_convertible_v<U*, T*>)
          : m_ctrl(p_other.m_ctrl)
          , m_ptr(static_cast<T*>(p_other.m_ptr)) {
            if (m_ctrl) {
                ptr_add_weak(m_ctrl);
            }
        }
 
        ~weak_ptr() {
            if (m_ctrl) {
                ptr_release_weak(m_ctrl);
            }
        }
 
        weak_ptr& operator=(weak_ptr const& p_other) noexcept {
            weak_ptr(p_other).swap(*this);
            return *this;
        }
 
        weak_ptr& operator=(weak_ptr&& p_other) noexcept {
            weak_ptr(std::move(p_other)).swap(*this);
            return *this;
        }
 
        weak_ptr& operator=(strong_ptr<T> const& p_strong) noexcept {
            weak_ptr(p_strong).swap(*this);
            return *this;
        }
 
        void swap(weak_ptr& p_other) noexcept {
            std::swap(m_ctrl, p_other.m_ctrl);
            std::swap(m_ptr, p_other.m_ptr);
        }
 
        [[nodiscard]] bool expired() const noexcept {
            return !m_ctrl ||
                   m_ctrl->strong_count.load(std::memory_order_relaxed) == 0;
        }
 
        [[nodiscard]] optional_ptr<T> lock() const noexcept;
 
        [[nodiscard]] auto use_count() const noexcept {
            return m_ctrl ? m_ctrl->strong_count.load(std::memory_order_relaxed)
                          : 0;
        }
 
    private:
        detail::ref_info* m_ctrl = nullptr;
        T* m_ptr = nullptr;
    };
 
    template<typename T>
    class optional_ptr {
    public:
        constexpr optional_ptr() noexcept {}
 
        constexpr optional_ptr(std::nullptr_t) noexcept {}
 
        constexpr optional_ptr(optional_ptr&& p_other) noexcept = delete;
 
        constexpr optional_ptr(strong_ptr<T> const& value) noexcept
          : m_value(value) {}
 
        constexpr optional_ptr(optional_ptr const& p_other) { *this = p_other; }
 
        template<typename U>
        constexpr optional_ptr(strong_ptr<U> const& p_value)
            requires(std::is_convertible_v<U*, T*>)
        {
            *this = p_value;
        }
 
        constexpr optional_ptr& operator=(optional_ptr&& other) noexcept =
          delete;
 
        constexpr optional_ptr& operator=(optional_ptr const& other) {
            if (this != &other) {
                if (is_engaged() && other.is_engaged()) {
                    m_value = other.m_value;
                }
                else if (is_engaged() && not other.is_engaged()) {
                    reset();
                }
                else if (not is_engaged() && other.is_engaged()) {
                    new (&m_value) strong_ptr<T>(other.m_value);
                }
            }
            return *this;
        }
 
        constexpr optional_ptr& operator=(strong_ptr<T> const& value) noexcept {
            if (is_engaged()) {
                m_value = value;
            }
            else {
                new (&m_value) strong_ptr<T>(value);
            }
            return *this;
        }
 
        template<typename U>
        constexpr optional_ptr& operator=(strong_ptr<U> const& p_value) noexcept
            requires(std::is_convertible_v<U*, T*>)
        {
            if (is_engaged()) {
                m_value = p_value;
            }
            else {
                new (&m_value) strong_ptr<U>(p_value);
            }
            return *this;
        }
 
        constexpr optional_ptr& operator=(std::nullptr_t) noexcept {
            reset();
            return *this;
        }
 
        ~optional_ptr() {
            if (is_engaged()) {
                m_value.~strong_ptr<T>();
            }
        }
 
        [[nodiscard]] constexpr bool has_value() const noexcept {
            return is_engaged();
        }
 
        constexpr explicit operator bool() const noexcept {
            return is_engaged();
        }
 
        [[nodiscard]] constexpr strong_ptr<T>& value() {
            if (!is_engaged()) {
                throw std::bad_optional_access();
            }
            return m_value;
        }
 
        [[nodiscard]] constexpr strong_ptr<T> const& value() const {
            if (!is_engaged()) {
                throw std::bad_optional_access();
            }
            return m_value;
        }
 
        [[nodiscard]] constexpr operator strong_ptr<T>() { return value(); }
 
        [[nodiscard]] constexpr operator strong_ptr<T>() const {
            return value();
        }
 
        template<typename U>
        [[nodiscard]] constexpr operator strong_ptr<U>()
            requires(std::is_convertible_v<T*, U*> && !std::is_same_v<T, U>)
        {
            if (!is_engaged()) {
                throw std::bad_optional_access();
            }
            // strong_ptr handles the polymorphic conversion
            return strong_ptr<U>(m_value);
        }
 
        template<typename U>
        [[nodiscard]] constexpr operator strong_ptr<U>() const
            requires(std::is_convertible_v<T*, U*> && !std::is_same_v<T, U>)
        {
            if (!is_engaged()) {
                // hal::safe_throw(hal::bad_optional_ptr_access(this));
                throw std::bad_optional_access();
            }
            // strong_ptr handles the polymorphic conversion
            return strong_ptr<U>(m_value);
        }
 
        [[nodiscard]] constexpr auto* operator->() {
            auto& ref = *(this->value());
            return &ref;
        }
 
        [[nodiscard]] constexpr auto* operator->() const {
            auto& ref = *(this->value());
            return &ref;
        }
 
        [[nodiscard]] constexpr auto& operator*() {
            auto& ref = *(this->value());
            return ref;
        }
 
        [[nodiscard]] constexpr auto& operator*() const {
            auto& ref = *(this->value());
            return ref;
        }
 
        constexpr void reset() noexcept {
            if (is_engaged()) {
                m_value.~strong_ptr<T>();
                m_raw_ptrs[0] = nullptr;
                m_raw_ptrs[1] = nullptr;
            }
        }
 
        template<typename... Args>
        constexpr strong_ptr<T>& emplace(Args&&... args) {
            reset();
            new (&m_value) strong_ptr<T>(std::forward<Args>(args)...);
            return m_value;
        }
 
        constexpr void swap(optional_ptr& other) noexcept {
            if (is_engaged() && other.is_engaged()) {
                std::swap(m_value, other.m_value);
            }
            else if (is_engaged() && !other.is_engaged()) {
                new (&other.m_value) strong_ptr<T>(std::move(m_value));
                reset();
            }
            else if (!is_engaged() && other.is_engaged()) {
                new (&m_value) strong_ptr<T>(std::move(other.m_value));
                other.reset();
            }
        }
 
    private:
        union {
            std::array<void*, 2> m_raw_ptrs = { nullptr, nullptr };
            strong_ptr<T> m_value;
        };
 
        // Ensure the strong_ptr layout matches our expectations
        static_assert(sizeof(m_value) == sizeof(m_raw_ptrs),
                      "strong_ptr must be exactly the size of two pointers");
 
        [[nodiscard]] constexpr bool is_engaged() const noexcept {
            return m_raw_ptrs[0] != nullptr || m_raw_ptrs[1] != nullptr;
        }
    };
 
    template<typename T>
    [[nodiscard]] inline optional_ptr<T> weak_ptr<T>::lock() const noexcept {
        if (expired()) {
            return nullptr;
        }
 
        // Try to increment the strong count
        auto current_count =
          m_ctrl->strong_count.load(std::memory_order_relaxed);
        while (current_count > 0) {
            // TODO(kammce): Consider if this is dangerous
            if (m_ctrl->strong_count.compare_exchange_weak(
                  current_count,
                  current_count + 1,
                  std::memory_order_acq_rel,
                  std::memory_order_relaxed)) {
                // Successfully incremented
                auto obj = strong_ptr<T>(m_ctrl, m_ptr);
                return obj;
            }
        }
 
        // Strong count is now 0
        return nullptr;
    }
 
    template<typename T>
    void swap(strong_ptr<T>& p_lhs, strong_ptr<T>& p_rhs) noexcept {
        p_lhs.swap(p_rhs);
    }
 
    template<typename T>
    void swap(weak_ptr<T>& p_lhs, weak_ptr<T>& p_rhs) noexcept {
        p_lhs.swap(p_rhs);
    }
 
    template<typename T, typename U>
    bool operator==(strong_ptr<T> const& p_lhs,
                    strong_ptr<U> const& p_rhs) noexcept {
        return p_lhs.operator->() == p_rhs.operator->();
    }
 
    template<typename T, typename U>
    bool operator!=(strong_ptr<T> const& p_lhs,
                    strong_ptr<U> const& p_rhs) noexcept {
        return !(p_lhs == p_rhs);
    }
    template<typename T, typename U>
    [[nodiscard]] bool operator==(optional_ptr<T> const& p_lhs,
                                  optional_ptr<U> const& p_rhs) noexcept {
        // If both are disengaged, they're equal
        if (!p_lhs.has_value() && !p_rhs.has_value()) {
            return true;
        }
 
        // If one is engaged and the other isn't, they're not equal
        if (p_lhs.has_value() != p_rhs.has_value()) {
            return false;
        }
 
        // Both are engaged, compare the underlying pointers
        return p_lhs.value().operator->() == p_rhs.value().operator->();
    }
 
    template<typename T, typename U>
    [[nodiscard]] bool operator!=(optional_ptr<T> const& p_lhs,
                                  optional_ptr<U> const& p_rhs) noexcept {
        return !(p_lhs == p_rhs);
    }
 
    template<typename T>
    [[nodiscard]] bool operator==(optional_ptr<T> const& p_lhs,
                                  std::nullptr_t) noexcept {
        return !p_lhs.has_value();
    }
 
    template<typename T>
    [[nodiscard]] bool operator==(std::nullptr_t,
                                  optional_ptr<T> const& p_rhs) noexcept {
        return !p_rhs.has_value();
    }
 
    template<typename T>
    [[nodiscard]] bool operator!=(optional_ptr<T> const& p_lhs,
                                  std::nullptr_t) noexcept {
        return p_lhs.has_value();
    }
 
    template<typename T>
    [[nodiscard]] bool operator!=(std::nullptr_t,
                                  optional_ptr<T> const& p_rhs) noexcept {
        return p_rhs.has_value();
    }
 
    class strong_ptr_only_token {
    private:
        strong_ptr_only_token() = default;
 
        template<class U, typename... Args>
        friend strong_ptr<U> make_strong_ptr(std::pmr::polymorphic_allocator<>,
                                             Args&&...);
    };
 
    template<class T, typename... Args>
    [[nodiscard]] inline strong_ptr<T> make_strong_ptr(
      std::pmr::polymorphic_allocator<> p_alloc,
      Args&&... p_args) {
        using rc_t = detail::rc<T>;
 
        rc_t* obj = nullptr;
 
        if constexpr (std::
                        is_constructible_v<T, strong_ptr_only_token, Args...>) {
            // Type expects token as first parameter
            obj = p_alloc.new_object<rc_t>(
              p_alloc, strong_ptr_only_token{}, std::forward<Args>(p_args)...);
        }
        else {
            // Normal type, construct without token
            obj =
              p_alloc.new_object<rc_t>(p_alloc, std::forward<Args>(p_args)...);
        }
 
        strong_ptr<T> result(&obj->m_info, &obj->m_object);
 
        // Initialize enable_strong_from_this if the type inherits from it
        if constexpr (std::is_base_of_v<enable_strong_from_this<T>, T>) {
            result->init_weak_this(result);
        }
 
        return result;
    }
} // namespace atlas::memory