Allocator.h

#pragma once

#include "Allocator.inl"

namespace SE {
template<bool SingleFrame>
void *TDynamicAllocator<SingleFrame>::AlignPointer(void *Pointer, uint32 Extra, uint32 Alignment) {
    byte *Address = (byte *)Pointer + Extra;  // Space for header.
    uintptr_t Aligned = Math::AlignUp((uintptr_t)Address, (uintptr_t)Alignment);
    return (void *)Aligned;
}
template<bool SingleFrame>
void *TDynamicAllocator<SingleFrame>::Alloc(uint32 InSize, uint32 InAlignment) {
    Impl::IncreaseAllocCount();
    if constexpr (SingleFrame) {
        const uint32 Alignment = Math::Max(InAlignment, Meta::DefaultAlignment);
        const uint32 Size = InSize + Alignment;
        byte *Pointer = (byte *)Impl::FrameAlloc(Size);
        byte *Aligned = (byte *)AlignPointer(Pointer, 0, Alignment);
        return Aligned;
    } else {
        if (InAlignment <= Meta::DefaultAlignment) {
            return malloc(InSize);
        } else {
            const uint32 Alignment = Math::Max(InAlignment, (uint32)sizeof(uint32));
            const uint32 Size = InSize + Alignment;
            byte *Pointer = (byte *)malloc(Size);
            byte *Aligned = (byte *)AlignPointer(Pointer, sizeof(uint32), Alignment);
            uint32 *Header = ((uint32 *)Aligned) - 1;
            *Header = Aligned - Pointer;
            return Aligned;
        }
    }
}
template<bool SingleFrame>
void *TDynamicAllocator<SingleFrame>::Realloc(void *InPointer, uint32 InSize, uint32 InOldSize, uint32 InAlignment) {
    Impl::IncreaseFreeCount();
    Impl::IncreaseAllocCount();
    if constexpr (SingleFrame) {
        uint32 Alignment = Math::Max(InAlignment, Meta::DefaultAlignment);
        uint32 Size = InSize + Alignment;
        byte *NewPointer = (byte *)Impl::FrameAlloc(Size);
        byte *NewAligned = (byte *)AlignPointer(NewPointer, 0, Alignment);
        memcpy(NewAligned, InPointer, Math::Min(InOldSize, InSize));

        return NewAligned;
    } else {
        if (InAlignment <= Meta::DefaultAlignment) {
            return realloc(InPointer, InSize);
        } else {
            byte *Aligned = (byte *)InPointer;
            uint32 *Header = ((uint32 *)Aligned) - 1;
            uint32 Offset = *Header;
            byte *Pointer = Aligned - Offset;

            uint32 Alignment = Math::Max(InAlignment, (uint32)sizeof(uint32));
            uint32 Size = InSize + Alignment;
            Pointer = (byte *)realloc(Pointer, Size);

            byte *NewAligned = (byte *)AlignPointer(Pointer, (uint32)sizeof(uint32), Alignment);
            if (NewAligned == Aligned) {
                return Aligned;
            }

            Aligned = Pointer + Offset;
            memmove(NewAligned, Aligned, InSize);

            uint32 *NewHeader = ((uint32 *)NewAligned) - 1;
            *NewHeader = NewAligned - Pointer;
            return NewAligned;
        }
    }
}
template<bool SingleFrame>
void TDynamicAllocator<SingleFrame>::Free(void *InPointer, uint32 Alignment) {
    if (InPointer) {
        Impl::IncreaseFreeCount();
    }
    if constexpr (!SingleFrame) {
        if (Alignment <= Meta::DefaultAlignment) {
            free(InPointer);
        } else if (InPointer) {
            byte *Aligned = (byte *)InPointer;
            uint32 *Header = ((uint32 *)Aligned) - 1;
            byte *Pointer = Aligned - *Header;
            free(Pointer);
        }
    }
}

template<uint32 InElementSize, uint32 InElementAlignment>
TBlockAllocator<InElementSize, InElementAlignment>::~TBlockAllocator() {
    for (Impl::FBlockAllocatorBlock *Block = BlockHead; Block;) {
        FHeapAllocator::Free(Block->Memory, InElementAlignment);
        Impl::FBlockAllocatorBlock *Next = Block->Next;
        FHeapAllocator::Free(Block, alignof(decltype(Block)));
        Block = Next;
    }
}
template<uint32 InElementSize, uint32 InElementAlignment>
void *TBlockAllocator<InElementSize, InElementAlignment>::Alloc() {
    if (!Head) {
        Head = AllocBlock();
    }
    void *Result = Head;
    Head = Head->Next;
    return Result;
}
template<uint32 InElementSize, uint32 InElementAlignment>
void TBlockAllocator<InElementSize, InElementAlignment>::Free(void *Pointer) {
    Impl::FBlockAllocatorChunkHeader *Chunk = (Impl::FBlockAllocatorChunkHeader *)Pointer;
    Chunk->Next = Head;
    Head = Chunk;
}
template<uint32 InElementSize, uint32 InElementAlignment>
Impl::FBlockAllocatorChunkHeader *TBlockAllocator<InElementSize, InElementAlignment>::AllocBlock() {
    Impl::FBlockAllocatorChunkHeader *FirstChunk = (Impl::FBlockAllocatorChunkHeader *)FHeapAllocator::Alloc(BlockSize, InElementAlignment);
    Impl::FBlockAllocatorBlock *Block =
        (Impl::FBlockAllocatorBlock *)FHeapAllocator::Alloc(sizeof(Impl::FBlockAllocatorBlock), alignof(Impl::FBlockAllocatorBlock *));
    Block->Memory = FirstChunk;
    if (!BlockTail) {
        Check(!BlockHead);
        Block->Next = nullptr;
        BlockHead = Block;
        BlockTail = Block;
    } else {
        Block->Next = nullptr;
        BlockTail->Next = Block;
        BlockTail = Block;
    }
    Impl::FBlockAllocatorChunkHeader *Chunk = FirstChunk;
    for (uint32 Index = 0; Index < ElementsPerBlock - 1; ++Index) {
        Chunk->Next = (Impl::FBlockAllocatorChunkHeader *)((byte *)Chunk + ChunkSize);
        Chunk = Chunk->Next;
    }
    Chunk->Next = nullptr;
    return FirstChunk;
}
}  // namespace SE

