Fix potential overflow when we're nearing 4G blocks..

src/arena.c

@@ -2,29 +2,22 @@
 #include "base_core.h"
 #include "platform.h"
 
-typedef struct MemoryBlockOptions MemoryBlockOptions;
-struct MemoryBlockOptions {
-    U32 size;
-    U32 commit_size;
-    U32 commit_now_size;
-    U32 page_size;
-};
-static MemoryBlock *memoryblock_allocate(const MemoryBlockOptions *opts)
-{
-    U32 aligned_size            = AlignUpPow2(opts->size, opts->page_size);
-    U32 aligned_commit_size     = AlignUpPow2(opts->commit_size, opts->page_size);
-    U32 aligned_commit_now_size = AlignUpPow2(opts->commit_now_size, opts->page_size);
 
-    MemoryBlock *block = memory_reserve(aligned_size);
+static MemoryBlock *memoryblock_allocate(U32 size, U32 commit_size, U32 page_size)
+{
+    Assert(IsAlignedTo(size, page_size));
+    Assert(IsAlignedTo(commit_size, page_size));
+
+    MemoryBlock *block = memory_reserve(size);
     Assert(block != NULL);
-    memory_commit(block, aligned_commit_now_size);
+    memory_commit(block, commit_size);
 
     block->base          = (U8 *)block;
-    block->size          = aligned_size;
+    block->size          = size;
     block->first_free    = MemoryBlock_Header_Size;
     block->arena_offset  = 0;
     block->prev          = NULL;
-    block->committed_end = aligned_commit_now_size;
+    block->committed_end = commit_size;
 
     ASAN_POISON_MEMORY_REGION(block->base, block->size);
     ASAN_UNPOISON_MEMORY_REGION(block, MemoryBlock_Header_Size);
@@ -37,22 +30,19 @@ void arena_init(Arena *arena, const ArenaOptions *opts)
     Assert(opts->size > MemoryBlock_Header_Size);
     Assert(opts->size < (U32_Max -  page_size)); // Reserving last piece so our math doesn't overflow
     Assert(opts->commit_size > 0);
-    Assert(opts->commit_size <= opts->size);
 
-    MemoryBlockOptions mb_opts;
-    mb_opts.page_size       = page_size;
-    mb_opts.commit_size     = opts->commit_size;
-    mb_opts.size            = opts->size;
-    mb_opts.commit_now_size = opts->commit_size;
-    MemoryBlock *block = memoryblock_allocate(&mb_opts);
+    U32 aligned_commit_size = AlignUp(opts->commit_size, page_size);
+    U32 aligned_size        = AlignUp(opts->size, page_size);
+    Assert(aligned_commit_size <= aligned_size);
+
+    MemoryBlock *block  = memoryblock_allocate(aligned_size, aligned_commit_size, page_size);
 
     arena->flags            = opts->flags;
     arena->checkpoint_level = 0;
     arena->current          = block;
-    arena->block_size       = opts->size;
-    arena->commit_size      = opts->commit_size;
+    arena->block_size       = aligned_size;
+    arena->commit_size      = aligned_commit_size;
     arena->page_size        = page_size;
-
 }
 
 void arena_free(Arena *arena)
@@ -69,35 +59,32 @@ void arena_free(Arena *arena)
 
 void *arena_push(Arena *arena, U32 size, U32 align)
 {
-    Assert(size < (U32_Max - platform_page_size()));
     Assert(IsPow2(align));
 
+
     MemoryBlock *block = arena->current;
-    U32 start_pos = AlignUpPow2(block->first_free, align);  // might align past block->size, handled in below check
+    U32 start_pos = AlignUp(block->first_free, align);  // might align past block->size, handled in below check
     U32 available = block->size - start_pos;
 
     if (start_pos > block->size || available < size) {
-        U32 block_size = arena->block_size; // original user requested size
+        U32 pad = AlignUp(MemoryBlock_Header_Size, align);
+        Assert(size <= (U32_Max - arena->page_size) - pad);
+        U32 needed = pad + size;
+        U32 new_block_size = arena->block_size;
         U32 new_commit_size = arena->commit_size;
-        if (size > block_size - MemoryBlock_Header_Size) {
-            block_size = AlignUpPow2(size + MemoryBlock_Header_Size,align);
-            new_commit_size = block_size;
+        if (needed > new_block_size) {
+            new_block_size = AlignUp(needed, arena->page_size);
+            new_commit_size = new_block_size;
         }
-        U32 commit_now_size = Max(new_commit_size, size);
 
-        MemoryBlockOptions mb_opts;
-        mb_opts.page_size   = arena->page_size;
-        mb_opts.commit_size = new_commit_size;
-        mb_opts.size        = block_size;
-        mb_opts.commit_now_size = commit_now_size;
-        MemoryBlock *new_block = memoryblock_allocate(&mb_opts);
+        MemoryBlock *new_block = memoryblock_allocate(new_block_size, new_commit_size, arena->page_size);
         // arena_offset tracks the virtual absolute position within an arena
         // This is not a contiguous space, we use it to free blocks when popping the arena
         new_block->arena_offset = block->arena_offset + block->size;
         new_block->prev         = block;
         arena->current          = new_block;
         block                   = new_block;
-        start_pos               = AlignUpPow2(block->first_free, align);
+        start_pos               = AlignUp(block->first_free, align);
     }
 
     U32 end_pos = start_pos + size;
@@ -106,7 +93,7 @@ void *arena_push(Arena *arena, U32 size, U32 align)
     U8 *mem = &block->base[start_pos];
     if (block->committed_end < end_pos) {
         // Commit more memory
-        U32 aligned_end_pos   = AlignUpPow2(start_pos + size, arena->page_size);
+        U32 aligned_end_pos   = AlignUp(start_pos + size, arena->page_size);
         U32 aligned_size      = RoundUpToMultiple(aligned_end_pos - block->committed_end, arena->commit_size);
         aligned_size = Min(aligned_size, block->size - block->committed_end);
 

src/base_core.h

@@ -104,11 +104,12 @@ typedef unsigned __int128 U128;
     (const type *)((const void *)(const unsigned char *)__mptr - offsetof(const type, member));                     \
   })
 
-#define AlignUpPow2(val, align)     (((val) + (align) - 1UL) & ~((align) - 1UL))
-#define AlignDownPow2(val, align)   ((val) & ~((align) - 1UL))
+// align needs to be pow2
+#define AlignUp(val, align)     (((val) + (align) - 1UL) & ~((align) - 1UL))
+#define AlignDown(val, align)   ((val) & ~((align) - 1UL))
 #define IsPow2(x)                   ((x) && !((x) & ((x) - 1UL)))
 #define IsPow2OrZero(x)             (!((x) & ((x) - 1UL)))
-#define IsAlignedToPow2(x, y)           (IsPow2(y) && !((x) & ((y) - 1UL)))
+#define IsAlignedTo(x, align)           (IsPow2(align) && !((x) & ((align) - 1UL)))
 
 #define GetBit(val, idx) (((val) >> (idx)) & 1)
 // Get n_bits starting at idx going left

src/main.c

@@ -1,8 +1,10 @@
 #include "base_core.h"
 #include "arena.h"
+#include "platform.h"
 #include <stdio.h>
 #include <stdlib.h>
 
+
 void arena_dump(const Arena *arena)
 {
     return;
@@ -51,8 +53,8 @@ static void arena_test_2(void)
     TRACEF("\n");
 
     ArenaOptions arena_params = {};
-    U32 arena_size = MB(1) + MemoryBlock_Header_Size;
-    U32 commit_size = KB(16);
+    U32 arena_size = MB(3) + MemoryBlock_Header_Size;
+    U32 commit_size = KB(72);
     U32 push_size = KB(32);
     arena_params.size = arena_size;
     arena_params.commit_size = commit_size;
@@ -66,7 +68,7 @@ static void arena_test_2(void)
 
     arena_dump(&arena);
     ArenaCheckpoint checkpoint = arena_temp_begin(&arena);
-    ptr = arena_push(&arena, MB(2), 8);
+    ptr = arena_push(&arena, MB(4)+3, 8);
     MemoryZero(ptr, MB(2));
 
     ptr = arena_push(&arena, commit_size * 3, 8);
@@ -92,8 +94,10 @@ static void arena_test_2(void)
     arena_reset_to(&arena, 1089);
     arena_dump(&arena);
     arena_reset_to(&arena, 0);
-    ptr = arena_push(&arena, MB(1), 8);
+    for (U64 i = 0; i < KB(1); i++) {
+        ptr = arena_push(&arena, MB(1)  + 3, 8);
         MemoryZero(ptr, MB(1));
+    }
     arena_free(&arena);
 }
 
@@ -120,9 +124,49 @@ static void arena_test_3(void)
     arena_free(&arena);
 }
 
+// Exercises the alignment-padding edge case: a fresh block has first_free at
+// MemoryBlock_Header_Size (64), so a large alignment pushes the user pointer
+// well past the header. Requesting (block_size - 64) bytes with page alignment
+// previously tripped the end_pos <= block->size assert because the alignment
+// gap was not accounted for. After the fix this should allocate a larger block.
+static void arena_test_4(void)
+{
+    TRACEF("\n");
+
+    U32 page_size = platform_page_size();
+    ArenaOptions arena_params = {};
+    U32 arena_size  = MB(1) + MemoryBlock_Header_Size;
+    U32 commit_size = KB(64);
+    arena_params.size = arena_size;
+    arena_params.commit_size = commit_size;
+    Arena arena;
+    arena_init(&arena, &arena_params);
+
+    // block_size is the page-aligned arena_size. Request nearly the whole block
+    // but with page alignment, so the alignment gap (a full page) forces a grow.
+    U32 block_size = arena.block_size;
+    U32 push_size  = block_size - MemoryBlock_Header_Size;
+
+    void *ptr = arena_push(&arena, push_size, page_size);
+    Assert(IsAlignedTo((uintptr_t)ptr, page_size));
+    MemoryZero(ptr, push_size);
+    printf("test_4: allocated %u bytes page-aligned at %p\n", push_size, ptr);
+    arena_dump(&arena);
+
+    // A second page-aligned push to make sure the cursor math still holds.
+    ptr = arena_push(&arena, page_size, page_size);
+    Assert(IsAlignedTo((uintptr_t)ptr, page_size));
+    MemoryZero(ptr, page_size);
+    printf("test_4: allocated %u bytes page-aligned at %p\n", page_size, ptr);
+
+    arena_free(&arena);
+}
+
 int main(int argc, char *argv[])
 {
-    arena_test_1();
+    // arena_test_1();
     arena_test_2();
-    arena_test_3();
+    // arena_test_3();
+    arena_test_4();
 }
+

src/platform.c

@@ -6,7 +6,7 @@
 #include <sys/mman.h>
 
 
-#define IsPageAligned(x) (IsAlignedToPow2(x, platform_page_size()))
+#define IsPageAligned(x) (IsAlignedTo(x, platform_page_size()))
 
 U32 platform_page_size(void)
 {