// SmallBlockAllocator.h
//------------------------------------------------------------------------------
#pragma once

// Defines
//------------------------------------------------------------------------------
// Enable SmallBlockAllocator unless:
//   * We are using AddressSanitizer, in which case it is better to have a
//     separate block for each allocation to catch out-of-bounds reads/writes.
//   * We are using MemorySanitizer, in which case it is better to not reuse
//     allocated blocks because doing so will hide uninitialized reads from
//     reused blocks from MemorySanitizer.
#if !defined( __has_feature )
    #define __has_feature( ... ) 0
#endif
#if !__has_feature( address_sanitizer ) && !__has_feature( memory_sanitizer ) && !defined( __SANITIZE_ADDRESS__ )
    #define SMALL_BLOCK_ALLOCATOR_ENABLED
#endif

// Includes
//------------------------------------------------------------------------------
#include "Core/Mem/MemPoolBlock.h"
#include "Core/Process/Mutex.h"

// SmallBlockAllocator
//------------------------------------------------------------------------------
class SmallBlockAllocator
{
public:
    // Attempt allocation. Returns nullptr if allocation can't be satisfied
    static void * Alloc( size_t size, size_t align );

    // Attempt to free. Returns false if not a bucket owned allocation
    static bool Free( void * ptr );

#if defined( DEBUG )
    static void DumpStats();
#endif

protected:
    static void InitBuckets();

    static const size_t BUCKET_MAX_ALLOC_SIZE = 256;
    static const size_t BUCKET_ALIGNMENT = 16;
    static const size_t BUCKET_NUM_BUCKETS = ( BUCKET_MAX_ALLOC_SIZE / BUCKET_ALIGNMENT );
    static const size_t BUCKET_ADDRESSSPACE_SIZE = ( 200 * 1024 * 1024 );
    static const size_t BUCKET_NUM_PAGES = ( BUCKET_ADDRESSSPACE_SIZE / MemPoolBlock::kMemPoolBlockPageSize );
    static const size_t BUCKET_MAPPING_TABLE_SIZE = BUCKET_NUM_PAGES;

    PRAGMA_DISABLE_PUSH_MSVC( 4324 ) // structure was padded due to alignment specifier
    class alignas( 64 ) MemBucket : public MemPoolBlock
    {
    public:
        MemBucket( size_t size, size_t align )
            : MemPoolBlock( size, align )
        {
        }

    protected:
        virtual void * AllocateMemoryForPage() override;

        friend class SmallBlockAllocator;
        Mutex m_Mutex;
    };
    PRAGMA_DISABLE_POP_MSVC

    // Address space used by allocators
    static void * s_BucketMemoryStart;
    static uint32_t s_BucketNextFreePageIndex; // Next free memory page to commit

    // The actual buckets (using this placeholder memory to avoid static init issues)
    class alignas( alignof( MemBucket ) ) MemBucketStorage
    {
    public:
        uint8_t m_Storage[ sizeof( MemBucket ) ] = { 0 };
    };
    static MemBucketStorage s_BucketMemBucketMemory[ BUCKET_NUM_BUCKETS ];
    static MemBucket * s_Buckets;

    // A table to allow 0(1) conversion of any address to the bucket that owns it
    static uint8_t s_BucketMappingTable[ BUCKET_MAPPING_TABLE_SIZE ];
};

//------------------------------------------------------------------------------
