ObjectPool.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include <type_traits>
#include "HAL/UnrealMemory.h"
#include "Math/UnrealMathUtility.h"
#include "Templates/AlignmentTemplates.h"
#include "Templates/UnrealTemplate.h"
 
namespace Chaos
{
    // SFINAE helpers to clean up the appearance in the class definition.
 
    // The type needs to have a destructor run when the pool pointer is freed
    template<typename T>
    using TRequiresDestructor = std::enable_if_t<!std::is_trivially_destructible_v<T>>;
 
    // The type needs no destruction and can just be abandoned on free
    template<typename T>
    using TTrivialDestruct = std::enable_if_t<std::is_trivially_destructible_v<T>>;
 
    template<typename ObjectType>
    class TObjectPool
    {
    public:
 
        // Alignment of the stored type
        constexpr static int32 ItemAlign = alignof(ObjectType);
 
        using FObject = ObjectType;
        using FPtr = ObjectType*;
 
        explicit TObjectPool(int32 InNumPerBlock, int32 InitialBlocks = 1)
            : NumPerBlock(InNumPerBlock)
        {
            ReserveBlocks(InitialBlocks);
        }
 
        ~TObjectPool()
        {
            Reset();
        }
 
        // No copy, no move, no default
        // #TODO possibly add move and copy if the internal type is copy constructible but likely better
        // to discourage copying. Moving to another pool should be fine.
        TObjectPool() = delete;
        TObjectPool(const TObjectPool&) = delete;
        TObjectPool(TObjectPool&&) = delete;
        TObjectPool& operator=(const TObjectPool&) = delete;
        TObjectPool& operator=(TObjectPool&&) = delete;
 
        template<typename... TArgs>
        FPtr Alloc(TArgs&&... Args)
        {
            // Need a new Block
            if(FreeCount == 0)
            {
                // Need a new block
                ObjectType* NewPtr = AddBlock().GetNextFree();
                Construct<ObjectType>(NewPtr, Forward<TArgs>(Args)...);
                --FreeCount;
                return NewPtr;
            }
 
            // We know there's a free item somewhere, find one - shuffling full blocks to the end of the list
            check(Blocks.Num() > 0);
 
            if (Blocks[0].IsFull())
            {
                // Find the fullest block that is not completely full
                int32 SelectedIndex = INDEX_NONE;
                int32 SelectedNumFree = TNumericLimits<int32>::Max();
                for (int32 Index = 0; Index < Blocks.Num(); ++Index)
                {
                    const int32 BlockNumFree = Blocks[Index].NumFree;
                    if ((BlockNumFree > 0) && (BlockNumFree < SelectedNumFree))
                    {
                        SelectedNumFree = BlockNumFree;
                        SelectedIndex = Index;
                    }
                }
 
                // Move the selected block to the front
                checkf(SelectedIndex != INDEX_NONE, TEXT("Could not find an empty block"));
                Swap(Blocks[0], Blocks[SelectedIndex]);
            }
 
            // Ensure any writes to Blocks[0] cannot be reordered so that the GetNextFree
            // call is definitely handled by the right block
            FPlatformMisc::MemoryBarrier();
 
            // Blocks[0] is now a block with at least one free element
            ObjectType* NewPtr = Blocks[0].GetNextFree();
            Construct<ObjectType>(NewPtr, Forward<TArgs>(Args)...);
            --FreeCount;
            return NewPtr;
        }
 
        void Free(FPtr Object)
        {
            int32 BlockIndex = FindBlock(Object);
            Blocks[BlockIndex].Free(Object);
            ++FreeCount;
        }
 
        void Reset()
        {
            for(FBlock& B : Blocks)
            {
                B.Reset();
            }
 
            FreeCount = Blocks.Num() * NumPerBlock;
        }
 
        int32 GetNumAllocatedBlocks() const
        {
            return Blocks.Num();
        }
 
        int32 GetNumPerBlock() const
        {
            return NumPerBlock;
        }
 
        int32 GetAllocatedSize() const
        {
            return (PaddedItemSize * NumPerBlock) * Blocks.Num();
        }
 
        void ShrinkTo(int32 NumDesiredEmptyBlocks)
        {
            int32 CurrentNumEmpty = 0;
            int32 Index = 0;
            while(Index < Blocks.Num())
            {
                if(Blocks[Index].IsEmpty())
                {
                    if(CurrentNumEmpty < NumDesiredEmptyBlocks)
                    {
                        ++CurrentNumEmpty;
                    }
                    else
                    {
                        Blocks.RemoveAtSwap(Index);
                        continue;
                    }
                }
                ++Index;
            }
        }
 
