1 files changed, 16 insertions, 0 deletions

diff --git a/memstack/src/memstack.c b/memstack/src/memstack.c
index 0848afb..84131ef 100644
--- a/memstack/src/memstack.c
+++ b/memstack/src/memstack.c
@@ -58,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);
 
@@ -82,6 +97,7 @@ void* memstack_alloc_aligned(memstack* stack, size_t bytes, size_t alignment) {
   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);