Heap Allocator now commits and releases pages dynamically.

Implemented Heap::free
Heap::reset() now clears m_usedSize

.github/workflows/cmake-multi-platform.yml

@@ -76,7 +76,8 @@ jobs:
       run: cmake --build ${{ steps.strings.outputs.build-output-dir }} --config ${{ matrix.build_type }}
 
     - name: Test
-      working-directory: ${{ steps.strings.outputs.build-output-dir }}
+      working-directory: ${{ steps.strings.outputs.build-output-dir }}/tests
       # Execute tests defined by the CMake configuration. Note that --build-config is needed because the default Windows generator is a multi-config generator (Visual Studio generator).
       # See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail
-      run: ctest --build-config ${{ matrix.build_type }}
+      run: ctest --output-on-failure --no-tests=error --build-config ${{ matrix.build_type }}
+

.gitignore

@@ -1 +1,2 @@
+.vs
 build

include/arena.h

@@ -56,6 +56,14 @@ public:
 
   // Reset the arena, potentially also releasing committed pages
   void reset();
+
+  // Get the actual used size.
+  // This may differ from the sum of allocations due to alignment requirements.
+  size_t getUsed() const;
+
+  // Get the actual maximum size.
+  // This may be greater than the maxSize passed into init, due to page size alignment.
+  size_t getMaxSize() const;
 private:
   Region m_region;
   size_t m_minSize;

include/block.h

@@ -35,6 +35,7 @@
 #include <stdint.h>
 
 #if defined(_WIN32) || defined(_WIN64)
+#define NOMINMAX
 #include <Windows.h>
 #endif
 

include/heap.h

@@ -46,7 +46,7 @@ public:
   Heap();
   ~Heap();
 
-  bool init(size_t maxSize = 1ULL << 32);
+  bool init(size_t minSize = 1ULL << 24, size_t maxSize = 1ULL << 32);
 
   // Allocate `size` bytes
   std::byte* alloc(size_t size);
@@ -60,14 +60,42 @@ public:
   // Free the allocation at `address`
   void free(std::byte* address);
 
+  // Reset the heap, freeing all memory and potentially releasing comitted pages.
+  void reset();
+
+  // Get the actual used size.
+  // This may differ from the sum of allocations due to alignment requirements.
+  size_t getUsed() const;
+
+  // Get the actual maximum size.
+  // This may be greater than the maxSize passed into init, due to page size alignment.
+  size_t getMaxSize() const;
+
 private:
   Region m_region;
+  size_t m_minSize;
+  size_t m_usedSize;
 
   // Add a free block to the index
   void insertFreeBlock(TLSFBlock* block);
 
   // Remove a free block from the index
   void removeFreeBlock(TLSFBlock* block);
+
+  // Find a free block that can hold at least size bytes
+  TLSFBlock* findFreeBlock(size_t size) const;
+
+  // Get the usable size of a block
+  size_t getBlockUsableSize(const TLSFBlock* block) const;
+
+  // Get the next block, in physical order
+  TLSFBlock* getNextPhysBlock(TLSFBlock* block);
+
+  // Returns true if block is a free block
+  bool isBlockFree(const TLSFBlock* block) const;
+
+  // Return true if block is the last block, in physical order
+  bool isBlockLast(const TLSFBlock* block) const;
 };
 
 } // namespace mimir

include/region.h

@@ -49,6 +49,10 @@ public:
   size_t getMaxSize() const;
   std::byte* getAddress() const;
 
+  // Helper functions for allocators
+  bool maybeGrow(size_t minSize);
+  void maybeShrink(size_t targetSize);
+
 private:
   std::byte* m_data;
   size_t m_committedSize;

src/arena.cpp

@@ -60,17 +60,25 @@ bool Arena::init(size_t minSize, size_t maxSize)
   return m_region.init(maxSize);
 }
 
