MemoryPool.h

#ifndef GX_MEMORYPOOL_H
#define GX_MEMORYPOOL_H

#include "Thread.h"

#define GX_ENABLE_MEMORYPOOL_STATS

namespace GX
{
    template<typename ValueType, typename ...ArgTypes >
    ValueType* original_new(ArgTypes ... args)
    {
        ValueType* result = (ValueType*) ::malloc(sizeof(ValueType));
        result = new (result) ValueType(args ...);
        return result;
    }

    template<typename ValueType>
    void original_delete(ValueType* pointer)
    {
        pointer->~ValueType();
        ::free(pointer);
    }

    template<U32 DataSize, U32 _ThunkSize>
    class MemoryPoolImpl
    {
    public:
        enum
        {
            ThunkSize = _ThunkSize
        };

#ifdef COCOS2D_DEBUG
        class Thunk;
        struct
#else
        union
#endif
        ThunkNode
        {
            U8  m_buffer[DataSize];
            U32 m_nextSlot;
#ifdef COCOS2D_DEBUG
            Thunk* m_pOwner;
#endif
        };

        struct Thunk
        {
        public:
            Thunk()
            : m_head(0)
            , m_pNext(nullptr)
#ifdef GX_ENABLE_MEMORYPOOL_STATS
            , m_used(0)
#endif
            {
                for(U32 i = 0; i < ThunkSize; i++)
                {
                    m_data[i].m_nextSlot = i < ThunkSize - 1 ? i + 1 : ~0;  // Consider the right sibling as the next valid slot
#ifdef COCOS2D_DEBUG
                    m_data[i].m_pOwner = this;
#endif
                }
            }

            ThunkNode* acquire()
            {
                ThunkNode* pResult = nullptr;
                
                if (m_head != ~0)
                {
                    GX_ASSERT(m_head < ThunkSize, "Memory pool thunk ruined");
                    if (m_head < ThunkSize)
                    {
                        pResult = m_data + m_head;
                        // Move head
                        m_head = pResult->m_nextSlot;
#ifdef GX_ENABLE_MEMORYPOOL_STATS
                        m_used++;
#endif
                    }
                }

                return pResult;
            }

            void release(ThunkNode* pNode)
            {
                if (pNode != nullptr)
                {
                    pNode->m_nextSlot = m_head;
                    m_head = pNode - m_data;
#ifdef GX_ENABLE_MEMORYPOOL_STATS
                    m_used--;
#endif
                }
            }

            void release(U8* pValue)
            {
                release(reinterpret_cast<ThunkNode*>(pValue));
            }

            bool owns(ThunkNode* pNode)
            {
                return pNode >= m_data && pNode < m_data + ThunkSize;
            }

#ifdef GX_ENABLE_MEMORYPOOL_STATS
            U32 used()
            {
                return m_used;
            }
#endif

        public:
            ThunkNode m_data[ThunkSize];
            Thunk* m_pNext;
            U32 m_head;
#ifdef GX_ENABLE_MEMORYPOOL_STATS
            U32 m_used;
#endif
        };

    public:
        MemoryPoolImpl()
        : m_pLastThunk(nullptr)
        , m_pHeadThunk(nullptr)
        {}

        U8* acquire()
        {
            CriticalSectionScopeLocker lock(m_cs);

            U8* pResult = nullptr;

            ThunkNode* pNode = nullptr;

            if (m_pLastThunk != nullptr)
            {
                pNode = m_pLastThunk->acquire();
            }

            if (pNode == nullptr)
            {
                Thunk* pIter = m_pHeadThunk;
                while(pIter != nullptr)
                {
                    if (m_pLastThunk != pIter)
                    {
                        pNode = pIter->acquire();
                        if (pNode != nullptr)
                        {
                            m_pLastThunk = pIter;
                            break;
                        }
                    }
                    pIter = pIter->m_pNext;
                }
            }

            if (pNode == nullptr)
            {
                m_pLastThunk = original_new<Thunk>();
                if (m_pLastThunk != nullptr)
                {
                    pNode = m_pLastThunk->acquire();

                    if (m_pHeadThunk == nullptr)
                    {
                        m_pHeadThunk = m_pLastThunk;
                    }
                    else
                    {
                        Thunk* pIter = m_pHeadThunk;
                        while (pIter->m_pNext != nullptr)
                        {
                            pIter = pIter->m_pNext;
                        }
                        pIter->m_pNext = m_pLastThunk;
                    }
                }
            }

            if (pNode != nullptr)
            {
                pResult = pNode->m_buffer;
            }

            GX_ASSERT(pResult != nullptr, "Memory alloc failed");

            return pResult;
        }

