3 files changed, 53 insertions, 0 deletions

diff --git a/memstack/include/memstack.h b/memstack/include/memstack.h
index 97a9d12..93cd2e6 100644
--- a/memstack/include/memstack.h
+++ b/memstack/include/memstack.h
@@ -32,6 +32,12 @@ void memstack_del(memstack*);
 /// Clear the stack.
 void memstack_clear(memstack*);
 
+/// Return the top of the stack.
+size_t memstack_watermark(const memstack*);
+
+/// Set the top of the stack.
+void memstack_set_watermark(memstack*, size_t watermark);
+
 /// Allocate a new block.
 ///
 /// Return null if the block does not fit in the remaining memory.
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);
diff --git a/memstack/test/memstack_test.c b/memstack/test/memstack_test.c
index 285cf46..5308be3 100644
--- a/memstack/test/memstack_test.c
+++ b/memstack/test/memstack_test.c
@@ -131,4 +131,35 @@ TEST_CASE(memstack_alloc_aligned) {
   memstack_del(&stack);
 }
 
+// Get and set the watermark.
+TEST_CASE(memstack_watermark) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  // Allocate N/2 ints.
+  for (int i = 0; i < NUM_INTS / 2; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+  }
+
+  const size_t watermark = memstack_watermark(&stack);
+
+  // Allocate the remaining N/2 ints.
+  for (int i = 0; i < NUM_INTS / 2; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+  }
+
+  // Now reset the watermark halfway through.
+  memstack_set_watermark(&stack, watermark);
+
+  // Allocate the remaining N/2 ints (again).
+  for (int i = 0; i < NUM_INTS / 2; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+  }
+
+  memstack_del(&stack);
+}
+
 int main() { return 0; }