+// Get the actual used size.
+// This may differ from the sum of allocations due to alignment requirements.
+size_t Arena::getUsed() const
+{
+  return m_usedSize;
+}
+
+// Get the actual maximum size.
+// This may be greater than the maxSize passed into init, due to page size alignment.
+size_t Arena::getMaxSize() const
+{
+  return m_region.getMaxSize();
+}
+
 std::byte* Arena::alloc(size_t size)
 {
   size_t neededSize = m_usedSize + size;
-  if (neededSize > m_region.getSize())
-  {
-    // Increase the size exponentially when we run out
-    size_t newSize = std::bit_ceil(neededSize) << 1;
-    if (newSize < m_minSize) {
-      newSize = m_minSize;
-    }
-    m_region.resize(newSize);
+  if (!m_region.maybeGrow(std::max(neededSize, m_minSize))) {
+    return nullptr;
   }
 
   std::byte* ret = m_region.getAddress() + m_usedSize;
@@ -80,8 +88,8 @@ std::byte* Arena::alloc(size_t size)
 
 std::byte* Arena::allocA16(size_t size)
 {
-  uint32_t nextByte = (m_usedSize + 15) & ~15;
-  uint32_t roundedSize = (size + 15) & ~15;
+  size_t nextByte = (m_usedSize + 15) & ~15;
+  size_t roundedSize = (size + 15) & ~15;
   if (alloc(roundedSize + nextByte - m_usedSize) == nullptr)     {
     return nullptr;
   }
@@ -91,8 +99,8 @@ std::byte* Arena::allocA16(size_t size)
 
 std::byte* Arena::allocA64(size_t size)
 {
-  uint32_t nextByte = (m_usedSize + 63) & ~63;
-  uint32_t roundedSize = (size + 63) & ~63;
+  size_t nextByte = (m_usedSize + 63) & ~63;
+  size_t roundedSize = (size + 63) & ~63;
   if (alloc(roundedSize + nextByte - m_usedSize) == nullptr) {
     return nullptr;
   }
@@ -112,13 +120,7 @@ void Arena::reset()
   m_watermark[kWatermarkLen - 1] = m_usedSize;
 
   size_t targetSize = std::bit_ceil(highWatermark) << 1;
-  if (targetSize < m_minSize) {
-    targetSize = m_minSize;
-  }
-  size_t thresholdSize = targetSize << 1; // The threshold for shrinking is greater than the target to implement hysteresis
-  if (m_region.getSize() > thresholdSize) {
-    m_region.resize(targetSize);
-  }
+  m_region.maybeShrink(std::max(m_minSize, targetSize));
 
   m_usedSize = 0;
 }

src/heap.cpp

@@ -32,6 +32,7 @@
 #include "../include/mimir.h"
 #include "mimir_impl.h"
 
+#include <algorithm>
 #include <cassert>
 #include <cstring>
 #include <bit>
@@ -40,7 +41,7 @@ namespace mimir {
 
 // Two-Level Segregated Fit (TLSF): http://www.gii.upv.es/tlsf/files/papers/ecrts04_tlsf.pdf
 
-const size_t kMinblockSize = 16;
+const size_t kMinBlockSize = 16;
 
 struct TLSFBlock
 {
@@ -94,6 +95,8 @@ uint64_t blockSizeToIndex(uint64_t size)
 */
 
 Heap::Heap()
+  : m_minSize(0)
+  , m_usedSize(0)
 {
 }
 
@@ -101,41 +104,79 @@ Heap::~Heap()
 {
 }
 
-bool Heap::init(size_t maxSize)
+bool Heap::init(size_t minSize, size_t maxSize)
 {
+  assert(maxSize >= minSize);
+  assert(m_usedSize == 0);
+
+  m_minSize = sizeof(TLSFIndex) + 16;
+  m_minSize = std::max(m_minSize, minSize);
+  m_minSize = KRAKEN_MEM_ROUND_UP_PAGE(m_minSize);
   if (!m_region.init(maxSize)) {
     return false;
   }
-  if (!m_region.resize(sizeof(TLSFIndex) + 16)) {
+  if (!m_region.resize(m_minSize)) {
     return false;
   }
-  TLSFIndex* index = (TLSFIndex*)m_region.getAddress();
-  memset(index, 0, sizeof(TLSFIndex));
-
-  // Start with one free block, filling the entire Region
-  TLSFBlock* block = (TLSFBlock*)(m_region.getAddress() + sizeof(TLSFIndex));
-  block->nextFree = nullptr;
-  block->prevFree = nullptr;
-  block->size = m_region.getMaxSize() - sizeof(TLSFIndex);
-  block->size |= 0b11; // T=1: Last Block, F=1: Free Block
-
-  insertFreeBlock(block);
-
+  reset();
   return true;
 }
 
+// Get the actual maximum size.
+// This may be greater than the maxSize passed into init, due to page size alignment.
+size_t Heap::getMaxSize() const
+{
+  return m_region.getSize();
+}
+
+// Get the actual used size.
+// This may differ from the sum of allocations due to alignment requirements.
+size_t Heap::getUsed() const
+{
+  return m_usedSize;
+}
+
+size_t Heap::getBlockUsableSize(const TLSFBlock* block) const
+{
+  return (block->size & ~0b11) - 16ULL;
+}
+
+TLSFBlock* Heap::getNextPhysBlock(TLSFBlock* block)
+{
+  if (isBlockLast(block)) {
+    return nullptr;
+  }
+
+  return (TLSFBlock*)((std::byte*)block + block->size);
+}
+
+// Return true if block is the last block in physical order
+bool Heap::isBlockLast(const TLSFBlock* block) const
+{
+  return (block->size & 0b10) == 0b10;
+}
+
+// Returns true if block is a free block
+bool Heap::isBlockFree(const TLSFBlock* block) const
+{
+  return (block->size & 0b01) == 0b01;
+}
+
 void Heap::insertFreeBlock(TLSFBlock* block)
 {
   // Add the block to the index
 
   TLSFIndex* index = (TLSFIndex*)m_region.getAddress();
 
-  size_t usableSize = (block->size & ~0b11) - 16;
-
-  size_t firstLevelIndex = std::bit_width(usableSize) - 5;
+  size_t usableSize = getBlockUsableSize(block);
+  size_t firstLevel = std::bit_width(usableSize);
+  size_t firstLevelIndex = firstLevel - 5;
   size_t secondLevelIndex = (usableSize >> firstLevelIndex) & 0b1111ULL;
+
+  // Update free bitmaps
   index->firstLevelFreeBitMap |= std::bit_floor(usableSize);
-  index->secondLevelFreeBitMap[firstLevelIndex] |= secondLevelIndex & 0b1111ULL;
+  index->secondLevelFreeBitMap[firstLevelIndex] |= secondLevelIndex;
+
   TLSFBlock* prevFirstFreeBlock = index->secondLevelFreeBlocks[firstLevelIndex][secondLevelIndex];
   index->secondLevelFreeBlocks[firstLevelIndex][secondLevelIndex] = block;
 
@@ -144,6 +185,9 @@ void Heap::insertFreeBlock(TLSFBlock* block)
     block->nextFree = prevFirstFreeBlock;
     prevFirstFreeBlock->prevFree = block;
   }
+
+  // Mark block as free
+  block->size |= 0b01; // T=1: Last Block, F=1: Free Block
 }
 
 void Heap::removeFreeBlock(TLSFBlock* block)
@@ -151,24 +195,25 @@ void Heap::removeFreeBlock(TLSFBlock* block)
   // Remove the block from the index
   TLSFIndex* index = (TLSFIndex*)m_region.getAddress();
 
-  size_t usableSize = (block->size & ~0b11) - 16;
-
+  size_t usableSize = getBlockUsableSize(block);
   size_t firstLevel = std::bit_width(usableSize);
   size_t firstLevelIndex = firstLevel - 5;
-  size_t secondLevelIndex = (usableSize >> (firstLevel - 5)) & 0b1111ULL;
+  size_t secondLevelIndex = (usableSize >> (firstLevelIndex)) & 0b1111ULL;
 
   if (block->prevFree != nullptr) {
     // Link the neighboring free blocks together
-    block->prevFree = block->nextFree;
+    block->prevFree->nextFree = block->nextFree;
     if (block->nextFree) {
-      block->nextFree = block->prevFree;
+      block->nextFree->prevFree = block->prevFree;
     }
   } else {
     // This block was the first for this level
     index->secondLevelFreeBlocks[firstLevelIndex][secondLevelIndex] = block->nextFree;
 
-    if (block->nextFree == nullptr) {
-      // We have removed all the blocks at this level.
+    if (block->nextFree != nullptr) {
+      block->nextFree->prevFree = nullptr;
+    } else {
+      // We have removed all the free blocks at this level.
 
       // Update second level bitmask.
       index->secondLevelFreeBitMap[firstLevelIndex] &= ~(1ULL << secondLevelIndex);
@@ -182,20 +227,15 @@ void Heap::removeFreeBlock(TLSFBlock* block)
     }
   }
 
-  if (block->size & 0b10ULL) {
-    // This block was the last in physical order.
-    if (block->prevPhys) {
-      // Update the prior block in physical order to mark it as the last block.
-      block->prevPhys->size &= 0b10ULL;
-    }
-  }
+  // Mark block as non-free
+  block->size &= ~0b01; // T=1: Last Block, F=1: Free Block
 }
 
-// Allocate `size` bytes
-std::byte* Heap::alloc(size_t size)
+// Find a free block that can hold at least size bytes
+TLSFBlock* Heap::findFreeBlock(size_t size) const
 {
-  if (size < kMinblockSize) {
-    size = kMinblockSize;
+  if (size < kMinBlockSize) {
+    size = kMinBlockSize;
   }
   TLSFIndex* index = (TLSFIndex*)m_region.getAddress();
 
@@ -235,52 +275,60 @@ std::byte* Heap::alloc(size_t size)
       return nullptr;
     }
 
-    // Any buffer at the second level will fit this allocation.
+    // Any buffer at the second level will fit this allocation. Pick the smallest buffer.
     selectedSecondLevel = std::countr_zero(index->secondLevelFreeBitMap[selectedFirstLevel - 4]);
   }
 
   // Take the first free block
   TLSFBlock* block = index->secondLevelFreeBlocks[selectedFirstLevel - 4][selectedSecondLevel];
 
-  size_t oldBlockSize = block->size & ~0b11ULL;
-  bool oldBlockWasLastPhysBlock = (block->size & ~0b10ULL) != 0;
+  return block;
+}
 
-  TLSFBlock* prevFreeBlock = block->prevFree;
-  TLSFBlock* nextFreeBlock = block->nextFree;
-  TLSFBlock* prevPhysBlock = block->prevPhys;
-  TLSFBlock* nextPhysBlock = nullptr;
-  if (!oldBlockWasLastPhysBlock) {
-    nextPhysBlock = (TLSFBlock*)(((std::byte*)block) + oldBlockSize + 16);
+// Allocate `size` bytes
+std::byte* Heap::alloc(size_t size)
+{
+  if (size < kMinBlockSize) {
+    size = kMinBlockSize;
   }
 
-  // The next free block at this level replaces this block in the index.
-  index->secondLevelFreeBlocks[selectedFirstLevel - 4][selectedSecondLevel] = nextFreeBlock;
+  TLSFBlock* block = findFreeBlock(size);
+  if (block == nullptr) {
+    // No free block found
+    return nullptr;
+  }
 
-  if (nextFreeBlock == nullptr) {
-    // This was the last free block at this level.
-    // Update the bitmap for the second level...
-    index->secondLevelFreeBitMap[selectedFirstLevel - 4] &= ~(1ULL << selectedSecondLevel);
+  // Ensure that we can commit pages to back the new block
+  // as well as a trailing free block header
+  size_t neededCommit = ((std::byte*)block - m_region.getAddress()) + size + 16ULL + 32ULL;
+  if (!m_region.maybeGrow(std::max(m_minSize, neededCommit))) {
+    // Out of system memory
+    return nullptr;
+  }
 
-    // Check if there are any remaining free blocks within the first level
-    if (index->secondLevelFreeBitMap[selectedFirstLevel - 4] == 0) {
+  removeFreeBlock(block);
 
-      // This was the last one.  Clear the first level bit as well.
-      index->firstLevelFreeBitMap &= ~(1ULL << selectedFirstLevel);
+  size_t usableSize = getBlockUsableSize(block);
+
+  if (size - usableSize >= sizeof(TLSFBlock)) {
+    // Shrink the block to size; allocate a new free block for the remainder
+
+    TLSFBlock* freeBlock = (TLSFBlock*)((std::byte*)block + size + 16ULL);
+    freeBlock->size = size - usableSize;
+    freeBlock->prevPhys = block;
+    freeBlock->prevFree = nullptr;
+    freeBlock->nextFree = nullptr;
+    if (isBlockLast(block))
+    { // LSB: T, F: T = Last physical block, F = Free block
+      // Block was the last physical block
+      block->size &= ~0b10;
+      freeBlock->size &= 0b10;
     }
-  } else {
-    // The next free block is now the first free block at this level.
-    nextFreeBlock->prevFree = nullptr;
+    block->size -= freeBlock->size & ~0b11;
+    insertFreeBlock(freeBlock);
   }
 
-  if (nextPhysBlock) {
-    // Link the next physical block to this one.
-    nextPhysBlock->prevPhys = block;
-  } else {
-    // This is the last physical block, so set the last block bit.
-    block->size &= 0b10ULL;
-  }
-
-  block->size = size;
+  m_usedSize += getBlockUsableSize(block);
   return (std::byte*)block + 16;
 }
 
@@ -299,6 +347,71 @@ std::byte* Heap::allocA64(size_t size)
 // Free the allocation at `address`
 void Heap::free(std::byte* address)
 {
+  TLSFBlock* block = (TLSFBlock*)(address - 16ULL);
+  m_usedSize -= getBlockUsableSize(block);
+
+  TLSFBlock* prevPhys = block->prevPhys;
+  if (prevPhys && isBlockFree(prevPhys)) {
+    // The previous block in physical order is a free block.
+
+    // Remove it from the index
+    removeFreeBlock(prevPhys);
+
+    // Resize it to glob together with the block to free
+    prevPhys->size += (block->size & ~0b11);
+    block = prevPhys;
+  }
+
+  // Reset the free block flag and the last block flag
+  block->size &= ~0b11;
+  
+  TLSFBlock* nextPhys = getNextPhysBlock(block);  
+  if (nextPhys == nullptr) {
+    // This is the last block
+    block->size &= 0b10;
+  } else if (isBlockFree(nextPhys)) {
+    // The next block in physical order is a free block.
+
+    // remove it from the index
+    removeFreeBlock(nextPhys);
+
+    // Resize free'd block to glob together with the nextPhys block
+    block->size += (nextPhys->size & ~0b11);
+    if (isBlockLast(nextPhys)) {
+      block->size &= 0b10;
+    }
+  }
+
+  // Add the free'd block to the index
+  insertFreeBlock(block);
+
+  if (isBlockLast(block)) {
+    // Maybe free pages
+    size_t neededSize = ((std::byte*)block - m_region.getAddress()) + 32ULL;
+    m_region.maybeShrink(neededSize);
+  }
+}
+
+// Reset (or initialize) the heap, freeing all memory and potentially releasing comitted pages.
+void Heap::reset()
+{
+  // Clear the index, representing no free blocks
+  TLSFIndex* index = (TLSFIndex*)m_region.getAddress();
+  memset(index, 0, sizeof(TLSFIndex));
+  m_usedSize = 0;
+
+  // Free comitted pages except minimum required
+  m_region.resize(m_minSize);
+
+  // Start with one free block, filling the entire Region
+  TLSFBlock* block = (TLSFBlock*)(m_region.getAddress() + sizeof(TLSFIndex));
+  block->prevPhys = nullptr;
+  block->nextFree = nullptr;
+  block->prevFree = nullptr;
+  block->size = m_region.getMaxSize();
+  block->size |= 0b10; // T=1: Last Block, F=1: Free Block
+
+  insertFreeBlock(block);
 }
 
 } // namespace mimir

src/mimir_impl.h

@@ -39,6 +39,7 @@
 #endif
 
 #if defined(_WIN32) || defined(_WIN64)
+#define NOMINMAX
 #include <Windows.h>
 #include <memoryapi.h>
 #endif

src/region.cpp

@@ -93,6 +93,10 @@ bool Region::resize(size_t size)
     return true;
   }
 
+  if (newSize > m_maxSize) {
+    return false;
+  }
+
   if (newSize > m_committedSize) {
     // ---- Growing ----
 #if defined(_WIN32) || defined(_WIN64)
@@ -150,4 +154,32 @@ std::byte* Region::getAddress() const
   return m_data;
 }
 
-};
+bool Region::maybeGrow(size_t minSize)
+{
+  if (minSize < m_committedSize) {
+    // Already big enough
+    return true;
+  }
+
+  if (minSize > m_maxSize) {
+    // We can't grow this big
+    return false;
+  }
+
+  // We need to grow.
+  // Increase the size exponentially, capped to m_maxSize
+  size_t newSize = std::bit_ceil(minSize) << 1;
+  newSize = std::min(m_maxSize, newSize);
+  
+  return resize(newSize);
+}
+
+void Region::maybeShrink(size_t targetSize)
+{
+  size_t thresholdSize = targetSize << 1; // The threshold for shrinking is greater than the target to implement hysteresis
+  if (getSize() > thresholdSize) {
+    resize(targetSize);
+  }
+}
+
+} // namespace mimir

tests/CMakeLists.txt

@@ -14,8 +14,10 @@ enable_testing()
 
 add_executable(
   mimir_test
+  "arena_test.cpp"
+  "heap_test.cpp"
   "region_test.cpp"
-)
+  "test_util.cpp")
 target_link_libraries(
   mimir_test
   GTest::gtest_main

tests/arena_test.cpp

@@ -0,0 +1,171 @@
+#include <gtest/gtest.h>
+#include <random>
+#include "mimir.h"
+
+#include "test_util.h"
+
+namespace mimir_tests {
+
+TEST(Arena, InitialState)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init());
+  EXPECT_EQ(arena.getUsed(), 0);
+}
+
+TEST(Arena, OutOfMemory)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init(1ULL << 24, 1ULL << 46));
+  EXPECT_TRUE(arena.alloc(1ULL << 49) == nullptr);
+}
+
+TEST(Arena, Overflow)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init(128, 1024));
+  EXPECT_TRUE(arena.alloc(512) != nullptr);
+  EXPECT_TRUE(arena.alloc(512) != nullptr);
+  EXPECT_TRUE(arena.alloc(arena.getMaxSize() - 1024 + 1) == nullptr);
+}
+
+TEST(Arena, Alignment16)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init());
+
+  for (int i = 0; i < 128; i++) {
+    for (int j = 1; j < 128; j++) {
+      EXPECT_TRUE(arena.alloc(i) != nullptr);
+      for (int k = 0; k < 16; k++) {
+        std::byte* addr = arena.allocA16(j);
+        EXPECT_TRUE(addr != nullptr);
+        EXPECT_EQ((size_t)addr & 0b1111, 0);
+        EXPECT_EQ(arena.getUsed() & 0b1111, 0);
+      }
+      arena.reset();
+    }
+  }
+}
+
+TEST(Arena, Alignment64)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init());
+
+  for (int i = 0; i < 128; i++) {
+    for (int j = 1; j < 128; j++) {
+      EXPECT_TRUE(arena.alloc(i) != nullptr);
+      for (int k = 0; k < 16; k++) {
+        std::byte* addr = arena.allocA64(j);
+        EXPECT_TRUE(addr != nullptr);
+        EXPECT_EQ((size_t)addr & 0b111111, 0);
+        EXPECT_EQ(arena.getUsed() & 0b111111, 0);
+      }
+      arena.reset();
+    }
+  }
+}
+
+TEST(Arena, OverflowByAlignmentA16)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init(128, 1024));
+
+  // Controlled test - first try with no alignment. This should pass.
+  EXPECT_TRUE(arena.alloc(alignDown(arena.getMaxSize(), 16) - 15) != nullptr);
+  EXPECT_TRUE(arena.alloc(arena.getMaxSize() - arena.getUsed()) != nullptr);
+  arena.reset();
+
+  // Now try with allocA16. This should fail.
+  EXPECT_TRUE(arena.alloc(alignDown(arena.getMaxSize(), 16) - 15) != nullptr);
+  EXPECT_TRUE(arena.allocA16(arena.getMaxSize() - arena.getUsed()) == nullptr);
+}
+
+TEST(Arena, OverflowByAlignmentA64)
+{
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init(128, 1024));
+
+  // Controlled test - first try with no alignment. This should pass.
+  EXPECT_TRUE(arena.alloc(alignDown(arena.getMaxSize(), 64) - 63) != nullptr);
+  EXPECT_TRUE(arena.alloc(arena.getMaxSize() - arena.getUsed()) != nullptr);
+  arena.reset();
+
+  // Now try with allocA64. This should fail.
+  EXPECT_TRUE(arena.alloc(alignDown(arena.getMaxSize(), 64) - 63) != nullptr);
+  EXPECT_TRUE(arena.allocA64(arena.getMaxSize() - arena.getUsed()) == nullptr);
+}
+
+TEST(Arena, FillAndReset)
+{
+  // We purposefully want to have the same random sequence each time
+
+  const int maxAllocCount = 128;
+  const int cycleCount = 128;
+  std::mt19937_64 rgen(1234);
+  std::uniform_int_distribution<int> count_dist(0, maxAllocCount);
+  std::uniform_int_distribution<size_t> alloc_dist1(1, 256);
+  std::uniform_int_distribution<size_t> alloc_dist2(1, 1024 * 256);
+  std::uniform_int_distribution<size_t> alloc_dist3(1, 1024 * 1024 * 256);
+
+  mimir::init();
+  mimir::Arena arena;
+  EXPECT_TRUE(arena.init(0, 1ULL << 32)); // Use up to 4GB
+  size_t allocSize[maxAllocCount];
+  std::byte* allocAddr[maxAllocCount];
+
+  for (int i = 0; i < cycleCount; i++) {
+    int allocCount = 0;
+    if (i != 5) {
+      // Make sure at least one cycle has 0 allocations
+      allocCount = count_dist(rgen);
+    }
+    for (int j = 0; j < allocCount; j++) {
+      allocSize[j] = alloc_dist1(rgen);
+    }
+    for (int j = 0; j < 32; j++) {
+      int k = count_dist(rgen);
+      if (k < allocCount) {
+        allocSize[k] = alloc_dist2(rgen);
+      }
+    }
+    
+    for (int j = 0; j < 3; j++) {
+      int k = count_dist(rgen);
+      if (k < allocCount) {
+        allocSize[k] = alloc_dist3(rgen);
+      }
+    }
+
+    size_t totalSize = 0;
+    for (int j = 0; j < allocCount; j++) {
+      totalSize += allocSize[j];
+      allocAddr[j] = arena.alloc(allocSize[j]);
+      EXPECT_TRUE(allocAddr[j] != nullptr);
+      if (allocAddr[j] != nullptr) {
+        writeData(allocAddr[j], allocSize[j], j);
+      }
+    }
+
+    for (int j = 0; j < allocCount; j++) {
+      if (allocAddr[j] != nullptr) {
+        EXPECT_TRUE(checkData(allocAddr[j], allocSize[j], j));
+      }
+    }
+
+    EXPECT_EQ(arena.getUsed(), totalSize);
+
+    arena.reset();
+    EXPECT_EQ(arena.getUsed(), 0);
+  }
+}
+
+}; // namespace mimir_tests