        bool release(ThunkNode* pNode)
        {
            CriticalSectionScopeLocker lock(m_cs);

            bool succeed = false;
            if (pNode != nullptr)
            {
                Thunk* pOwner = nullptr;
                Thunk* pOwnerPrev = nullptr;
                Thunk* pIter = m_pHeadThunk;

                while(pIter != nullptr)
                {
                    pOwner = pIter;

                    if (pOwner->owns(pNode))
                    {
                        pOwner->release(pNode);
                        succeed = true;
                        break;
                    }
                    else
                    {
                        pOwnerPrev = pIter;
                        pIter = pIter->m_pNext;
                    }
                }

#ifdef GX_ENABLE_MEMORYPOOL_STATS
                if (succeed && pOwner->used() == 0)
                {
                    if (pOwnerPrev != nullptr)
                    {
                        pOwnerPrev->m_pNext = pOwner->m_pNext;
                    }
                    else
                    {
                        pOwnerPrev = m_pHeadThunk = pOwner->m_pNext;
                    }

                    if (m_pLastThunk == pOwner)
                    {
                        m_pLastThunk = pOwnerPrev;
                    }

                    original_delete(pOwner);
                }
#endif
            }

            return succeed;
        }

        bool release(U8* pValue)
        {
            return release(reinterpret_cast<ThunkNode*>(pValue));
        }

#ifdef GX_ENABLE_MEMORYPOOL_STATS
        void dumpStats()
        {
            CriticalSectionScopeLocker lock(m_cs);
            
            GX_LOG("MemoryPool<{:d}, {:d}>", DataSize, ThunkSize);
            U32 thunkCount = 0;

            Thunk* pIter = m_pHeadThunk;
            while(pIter != nullptr)
            {
                GX_LOG("  Thunk #{:02d}: Used: {:d} ({:.2f}%)", thunkCount, pIter->used(), (F32)pIter->used() / ThunkSize * 100);
                pIter = pIter->m_pNext;
                thunkCount++;
            }
            GX_LOG("  Total: {:d} KB, {:d} MB", (thunkCount * DataSize * ThunkSize) >> 10, (thunkCount * DataSize * ThunkSize) >> 20);
        }
#endif

    private:
        Thunk*  m_pHeadThunk;
        Thunk*  m_pLastThunk;
        CriticalSection m_cs;
    };

    template<int DataSize, int ThunkSize = (0x80 << 10) / DataSize>
    struct MemoryPool : public MemoryPoolImpl<DataSize, ThunkSize>
    {
        static MemoryPool& shared_pool()
        {
            static MemoryPool s_pool;
            return s_pool;
        }
    };

    class MemoryPoolAllocator
    {
    public:
        template<typename DataType, typename ... ArgTypes>
        static DataType* alloc(ArgTypes ... args)
        {
            DataType* pResult = reinterpret_cast<DataType*>(allocImpl(Meta::NextPOT<sizeof(DataType)>::Result));
            new (pResult) DataType(args ...);
            return pResult;
        }

        template<typename DataType>
        static void free(DataType* pValue)
        {
            if (pValue != nullptr)
            {
                pValue->~DataType();
                freeImpl(reinterpret_cast<U8*>(pValue), Meta::NextPOT<sizeof(DataType)>::Result);
            }
        }

    public:
        static U8* allocImpl(U32 size);
        static void freeImpl(U8* pData, U32 assumeSize);
        static void dumpStats();
    };

//#define GX_MemoryPoolNew(DataType, ...) GX::MemoryPoolAllocator::alloc<DataType>(__VA_ARGS__)
//#define GX_MemoryPoolDelete(Pointer)    GX::MemoryPoolAllocator::free<typename GX::Meta::RemoveAll<decltype(Pointer)>::Result>(Pointer)
}

#endif