1 files changed, 38 insertions, 1 deletions

diff --git a/memstack/src/memstack.c b/memstack/src/memstack.c
index 10d1e30..84131ef 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);
@@ -48,6 +58,21 @@ void memstack_clear(memstack* stack) {
   memset(stack->base, 0, stack->capacity);
 }
 
+size_t memstack_watermark(const memstack* stack) {
+  assert(stack);
+  return stack->watermark - stack->base;
+}
+
+void memstack_set_watermark(memstack* stack, size_t watermark) {
+  assert(stack);
+  const bool fits = (watermark < stack->capacity);
+  if (stack->trap && !fits) {
+    FAIL("memstack watermark update failed, bad watermark");
+  }
+  assert(fits);
+  stack->watermark = stack->base + watermark;
+}
+
 void* memstack_alloc(memstack* stack, size_t bytes) {
   assert(stack);
 
@@ -59,13 +84,25 @@ 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->capacity -= (new_watermark - stack->watermark);
+  stack->watermark = new_watermark;
+
+  return memstack_alloc(stack, bytes);
+}
+
 size_t memstack_size(const memstack* stack) {
   assert(stack);
   return stack->watermark - stack->base;