tiny_allocator.cpp

//
// Simple allocator for small memory blocks (from 1 to 255 bytes).
// It keeps freelists of small blocks, built on top of 8K pages.
// The allocator only grows - empty pages are never freed back
// to the system.
//
// This allocator could easily be made lockless with atomic freelists.
//

#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <cstdio>

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

#if defined(__clang__) && defined(__has_feature)
    #if __has_feature(address_sanitizer)
        #include <sanitizer/asan_interface.h>
        #define TINY_ALLOC_ASAN 1
    #endif // address_sanitizer
#endif // __clang__ && __has_feature

#if !defined(TINY_ALLOC_ASAN)
    #define ASAN_POISON_MEMORY_REGION(addr, size)   /* nothing */
    #define ASAN_UNPOISON_MEMORY_REGION(addr, size) /* nothing */
    #define ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE      /* nothing */
#else // TINY_ALLOC_ASAN
    // ASAN_POISON_MEMORY_REGION()   already defined in asan_interface.h
    // ASAN_UNPOISON_MEMORY_REGION() already defined in asan_interface.h
    #define ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE __attribute__((no_sanitize("address")))
#endif // TINY_ALLOC_ASAN

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

class TinyAllocator final
{
private:

    enum Constants
    {
        kHeaderSize      = 2,
        kAlign           = 8,
        kTinyAllocTag    = 0xAA,
        kInvalidAllocTag = 0xDD,
        kPageSize        = 8192,
    };

    struct Page
    {
        Page * Next;
        std::uint8_t Memory[kPageSize - sizeof(Page *)];
    };

    static_assert(sizeof(Page) == kPageSize, "Wrong size for TinyAllocator::Page");

private:

    static constexpr std::size_t AlignSize(std::size_t bytes)
    {
        return ((bytes + kAlign - 1) & ~(kAlign - 1));
    }

    static constexpr std::size_t TinyAlign(std::size_t bytes)
    {
        return AlignSize(bytes + kHeaderSize) - kHeaderSize;
    }

    static Page * AllocPage()
    {
        Page * p = static_cast<Page *>(std::malloc(sizeof(Page)));
        if (p == nullptr)
        {
            return nullptr;
        }

        p->Next = nullptr;
        std::memset(p->Memory, kInvalidAllocTag, sizeof(Page::Memory));

        ASAN_POISON_MEMORY_REGION(p, sizeof(Page));
        return p;
    }

    static void FreePage(Page * p)
    {
        ASAN_UNPOISON_MEMORY_REGION(p, sizeof(Page));
        std::free(p);
    }

    Page * m_FirstUsedPage;
    Page * m_CurrentPage;
    std::size_t m_CurrentPageOffset;
    void * m_FirstFree[(256 / kAlign) + 1];

public:

    TinyAllocator()
        : m_FirstUsedPage{ nullptr }
        , m_CurrentPage{ AllocPage() }
        , m_CurrentPageOffset{ TinyAlign(0) }
    {
        std::memset(m_FirstFree, 0, sizeof(m_FirstFree));
    }

    ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE ~TinyAllocator()
    {
        for (Page * p = m_FirstUsedPage; p != nullptr;)
        {
            Page * next = p->Next;
            FreePage(p);
            p = next;
        }
    }

    ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void * Allocate(std::size_t bytes)
    {
        if (bytes > 255)
        {
            std::fprintf(stderr, "TinyAllocator::Allocate - Max allocation size is 255. Requested %zu bytes", bytes);
            std::abort();
        }

        // We need the at least the size of a pointer for the free list.
        if (bytes < sizeof(void *))
        {
            bytes = sizeof(void *);
        }

        // Increase the number of bytes if necessary to make sure the next tiny allocation is aligned.
        bytes = TinyAlign(bytes);

        // Append to the free list:
        auto * tinyBlock = static_cast<std::uint8_t *>(m_FirstFree[bytes / kAlign]);
        if (tinyBlock != nullptr)
        {
            std::size_t * link = reinterpret_cast<std::size_t *>(tinyBlock + kHeaderSize);
            tinyBlock[1] = kTinyAllocTag; // Allocation identifier
            m_FirstFree[bytes / kAlign] = reinterpret_cast<void *>(*link);
            void * userPtr = reinterpret_cast<void *>(link);

            ASAN_UNPOISON_MEMORY_REGION(userPtr, bytes);
            return userPtr;
        }

        // Check if we need to allocate a new page:
        const std::size_t bytesLeftInCurrPage = sizeof(Page::Memory) - m_CurrentPageOffset;
        if (bytes >= bytesLeftInCurrPage)
        {
            Page * newPage = AllocPage();
            if (newPage == nullptr)
            {
                return nullptr;
            }

            m_CurrentPage->Next = m_FirstUsedPage;
            m_FirstUsedPage     = m_CurrentPage;
            m_CurrentPage       = newPage;

            // Make sure the first allocation is aligned.
            m_CurrentPageOffset = TinyAlign(0);
        }

        // Allocate from page:
        tinyBlock = m_CurrentPage->Memory + m_CurrentPageOffset;
        tinyBlock[0] = std::uint8_t(bytes / kAlign); // Write # of bytes/kAlign
        tinyBlock[1] = kTinyAllocTag;                // Allocation identifier
        m_CurrentPageOffset += bytes + kHeaderSize;  // Increase the offset on the current page
        void * userPtr = (tinyBlock + kHeaderSize);  // Skip the first two bytes header

        ASAN_UNPOISON_MEMORY_REGION(userPtr, bytes);
        return userPtr;
    }

    ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void Free(void * ptr)
    {
        auto * tinyBlock = static_cast<std::uint8_t *>(ptr) - kHeaderSize;

        // Check allocation tag:
        if (tinyBlock[1] != kTinyAllocTag)
        {
            std::fprintf(stderr, "TinyAllocator::Free - Invalid header tag at %p", ptr);
            std::abort();
        }

        // Clear tag:
        tinyBlock[1] = kInvalidAllocTag;

        // Index into the table with free tiny memory blocks:
        const std::size_t index = tinyBlock[0];

        // Check if the index is valid:
        if (index > (256 / kAlign))
        {
            std::fprintf(stderr, "TinyAllocator::Free - Invalid list index at %p", ptr);
            std::abort();
        }

        std::size_t * link = reinterpret_cast<std::size_t *>(ptr);
        *link = reinterpret_cast<std::size_t>(m_FirstFree[index]);
        m_FirstFree[index] = tinyBlock;

        ASAN_POISON_MEMORY_REGION(ptr, index * kAlign);
    }

    // Disallow copy.
    TinyAllocator(const TinyAllocator & other) = delete;
    TinyAllocator& operator = (const TinyAllocator & other) = delete;
};

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

int main()
{
    TinyAllocator tAlloc;
    void* ptrs[256];

    std::printf("Allocating...\n");
    for (int i = 0; i < 256; ++i)
    {
        ptrs[i] = tAlloc.Allocate(i);
        std::memset(ptrs[i], 0xCD, i);
    }

    std::printf("Freeing...\n");
    for (int i = 0; i < 256; ++i)
    {
        tAlloc.Free(ptrs[i]);
    }

    char * p64 = static_cast<char *>(tAlloc.Allocate(64));
    for (int i = 0; i < 64; ++i)
    {
        p64[i] = 'X';
    }
    tAlloc.Free(p64);
    //p64[0] = 'Y'; ASan ERROR: use after free

    std::printf("Done!\n");
}