// | |
// General purpose minimal heap allocator. | |
// | |
// Keeps a freelist of chunks of arbitrary size. First round of allocations | |
// will use the whole heap block before it starts going into the freelist. | |
// Freeing just pushes the block back into the freelist and is constant time. | |
// Allocating might search the freelist if the whole heap was already allocated | |
// once. | |
// | |
// To reduce fragmentation we could merge free nodes either on free or allocate. | |
// Allocating from the freelist could also split nodes if the node size is bigger | |
// than the requested size. Right now it just returns the first node >= the alloc | |
// size. | |
// | |
#include <cassert> | |
#include <cstdint> | |
#include <cstddef> | |
#include <cstring> | |
#include <cstdio> | |
// ---------------------------------------------------------------------------- | |
#if defined(__clang__) && defined(__has_feature) | |
#if __has_feature(address_sanitizer) | |
#include <sanitizer/asan_interface.h> | |
#define MINIHEAP_ASAN 1 | |
#endif // address_sanitizer | |
#endif // __clang__ && __has_feature | |
#if !defined(MINIHEAP_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 // MINIHEAP_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 // MINIHEAP_ASAN | |
// ---------------------------------------------------------------------------- | |
class MiniHeapAllocator final | |
{ | |
public: | |
MiniHeapAllocator(void * heapBuffer, const std::uint32_t heapSize) | |
{ | |
assert(heapBuffer != nullptr); | |
assert(heapSize != 0); | |
m_HeapStart = static_cast<std::uint8_t *>(heapBuffer); | |
m_HeapSize = heapSize; | |
m_HeapUsed = 0; | |
m_FirstNode = static_cast<Node *>(heapBuffer); | |
m_FirstFree = nullptr; | |
// Keep the whole heap initially poisoned. | |
ASAN_POISON_MEMORY_REGION(m_HeapStart, m_HeapSize); | |
} | |
~MiniHeapAllocator() | |
{ | |
// Unpoison so it can be freed by the user. | |
ASAN_UNPOISON_MEMORY_REGION(m_HeapStart, m_HeapSize); | |
} | |
ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void * Allocate(std::size_t bytes) | |
{ | |
if (bytes < kMinAllocSize) | |
{ | |
bytes = kMinAllocSize; | |
} | |
// Every allocation is prefixed by a Node header. | |
bytes += sizeof(Node); | |
// First batch of allocations will just use the whole memory | |
// range available before we even start recycling. | |
if ((m_HeapUsed + bytes) <= m_HeapSize) | |
{ | |
Node * newAlloc = reinterpret_cast<Node *>(m_HeapStart + m_HeapUsed); | |
m_HeapUsed += bytes; | |
newAlloc->Guard = kGuardValue; | |
newAlloc->Size = bytes; | |
newAlloc->IsFree = false; | |
ASAN_UNPOISON_MEMORY_REGION(newAlloc + 1, bytes - sizeof(Node)); | |
return (newAlloc + 1); // User block follows Node header. | |
} | |
// Try recycling from the freelist: | |
FreeNode * bestMatch = nullptr; | |
for (FreeNode * iter = m_FirstFree; iter != nullptr; iter = iter->Next) | |
{ | |
assert(iter->Guard == kGuardValue); | |
assert(iter->IsFree == true); | |
if (iter->Size >= bytes) | |
{ | |
// TODO: Could split the FreeNode if size > bytes requested. | |
bestMatch = iter; | |
break; | |
} | |
} | |
// Cannot accommodate the block. | |
if (bestMatch == nullptr) | |
{ | |
// TODO: Could try to merge FreeNodes and retry. | |
return nullptr; | |
} | |
// Remove from the freelist: | |
if (bestMatch == m_FirstFree) | |
{ | |
m_FirstFree = m_FirstFree->Next; | |
} | |
else | |
{ | |
FreeNode * prevNode = bestMatch->Prev; | |
FreeNode * nextNode = bestMatch->Next; | |
if (prevNode != nullptr) | |
{ | |
assert(prevNode->Next == bestMatch); | |
prevNode->Next = nextNode; | |
} | |
if (nextNode != nullptr) | |
{ | |
assert(nextNode->Prev == bestMatch); | |
nextNode->Prev = prevNode; | |
} | |
} | |
bestMatch->IsFree = false; | |
ASAN_UNPOISON_MEMORY_REGION(static_cast<Node *>(bestMatch) + 1, bytes - sizeof(Node)); | |
return (static_cast<Node *>(bestMatch) + 1); // User block follows node header. | |
} | |
ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void Free(void * ptr) | |
{ | |
Node * header = static_cast<Node *>(ptr) - 1; | |
assert(IsValidPtr(header)); | |
assert(header->Guard == kGuardValue); | |
assert(header->Size >= sizeof(FreeNode)); | |
assert(header->IsFree == false); | |
ASAN_POISON_MEMORY_REGION(header, header->Size); | |
// Push into the freelist: | |
FreeNode * toFree = static_cast<FreeNode *>(header); | |
if (m_FirstFree != nullptr) | |
{ | |
assert(m_FirstFree->Prev == nullptr); | |
m_FirstFree->Prev = toFree; | |
} | |
toFree->Next = m_FirstFree; | |
toFree->Prev = nullptr; | |
m_FirstFree = toFree; | |
} | |
bool IsValidPtr(const void * const ptr) const | |
{ | |
auto * testAddr = static_cast<const std::uint8_t *>(ptr); | |
auto * heapEnd = m_HeapStart + m_HeapSize; | |
return (testAddr >= m_HeapStart && testAddr < heapEnd); | |
} | |
private: | |
static constexpr std::uint32_t kMinAllocSize = 16; // FreeNode Prev,Next | |
static constexpr std::uint32_t kGuardValue = 0xDEADBEEF; | |
struct Node | |
{ | |
std::uint32_t Guard; // Used by asserts | |
std::uint32_t Size : 31; | |
std::uint32_t IsFree : 1; | |
}; | |
struct FreeNode : Node | |
{ | |
FreeNode * Prev; | |
FreeNode * Next; | |
}; | |
static_assert(sizeof(Node) == 8, "MiniHeapAllocator::Node size expected to be 16 bytes"); | |
static_assert(sizeof(FreeNode) == 24, "MiniHeapAllocator::FreeNode size expected to be 24 bytes"); | |
private: | |
std::uint8_t * m_HeapStart; | |
std::uint32_t m_HeapSize; | |
std::uint32_t m_HeapUsed; | |
Node * m_FirstNode; | |
FreeNode * m_FirstFree; | |
}; | |
// ---------------------------------------------------------------------------- | |
int main() | |
{ | |
constexpr int kHeapSize = 1 * 1024 * 1024; | |
MiniHeapAllocator miniHeap{ new std::uint8_t[kHeapSize], kHeapSize }; | |
const int sizes[] = { 4, 8, 16, 24, 32, 64, 128, 256, 512 }; | |
constexpr int kMaxSizes = sizeof(sizes) / sizeof(sizes[0]); | |
constexpr int kMaxPtrs = 2048; | |
// 1st pass: | |
{ | |
void * ptrs[kMaxPtrs] = {0}; | |
std::printf("Allocating...\n"); | |
for (int i = 0, s = 0; i < kMaxPtrs; ++i) | |
{ | |
ptrs[i] = miniHeap.Allocate(sizes[s]); | |
if (ptrs[i] == nullptr) | |
{ | |
break; | |
} | |
std::memset(ptrs[i], 'X', sizes[s]); | |
s = (s + 1) % kMaxSizes; | |
} | |
std::printf("Freeing...\n"); | |
for (int i = 0; i < kMaxPtrs; ++i) | |
{ | |
if (ptrs[i] != nullptr) | |
{ | |
miniHeap.Free(ptrs[i]); | |
} | |
} | |
} | |
// 2nd pass: | |
{ | |
void * ptrs[kMaxPtrs] = {0}; | |
std::printf("Allocating (again)...\n"); | |
for (int i = 0, s = 0; i < kMaxPtrs; ++i) | |
{ | |
ptrs[i] = miniHeap.Allocate(sizes[s]); | |
if (ptrs[i] == nullptr) | |
{ | |
break; | |
} | |
std::memset(ptrs[i], 'X', sizes[s]); | |
s = (s + 1) % kMaxSizes; | |
} | |
std::printf("Freeing (again)...\n"); | |
for (int i = 0; i < kMaxPtrs; ++i) | |
{ | |
if (ptrs[i] != nullptr) | |
{ | |
miniHeap.Free(ptrs[i]); | |
} | |
} | |
} | |
std::printf("Done!\n"); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment