ReferenceCounter.h

#pragma once
#include <UnCompute/Base/Base.h>
#include <UnCompute/Memory/IAllocator.h>
 
namespace UN
{
    class IObject;
 
    class ReferenceCounter final
    {
        std::atomic<Int32> m_StrongRefCount;
        mutable IAllocator* m_pAllocator;
 
    public:
        inline explicit ReferenceCounter(IAllocator* pAllocator)
            : m_StrongRefCount(0)
            , m_pAllocator(pAllocator)
        {
        }
 
        inline UInt32 AddStrongRef()
        {
            return ++m_StrongRefCount;
        }
 
        template<class F>
        inline UInt32 ReleaseStrongRef(F&& destroyCallback)
        {
            if (--m_StrongRefCount == 0)
            {
                destroyCallback();
                m_pAllocator->Deallocate(this);
            }
 
            return m_StrongRefCount;
        }
 
        inline UInt32 GetStrongRefCount() const
        {
            return m_StrongRefCount;
        }
    };
} // namespace UN