AronaCore/core/misc/ref_counting.h

#pragma once
#include <assert.h>

/**
 * A smart pointer to an object which implements AddRef/Release.
 */
template<typename ReferencedType>
class ref_count_ptr {
    typedef ReferencedType* reference_type;

public:
    ref_count_ptr(): reference_(nullptr) {
    }

    explicit ref_count_ptr(ReferencedType* in_reference, bool bAddRef = true) {
        reference_ = in_reference;
        if (reference_ && bAddRef) {
            reference_->AddRef();
        }
    }

    ref_count_ptr(const ref_count_ptr& copy) {
        reference_ = copy.reference_;
        if (reference_) {
            reference_->AddRef();
        }
    }

    template<typename CopyReferencedType>
    explicit ref_count_ptr(const ref_count_ptr<CopyReferencedType>& copy) {
        reference_ = static_cast<ReferencedType*>(copy.get_reference());
        if (reference_) {
            reference_->AddRef();
        }
    }

    ref_count_ptr(ref_count_ptr&& move) noexcept {
        reference_ = move.reference_;
        move.reference_ = nullptr;
    }

    template<typename MoveReferencedType>
    explicit ref_count_ptr(ref_count_ptr<MoveReferencedType>&& move) {
        reference_ = static_cast<ReferencedType*>(move.get_reference());
        move.reference_ = nullptr;
    }

    ~ref_count_ptr() {
        if (reference_) {
            reference_->Release();
        }
    }

    ref_count_ptr& operator=(ReferencedType* in_reference) {
        if (reference_ != in_reference) {
            // Call AddRef before Release, in case the new reference is the same as the old reference.
            ReferencedType* old_reference = reference_;
            reference_ = in_reference;
            if (reference_) {
                reference_->AddRef();
            }
            if (old_reference) {
                old_reference->Release();
            }
        }
        return *this;
    }

    ref_count_ptr& operator=(const ref_count_ptr& in_ptr) {
        return *this = in_ptr.reference_;
    }

    template<typename CopyReferencedType>
    ref_count_ptr& operator=(const ref_count_ptr<CopyReferencedType>& in_ptr) {
        return *this = in_ptr.GetReference();
    }

    ref_count_ptr& operator=(ref_count_ptr&& in_ptr) noexcept {
        if (this != &in_ptr) {
            ReferencedType* old_reference = reference_;
            reference_ = in_ptr.reference_;
            in_ptr.reference_ = nullptr;
            if (old_reference) {
                old_reference->Release();
            }
        }
        return *this;
    }

    template<typename MoveReferencedType>
    ref_count_ptr& operator=(ref_count_ptr<MoveReferencedType>&& in_ptr) {
        // InPtr is a different type (or we would have called the other operator), so we need not test &InPtr != this
        ReferencedType* old_reference = reference_;
        reference_ = in_ptr.reference_;
        in_ptr.reference_ = nullptr;
        if (old_reference) {
            old_reference->Release();
        }
        return *this;
    }

    ReferencedType* operator->() const {
        return reference_;
    }

    operator reference_type() const {
        return reference_;
    }

    ReferencedType** get_init_reference() {
        *this = nullptr;
        return &reference_;
    }

    ReferencedType* get_reference() const {
        return reference_;
    }

    friend bool is_valid_ref(const ref_count_ptr& in_reference) {
        return in_reference.reference_ != nullptr;
    }

    bool is_valid() const {
        return reference_ != nullptr;
    }

    void safe_release() {
        *this = nullptr;
    }

    unsigned int get_ref_count() {
        unsigned int result = 0;
        if (reference_) {
            result = reference_->GetRefCount();
            assert(result > 0);
            // you should never have a zero ref count if there is a live ref counted pointer (*this is live)
        }
        return result;
    }

    void swap(ref_count_ptr& in_ptr) noexcept // this does not change the reference count, and so is faster
    {
        ReferencedType* old_reference = reference_;
        reference_ = in_ptr.reference_;
        in_ptr.reference_ = old_reference;
    }

    // void Serialize(FArchive& Ar)
    // {
    // 	reference_type PtrReference = Reference;
    // 	Ar << PtrReference;
    // 	if(Ar.IsLoading())
    // 	{
    // 		*this = PtrReference;
    // 	}
    // }

private:
    ReferencedType* reference_;

    template<typename OtherType>
    friend class ref_count_ptr;

public:
    bool operator==(const ref_count_ptr& b) const {
        return get_reference() == b.get_reference();
    }

    bool operator==(ReferencedType* b) const {
        return get_reference() == b;
    }
};