        void ReserveBlocks(int32 NumBlocks)
        {
            while(NumBlocks > Blocks.Num())
            {
                AddBlock();
            }
        }
 
        void ReserveItems(int32 NumItems)
        {
            const int32 ItemsSize = NumItems * PaddedItemSize;
            const int32 BlockSize = NumPerBlock * PaddedItemSize;
            const int32 NumBlocks = int32(FMath::CeilToInt(static_cast<double>(ItemsSize) / BlockSize));
            ReserveBlocks(NumBlocks);
        }
 
        float GetRatio() const
        {
            return static_cast<float>(sizeof(ObjectType)) / PaddedItemSize;
        }
 
        int32 GetNumAllocated() const
        {
            return GetCapacity() - GetNumFree();
        }
 
        int32 GetNumFree() const
        {
            return FreeCount;
        }
 
        int32 GetCapacity() const
        {
            return Capacity;
        }
 
    private:
 
        // Storage for an item, and an index of the next free item if it's currently on the free list
        template<typename T_ = ObjectType, typename Destructible = void>
        struct alignas(ItemAlign) TItem
        {
            TItem()
                : NextFree(INDEX_NONE)
            {}
 
            // Object must always remain as the first item to enable casting between Item and Object
            ObjectType Object;
 
            // Free-list tracking, index of the next free item
            int32 NextFree;
 
            // Set whether the object is live - encapsulated in a function as the base template performs no action
            void SetLive(bool)
            {}
        };
 
        // Specialized type for types that require a destructor as we need to add an extra
        // flag for live/dead objects (calling Free will destruct the object but Reset/DestroyAll
        // will need to destruct any non-free objects). This is specialized so that if the type
        // is trivially destructible we can avoid adding the extra member.
        template<typename T_>
        struct alignas(ItemAlign) TItem<T_, TRequiresDestructor<T_>>
        {
            TItem()
                : NextFree(INDEX_NONE)
            {}
 
            // Object must always remain as the first item to enable casting between Item and Object
            ObjectType Object;
 
            // Free-list tracking, index of the next free item
            int32 NextFree;
 
            // Whether the object is live (constructed) and will require destruction
            bool bLive : 1;
 
            // Set whether the object is live - encapsulated in a function as the base template performs no action
            // @param bInLive - new live flag value
            void SetLive(bool bInLive)
            {
                bLive = bInLive;
            }
        };
 
        // Helper alias for the item in this pool
        using FItem = TItem<ObjectType>;
 
        // Size of the item, plus padding up to the correct alignment so we can allocate the whole
        // block with the correct size.
        constexpr static int32 PaddedItemSize = Align(int32(sizeof(FItem)), ItemAlign);
 
        // A block is an area of memory large enough to hold NumInBlock items, correctly aligned and
        // provide items on demand to the pool
        struct FBlock
        {
            FBlock() = delete;
 
            explicit FBlock(int32 InNum)
                : Begin(static_cast<FItem*>(FMemory::Malloc(PaddedItemSize * InNum, ItemAlign)))
                , Next(Begin)
                , NumInBlock(InNum)
                , NumFree(NumInBlock)
            {
            }
 
            ~FBlock()
            {
                check(Begin);
 
                // Return our memory block
                FMemory::Free(Begin);
            }
 
            FPtr GetNextFree()
            {
                if(FreeList != INDEX_NONE)
                {
                    // Something in the free list, prefer this
                    FItem* NewItem = Begin + FreeList;
                    FreeList = NewItem->NextFree;
                    NewItem->NextFree = INDEX_NONE;
                    --NumFree;
 
                    return &NewItem->Object;
                }
 
                if(NumValid < NumInBlock)
                {
                    FItem* NewItem = Next++;
                    NewItem->NextFree = INDEX_NONE;
                    NewItem->SetLive(false);
 
                    --NumFree;
                    ++NumValid;
 
                    return &NewItem->Object;
                }
 
                checkf(false, TEXT("Attempt to request a free item from a full block (Freelist is empty, NumValid is %d, NumInBlock is %d, NumFree is %d)"), NumValid, NumInBlock, NumFree);
 
                return nullptr;
            }
 
            void Free(ObjectType* Object)
            {
                checkf(Owns(Object), TEXT("A pointer was passed to an object pool block that it does not own."));
 
                // If required, call the object destructor
                ConditionalDestruct<ObjectType>(Object);
 
                // Get the offset into this block (check above ensures this is correct)
                FItem* const AsItem = reinterpret_cast<FItem*>(Object);
                const UPTRINT Offset = reinterpret_cast<UPTRINT>(AsItem) - reinterpret_cast<UPTRINT>(Begin);
                const UPTRINT IndexInBlock = Offset / PaddedItemSize;
 
                // If there's already a freelist, add it to our next item
                if(FreeList != INDEX_NONE)
                {
                    AsItem->NextFree = FreeList;
                }
 
                FreeList = (int32)IndexInBlock;
                ++NumFree;
            }
 