tests/heap_test.cpp

@@ -0,0 +1,23 @@
+#include <gtest/gtest.h>
+#include "mimir.h"
+
+#include "test_util.h"
+
+namespace mimir_tests {
+
+TEST(Heap, InitialState)
+{
+  mimir::init();
+  mimir::Heap heap;
+  EXPECT_EQ(heap.getUsed(), 0);
+}
+
+TEST(Heap, OutOfMemory)
+{
+  mimir::init();
+  mimir::Heap heap;
+  EXPECT_TRUE(heap.init(1ULL << 24, 1ULL << 46));
+  EXPECT_TRUE(heap.alloc(1ULL << 49) == nullptr);
+}
+
+}; // namespace mimir_tests

tests/region_test.cpp

@@ -1,25 +1,10 @@
 #include <gtest/gtest.h>
 #include "mimir.h"
 
+#include "test_util.h"
+
 namespace mimir_tests {
 
-void writeData(std::byte* dest, size_t count, int seed)
-{
-  for (int i = 0; i < count; i++) {
-    dest[i] = std::byte((seed * 3162 + i * 63163) % 0xff);
-  }
-}
-
-bool checkData(std::byte* src, size_t count, int seed)
-{
-  for (int i = 0; i < count; i++) {
-    if (src[i] != std::byte((seed * 3162 + i * 63163) % 0xff)) {
-      return false;
-    }
-  }
-  return true;
-}
-
 TEST(Region, InitialState)
 {
   mimir::init();
@@ -41,7 +26,7 @@ TEST(Region, OutOfMemory)
   mimir::init();
   mimir::Region region;
   EXPECT_TRUE(region.init(1ULL << 46)); // Attempt to reserve 64TB of pages
-  EXPECT_FALSE(region.resize(1ULL << 45)); // Attempt to commit 32TB of pages (Assuming we don't have 32TB of ram!)
+  EXPECT_FALSE(region.resize(1ULL << 60)); // Attempt to commit an extremely large number of pages
 }
 
 TEST(Region, Min1GB)
@@ -70,12 +55,12 @@ TEST(Region, Resizing)
   EXPECT_TRUE(region.resize(1ULL << 9));
 }
 
