From ff565e8d422c5e58558d1f7682ba0c9756e5be27 Mon Sep 17 00:00:00 2001
From: 3gg <3gg@shellblade.net>
Date: Mon, 14 Jul 2025 09:30:08 -0700
Subject: Add function to allocate aligned blocks on the memstack

---
 memstack/include/memstack.h   | 10 ++++++++++
 memstack/src/memstack.c       | 23 ++++++++++++++++++++++-
 memstack/test/memstack_test.c | 27 +++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/memstack/include/memstack.h b/memstack/include/memstack.h
index 9a8a7ee..97a9d12 100644
--- a/memstack/include/memstack.h
+++ b/memstack/include/memstack.h
@@ -33,13 +33,23 @@ void memstack_del(memstack*);
 void memstack_clear(memstack*);
 
 /// Allocate a new block.
+///
 /// Return null if the block does not fit in the remaining memory.
+///
 /// When there is no space left in the stack, allocation can either trap
 /// (default) or gracefully return null. Call mem_enable_traps() to toggle this
 /// behaviour.
+///
 /// Newly allocated blocks are conveniently zeroed out.
 void* memstack_alloc(memstack*, size_t bytes);
 
+/// Allocate a new aligned block.
+///
+/// An alignment of 0 is allowed for convenience and has the same effect as 1.
+///
+/// Has the same properties as memstack_alloc().
+void* memstack_alloc_aligned(memstack*, size_t bytes, size_t alignment);
+
 /// Return the stack's used size in bytes.
 size_t memstack_size(const memstack*);
 
diff --git a/memstack/src/memstack.c b/memstack/src/memstack.c
index 10d1e30..0848afb 100644
--- a/memstack/src/memstack.c
+++ b/memstack/src/memstack.c
@@ -5,6 +5,16 @@
 #include <stdlib.h>
 #include <string.h>
 
+static bool is_pow2_or_0(size_t x) { return (x & (x - 1)) == 0; }
+
+/// Align the given address to the next address that is a multiple of the
+/// alignment. If the given address is already aligned, return the address.
+static uint8_t* align(uint8_t* address, size_t alignment) {
+  assert(is_pow2_or_0(alignment));
+  const size_t mask = alignment - 1;
+  return (uint8_t*)(((uintptr_t)address + mask) & ~mask);
+}
+
 bool memstack_make(memstack* stack, size_t capacity, void* memory) {
   assert(stack);
   assert(capacity >= 1);
@@ -59,13 +69,24 @@ void* memstack_alloc(memstack* stack, size_t bytes) {
   }
 
   // Allocate the block.
-  uint8_t* block = stack->watermark;
+  uint8_t* const block = stack->watermark;
   stack->watermark += bytes;
   assert(memstack_size(stack) <= stack->capacity);
 
   return block;
 }
 
+void* memstack_alloc_aligned(memstack* stack, size_t bytes, size_t alignment) {
+  assert(stack);
+
+  uint8_t* const new_watermark = align(stack->watermark, alignment);
+  assert(new_watermark >= stack->watermark);
+  assert((size_t)(new_watermark - stack->base) <= stack->capacity);
+  stack->watermark = new_watermark;
+
+  return memstack_alloc(stack, bytes);
+}
+
 size_t memstack_size(const memstack* stack) {
   assert(stack);
   return stack->watermark - stack->base;
diff --git a/memstack/test/memstack_test.c b/memstack/test/memstack_test.c
index 5e9b493..285cf46 100644
--- a/memstack/test/memstack_test.c
+++ b/memstack/test/memstack_test.c
@@ -32,6 +32,8 @@ TEST_CASE(memstack_allocate_until_full) {
   }
 
   TEST_TRUE(memstack_size(&stack) == CAPACITY);
+
+  memstack_del(&stack);
 }
 
 // Allocate all N ints, then free them.
@@ -48,6 +50,8 @@ TEST_CASE(memstack_fill_then_free) {
   memstack_clear(&stack);
 
   TEST_EQUAL(memstack_size(&stack), 0);
+
+  memstack_del(&stack);
 }
 
 // Attempt to allocate blocks past the maximum stack size.
@@ -68,6 +72,8 @@ TEST_CASE(memstack_allocate_beyond_max_size) {
   }
 
   TEST_TRUE(memstack_size(&stack) == CAPACITY);
+
+  memstack_del(&stack);
 }
 
 // Free blocks should always remain zeroed out.
@@ -81,6 +87,8 @@ TEST_CASE(memstack_zero_free_blocks_after_creation) {
     TEST_TRUE(block != nullptr);
     TEST_EQUAL(*block, 0);
   }
+
+  memstack_del(&stack);
 }
 
 // Free blocks should always remain zeroed out.
@@ -102,6 +110,25 @@ TEST_CASE(memstack_zero_free_block_after_free) {
     TEST_TRUE(block != nullptr);
     TEST_EQUAL(*block, 0);
   }
+
+  memstack_del(&stack);
+}
+
+// Aligned allocations should be properly aligned.
+TEST_CASE(memstack_alloc_aligned) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  // -1 because the base address of the memory storage might be unaligned.
+  for (int i = 0; i < NUM_INTS - 1; ++i) {
+    const int* block =
+        memstack_alloc_aligned(&stack, sizeof(int), alignof(int));
+    TEST_TRUE(block != nullptr);
+    TEST_EQUAL(*block, 0);
+    TEST_EQUAL((uintptr_t)block % alignof(int), 0);
+  }
+
+  memstack_del(&stack);
 }
 
 int main() { return 0; }
-- 
cgit v1.2.3