            void Reset()
            {
                DestructAll();
 
                Next = Begin;
                NumFree = NumInBlock;
                FreeList = INDEX_NONE;
                NumValid = 0;
            }
 
            void DestructAll()
            {
                for(int32 i = 0; i < NumValid; ++i)
                {
                    FItem* ToDestroy = Begin + i;
                    ConditionalDestruct<ObjectType>(&ToDestroy->Object);
                }
            }
 
            bool Owns(ObjectType* InPtr) const
            {
                FItem* AsItem = reinterpret_cast<FItem*>(InPtr);
 
                return AsItem >= Begin && AsItem < Begin + NumInBlock;
            }
 
            bool IsEmpty() const
            {
                return NumFree == NumInBlock;
            }
 
            bool IsFull() const
            {
                return NumFree == 0;
            }
 
            // Ptr to the beginning of the block
            FItem* Begin = nullptr;
            // Ptr to the next free item when block is filling initially
            FItem* Next = nullptr;
 
            // Number of items the block can hold
            const int32 NumInBlock;
            // Count of free items in the block
            int32 NumFree = 0;
            // The maximum number of valid items in the block (items are valid if they were ever allocated / constructed)
            int32 NumValid = 0;
            // Head of the free list for this block
            int32 FreeList = INDEX_NONE;
        };
 
        FBlock& AddBlock()
        {
            Blocks.Emplace(NumPerBlock);
 
            // Update pool capacity and free count
            Capacity += NumPerBlock;
            FreeCount += NumPerBlock;
 
            return Blocks.Last();
        }
 
        int32 FindBlock(ObjectType* InPtr)
        {
            for(int32 i = 0; i < Blocks.Num(); ++i)
            {
                if(Blocks[i].Owns(InPtr))
                {
                    return i;
                }
            }
 
            checkf(false, TEXT("A pointer was passed to an object pool method that it does not own."));
            return INDEX_NONE;
        }
 
        template<
            typename ObjectType_ = ObjectType, 
            TTrivialDestruct<ObjectType_>* = nullptr, 
            typename... TArgs>
        static ObjectType* Construct(ObjectType* At, TArgs&&... Args)
        {
            new (At) ObjectType(Forward<TArgs>(Args)...);
 
            return At;
        }
 
        template<
            typename ObjectType_ = ObjectType, 
            TRequiresDestructor<ObjectType_>* = nullptr, 
            typename... TArgs>
        static ObjectType* Construct(ObjectType* At, TArgs&&... Args)
        {
            new (At) ObjectType(Forward<TArgs>(Args)...);
 
            FItem* AsItem = reinterpret_cast<FItem*>(At);
            AsItem->bLive = true;
 
            return At;
        }
 
        template<
            typename ObjectType_ = ObjectType, 
            TRequiresDestructor<ObjectType_>* = nullptr>
        static void ConditionalDestruct(ObjectType* At)
        {
            checkSlow(At);
 
            // Check the live flag, items on the free list will already have been destructed
            // and we should avoid calling the destructor twice
            FItem* AsItem = reinterpret_cast<FItem*>(At);
            if(AsItem->bLive)
            {
                At->~ObjectType();
                AsItem->bLive = false;
            }
        }
 
        template<
            typename ObjectType_ = ObjectType, 
            TTrivialDestruct<ObjectType_>* = nullptr>
        static void ConditionalDestruct(ObjectType*)
        {}
 
        TArray<FBlock> Blocks;
 
        int32 NumPerBlock = 0;
        int32 Capacity = 0;
        int32 FreeCount = 0;
    };
 
 
    template<typename ObjectPoolType>
    class TObjectPoolDeleter
    {
    public:
        using FObjectPool = ObjectPoolType;
        using FObject = typename FObjectPool::FObject;
 
        TObjectPoolDeleter()
            : Pool(nullptr)
        {
        }
 
        TObjectPoolDeleter(FObjectPool& InPool)
            : Pool(&InPool)
        {
        }
 
        void operator()(FObject* Object)
        {
            if (Object != nullptr)
            {
                check(Pool != nullptr);
                Pool->Free(Object);
            }
        }
 
    private:
        FObjectPool* Pool;
    };
}
 
#if UE_ENABLE_INCLUDE_ORDER_DEPRECATED_IN_5_5
#include "Templates/IsTriviallyDestructible.h"
#endif