-TEST(Region, OverSized)
+TEST(Region, Oversized)
 {
   mimir::init();
   mimir::Region region;
 
-  EXPECT_TRUE(region.resize(1ULL << 4));
+  EXPECT_TRUE(region.init(1ULL << 4));
   size_t prevSize = region.getSize();
 
   EXPECT_FALSE(region.resize(1ULL << 20));

tests/test_util.cpp

@@ -0,0 +1,50 @@
+#include "test_util.h"
+
+#include <algorithm>
+#include <cstring>
+
+namespace mimir_tests {
+
+const size_t kPatternLen = 32;
+
+void getPattern(std::byte* pattern, int seed)
+{
+  for (int i = 0; i < kPatternLen; i++) {
+    pattern[i] = std::byte((seed * 3162 + i * 63163) % 0xff);
+  }
+}
+
+void writeData(std::byte* dest, size_t count, int seed)
+{
+  std::byte pattern[kPatternLen];
+  getPattern(pattern, seed);
+
+  for (int i = 0; i < count; i += kPatternLen) {
+    memcpy(dest + i, pattern, std::min(count - i, kPatternLen));
+  }
+}
+
+bool checkData(std::byte* src, size_t count, int seed)
+{
+  std::byte pattern[kPatternLen];
+  getPattern(pattern, seed);
+
+  for (int i = 0; i < count; i += kPatternLen) {
+    if (memcmp(src + i, pattern, std::min(count - i, kPatternLen)) != 0) {
+      return false;
+    }
+  }
+  return true;
+}
+
+size_t alignUp(size_t val, size_t alignment)
+{
+  return alignDown(val - 1, alignment) + alignment;
+}
+
+size_t alignDown(size_t val, size_t alignment)
+{
+  return val & ~(alignment - 1);
+}
+
+}; // namespace mimir_tests

tests/test_util.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <cstddef>
+
+namespace mimir_tests {
+
+void writeData(std::byte* dest, size_t count, int seed);
+bool checkData(std::byte* src, size_t count, int seed);
+size_t alignUp(size_t val, size_t alignment);
+size_t alignDown(size_t val, size_t alignment);
+
+}; // namespace mimir_tests