Allocator.inl

#pragma once

#include "DebugUtility/DebugUtility.h"
#include "Math/Math.h"
#include "Meta/Meta.h"

#include <cstdlib>
#include <cstring>

namespace SE {

namespace Impl {
void TickFrameMemory();
void *FrameAlloc(uint32 Size);

void IncreaseAllocCount();
void IncreaseFreeCount();

void *PageAlloc(uint64 Size);

inline constexpr uint32 MaxThreads = 16;
}  // namespace Impl

template<bool SingleFrame>
struct TDynamicAllocator {
    static void *AlignPointer(void *Pointer, uint32 Extra, uint32 Alignment);
    static void *Alloc(uint32 InSize, uint32 InAlignment);
    static void *Realloc(void *InPointer, uint32 InSize, uint32 InOldSize, uint32 InAlignment);
    static void Free(void *InPointer, uint32 Alignment);
};
// Do not store FFrameAllocator-based collections, their MaxSize corrupts the state.
using FFrameAllocator = TDynamicAllocator<true>;
using FHeapAllocator = TDynamicAllocator<false>;

template<uint32 InMaxSize>
struct TStackAllocator {
    static constexpr uint32 MaxSize = InMaxSize;
};

namespace Impl {
struct FBlockAllocatorBlock {
    void *Memory;
    FBlockAllocatorBlock *Next;
};
struct FBlockAllocatorChunkHeader {
    FBlockAllocatorChunkHeader *Next;
};
struct FBlockAllocatorChunkPayload {
    FBlockAllocatorBlock *Block;
};
}  // namespace Impl

// NOTE: Do not attempt supporting shrinking through block active chunk counter as chunk linked list has random order.
template<uint32 InElementSize, uint32 InElementAlignment>
class TBlockAllocator {
public:
    TBlockAllocator() = default;
    
    TBlockAllocator(const TBlockAllocator &) = delete;
    TBlockAllocator operator=(const TBlockAllocator &) = delete;
    
    TBlockAllocator(TBlockAllocator &&Other) {
        Swap(Head, Other.Head);
        Swap(BlockHead, Other.BlockHead);
        Swap(BlockTail, Other.BlockTail);
    }
    
    TBlockAllocator &operator=(TBlockAllocator &&Other) {
        Swap(Head, Other.Head);
        Swap(BlockHead, Other.BlockHead);
        Swap(BlockTail, Other.BlockTail);
        return *this;
    }
    
    ~TBlockAllocator();

    void *Alloc();
    void Free(void *Pointer);

private:
    static constexpr uint32 MinElementSize = Math::Max(InElementSize, uint32(sizeof(Impl::FBlockAllocatorChunkHeader)));
    static constexpr uint32 ChunkSize = Math::AlignUp(MinElementSize, InElementAlignment);
    static constexpr uint32 DesiredBlockSize = 4096;
    static constexpr uint32 ElementsPerBlock = Math::Max(DesiredBlockSize / ChunkSize, (uint32)1);
    static constexpr uint32 BlockSize = ElementsPerBlock * ChunkSize;

    Impl::FBlockAllocatorChunkHeader *Head = nullptr;
    Impl::FBlockAllocatorBlock *BlockHead = nullptr;
    Impl::FBlockAllocatorBlock *BlockTail = nullptr;

    Impl::FBlockAllocatorChunkHeader *AllocBlock();
};

}  // namespace SE