8 files changed, 472 insertions, 2 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 04e0e4e..d1fb3ab 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -9,9 +9,10 @@ add_subdirectory(cstring)
 add_subdirectory(error)
 add_subdirectory(filesystem)
 add_subdirectory(list)
-add_subdirectory(mem)
 add_subdirectory(log)
+add_subdirectory(mem)
 add_subdirectory(mempool)
+add_subdirectory(memstack)
 add_subdirectory(plugin)
 add_subdirectory(random)
 add_subdirectory(timer)
diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index 245173b..04d0313 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -106,8 +106,8 @@
 /// It is valid to mempool_free() the object at each step of the iteration.
 #define mempool_foreach(POOL, ITER, BODY)                                  \
   {                                                                        \
-    size_t i = (POOL)->pool.used;                                          \
     if ((POOL)->pool.num_used_blocks > 0) {                                \
+      size_t i = (POOL)->pool.used;                                        \
       do {                                                                 \
         if ((POOL)->pool.block_info[i].used) {                             \
           __typeof__((POOL)->object[0])* ITER =                            \
diff --git a/memstack/CMakeLists.txt b/memstack/CMakeLists.txt
new file mode 100644
index 0000000..9ad1aa1
--- /dev/null
+++ b/memstack/CMakeLists.txt
@@ -0,0 +1,30 @@
+cmake_minimum_required(VERSION 3.5)
+
+project(memstack)
+
+set(CMAKE_C_STANDARD 23)
+set(CMAKE_C_STANDARD_REQUIRED On)
+set(CMAKE_C_EXTENSIONS Off)
+
+# Library
+
+add_library(memstack
+  src/memstack.c)
+
+target_include_directories(memstack PUBLIC
+  include)
+
+target_link_libraries(memstack PRIVATE
+  cassert)
+
+target_compile_options(memstack PRIVATE -Wall -Wextra)
+
+# Test
+
+add_executable(memstack_test
+  test/memstack_test.c)
+
+target_link_libraries(memstack_test
+  memstack)
+
+target_compile_options(memstack_test PRIVATE -DUNIT_TEST -DNDEBUG -Wall -Wextra)
diff --git a/memstack/README.md b/memstack/README.md
new file mode 100644
index 0000000..7eb950e
--- /dev/null
+++ b/memstack/README.md
@@ -0,0 +1,15 @@
+# Mempool
+
+A memory pool of fixed-sized blocks with O(1) allocation and deallocation.
+
+Each block in the pool is tagged with a boolean value that indicates whether the
+block is free or in use, as well as a pointer to the next free/used block.
+Blocks thus form two lists, a free list and a used list. The allocator
+maintains the two lists when allocating/deallocating blocks.
+
+The allocator's internal data is stored separately from the block data so that
+the data remains tightly packed.
+
+Free blocks in the mempool always remain zeroed out for convenience of
+programming and debugging. If the application's data structures are valid when
+zeroed out, then they do not need to be explicitly initialized.
diff --git a/memstack/include/memstack.h b/memstack/include/memstack.h
new file mode 100644
index 0000000..9a8a7ee
--- /dev/null
+++ b/memstack/include/memstack.h
@@ -0,0 +1,50 @@
+/*
+ * Stack-based allocator.
+ */
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+/// Stack-based allocator.
+typedef struct memstack {
+  size_t   capacity;  // Total size available.
+  uint8_t* base;      // Base pointer to memory.
+  uint8_t* watermark; // Pointer to next free area of memory.
+  bool     owned;     // True if memory is owned by the memstack.
+  bool     trap;      // Whether to trap when allocating beyond capacity.
+} memstack;
+
+/// Create a stack-based allocator.
+///
+/// `stack` may be user-provided or null.
+/// - If null, the allocator malloc()s the memory for them.
+/// - If given, `stack` must be at least `capacity` bytes.
+///
+/// The memory is zeroed out for convenience.
+bool memstack_make(memstack*, size_t capacity, void* memory);
+
+/// Destroy the stack.
+///
+/// If the allocator owns the memory, then this function frees it.
+void memstack_del(memstack*);
+
+/// Clear the stack.
+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);
+
+/// Return the stack's used size in bytes.
+size_t memstack_size(const memstack*);
+
+/// Return the stack's total capacity in bytes.
+size_t memstack_capacity(const memstack*);
+
+/// Set whether to trap when attempting to allocate beyond capacity.
+void memstack_enable_traps(memstack*, bool);
diff --git a/memstack/src/memstack.c b/memstack/src/memstack.c
new file mode 100644
index 0000000..10d1e30
--- /dev/null
+++ b/memstack/src/memstack.c
@@ -0,0 +1,82 @@
+#include "memstack.h"
+
+#include <cassert.h>
+
+#include <stdlib.h>
+#include <string.h>
+
+bool memstack_make(memstack* stack, size_t capacity, void* memory) {
+  assert(stack);
+  assert(capacity >= 1);
+
+  // Allocate memory if not user-provided.
+  uint8_t* stack_memory = memory;
+  if (!stack_memory) {
+    stack_memory = calloc(1, capacity);
+    if (stack_memory == nullptr) {
+      return false;
+    }
+  }
+  assert(stack_memory);
+
+  stack->capacity  = capacity;
+  stack->base      = stack_memory;
+  stack->watermark = stack_memory;
+  stack->owned     = (stack_memory != memory);
+  stack->trap      = true;
+
+  return true;
+}
+
+void memstack_del(memstack* stack) {
+  assert(stack);
+
+  if (stack->owned && (stack->base != nullptr)) {
+    free(stack->base);
+    stack->base  = nullptr;
+    stack->owned = false;
+  }
+
+  stack->capacity  = 0;
+  stack->watermark = stack->base;
+}
+
+void memstack_clear(memstack* stack) {
+  assert(stack);
+
+  stack->watermark = stack->base;
+  memset(stack->base, 0, stack->capacity);
+}
+
+void* memstack_alloc(memstack* stack, size_t bytes) {
+  assert(stack);
+
+  if ((memstack_size(stack) + bytes) > stack->capacity) {
+    if (stack->trap) {
+      FAIL("memstack allocation failed, increase the stack's capacity.");
+    }
+    return nullptr; // Block does not fit in remaining memory.
+  }
+
+  // Allocate the block.
+  uint8_t* block = stack->watermark;
+  stack->watermark += bytes;
+  assert(memstack_size(stack) <= stack->capacity);
+
+  return block;
+}
+
+size_t memstack_size(const memstack* stack) {
+  assert(stack);
+  return stack->watermark - stack->base;
+}
+
+size_t memstack_capacity(const memstack* stack) {
+  assert(stack);
+  return stack->capacity;
+}
+
+void memstack_enable_traps(memstack* stack, bool enable) {
+  assert(stack);
+  stack->trap = enable;
+}
diff --git a/memstack/test/memstack_test.c b/memstack/test/memstack_test.c
new file mode 100644
index 0000000..5e9b493
--- /dev/null
+++ b/memstack/test/memstack_test.c
@@ -0,0 +1,107 @@
+#include "memstack.h"
+
+#include "test.h"
+
+#define NUM_INTS 10
+#define CAPACITY (NUM_INTS * sizeof(int))
+
+// Create and destroy a statically-backed stack.
+TEST_CASE(memstack_create) {
+  int memory[CAPACITY];
+
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, memory);
+  memstack_del(&stack);
+}
+
+// Create and destroy a dynamically-backed stack.
+TEST_CASE(mem_create_dyn) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+  memstack_del(&stack);
+}
+
+// Allocate all N ints.
+TEST_CASE(memstack_allocate_until_full) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  for (int i = 0; i < NUM_INTS; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+  }
+
+  TEST_TRUE(memstack_size(&stack) == CAPACITY);
+}
+
+// Allocate all N ints, then free them.
+TEST_CASE(memstack_fill_then_free) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  int* blocks[NUM_INTS] = {nullptr};
+  for (int i = 0; i < NUM_INTS; ++i) {
+    blocks[i] = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(blocks[i] != nullptr);
+  }
+
+  memstack_clear(&stack);
+
+  TEST_EQUAL(memstack_size(&stack), 0);
+}
+
+// Attempt to allocate blocks past the maximum stack size.
+// The stack should handle the failed allocations gracefully.
+TEST_CASE(memstack_allocate_beyond_max_size) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+  memstack_enable_traps(&stack, false);
+
+  // Fully allocate the stack.
+  for (int i = 0; i < NUM_INTS; ++i) {
+    TEST_TRUE(memstack_alloc(&stack, sizeof(int)) != nullptr);
+  }
+
+  // Past the end.
+  for (int i = 0; i < NUM_INTS; ++i) {
+    TEST_EQUAL(memstack_alloc(&stack, sizeof(int)), nullptr);
+  }
+
+  TEST_TRUE(memstack_size(&stack) == CAPACITY);
+}
+
+// Free blocks should always remain zeroed out.
+// This tests the invariant right after creating the stack.
+TEST_CASE(memstack_zero_free_blocks_after_creation) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  for (int i = 0; i < NUM_INTS; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+    TEST_EQUAL(*block, 0);
+  }
+}
+
+// Free blocks should always remain zeroed out.
+// This tests the invariant after clearing the stack and allocating a new block.
+TEST_CASE(memstack_zero_free_block_after_free) {
+  memstack stack = {0};
+  memstack_make(&stack, CAPACITY, nullptr);
+
+  for (int i = 0; i < NUM_INTS; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+    TEST_EQUAL(*block, 0);
+  }
+
+  memstack_clear(&stack);
+
+  for (int i = 0; i < NUM_INTS; ++i) {
+    const int* block = memstack_alloc(&stack, sizeof(int));
+    TEST_TRUE(block != nullptr);
+    TEST_EQUAL(*block, 0);
+  }
+}
+
+int main() { return 0; }
diff --git a/memstack/test/test.h b/memstack/test/test.h
new file mode 100644
index 0000000..fd8dc22
--- /dev/null
+++ b/memstack/test/test.h
@@ -0,0 +1,185 @@
+// SPDX-License-Identifier: MIT
+#pragma once
+
+#ifdef UNIT_TEST
+
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#if defined(__DragonFly__) || defined(__FreeBSD__) || defined(__FreeBSD_kernel__) ||     \
+    defined(__NetBSD__) || defined(__OpenBSD__)
+#define USE_SYSCTL_FOR_ARGS 1
+// clang-format off
+#include <sys/types.h>
+#include <sys/sysctl.h>
+// clang-format on
+#include <unistd.h>        // getpid
+#endif
+
+struct test_file_metadata;
+
+struct test_failure {
+        bool present;
+        const char *message;
+        const char *file;
+        int line;
+};
+
+struct test_case_metadata {
+        void (*fn)(struct test_case_metadata *, struct test_file_metadata *);
+        struct test_failure failure;
+        const char *name;
+        struct test_case_metadata *next;
+};
+
+struct test_file_metadata {
+        bool registered;
+        const char *name;
+        struct test_file_metadata *next;
+        struct test_case_metadata *tests;
+};
+
+struct test_file_metadata __attribute__((weak)) * test_file_head;
+
+#define SET_FAILURE(_message)                                                             \
+        metadata->failure = (struct test_failure) {                                       \
+                .message = _message, .file = __FILE__, .line = __LINE__, .present = true, \
+        }
+
+#define TEST_EQUAL(a, b)                                                                 \
+        do {                                                                             \
+                if ((a) != (b)) {                                                        \
+                        SET_FAILURE(#a " != " #b);                                       \
+                        return;                                                          \
+                }                                                                        \
+        } while (0)
+
+#define TEST_TRUE(a)                                                                     \
+        do {                                                                             \
+                if (!(a)) {                                                              \
+                        SET_FAILURE(#a " is not true");                                  \
+                        return;                                                          \
+                }                                                                        \
+        } while (0)
+
+#define TEST_STREQUAL(a, b)                                                              \
+        do {                                                                             \
+                if (strcmp(a, b) != 0) {                                                 \
+                        SET_FAILURE(#a " != " #b);                                       \
+                        return;                                                          \
+                }                                                                        \
+        } while (0)
+
+#define TEST_CASE(_name)                                                                  \
+        static void __test_h_##_name(struct test_case_metadata *,                         \
+                                     struct test_file_metadata *);                        \
+        static struct test_file_metadata __test_h_file;                                   \
+        static struct test_case_metadata __test_h_meta_##_name = {                        \
+            .name = #_name,                                                               \
+            .fn = __test_h_##_name,                                                       \
+        };                                                                                \
+        static void __attribute__((constructor(101))) __test_h_##_name##_register(void) { \
+                __test_h_meta_##_name.next = __test_h_file.tests;                         \
+                __test_h_file.tests = &__test_h_meta_##_name;                             \
+                if (!__test_h_file.registered) {                                          \
+                        __test_h_file.name = __FILE__;                                    \
+                        __test_h_file.next = test_file_head;                              \
+                        test_file_head = &__test_h_file;                                  \
+                        __test_h_file.registered = true;                                  \
+                }                                                                         \
+        }                                                                                 \
+        static void __test_h_##_name(                                                     \
+            struct test_case_metadata *metadata __attribute__((unused)),                  \
+            struct test_file_metadata *file_metadata __attribute__((unused)))
+
+extern void __attribute__((weak)) (*test_h_unittest_setup)(void);
+/// Run defined tests, return true if all tests succeeds
+/// @param[out] tests_run if not NULL, set to whether tests were run
+static inline void __attribute__((constructor(102))) run_tests(void) {
+        bool should_run = false;
+#ifdef USE_SYSCTL_FOR_ARGS
+        int mib[] = {
+                CTL_KERN,
+#if defined(__NetBSD__) || defined(__OpenBSD__)
+                KERN_PROC_ARGS,
+                getpid(),
+                KERN_PROC_ARGV,
+#else
+                KERN_PROC,
+                KERN_PROC_ARGS,
+                getpid(),
+#endif
+        };
+        char *arg = NULL;
+        size_t arglen;
+        sysctl(mib, sizeof(mib) / sizeof(mib[0]), NULL, &arglen, NULL, 0);
+        arg = malloc(arglen);
+        sysctl(mib, sizeof(mib) / sizeof(mib[0]), arg, &arglen, NULL, 0);
+#else
+        FILE *cmdlinef = fopen("/proc/self/cmdline", "r");
+        char *arg = NULL;
+        int arglen;
+        fscanf(cmdlinef, "%ms%n", &arg, &arglen);
+        fclose(cmdlinef);
+#endif
+        for (char *pos = arg; pos < arg + arglen; pos += strlen(pos) + 1) {
+                if (strcmp(pos, "--unittest") == 0) {
+                        should_run = true;
+                        break;
+                }
+        }
+        free(arg);
+
+        if (!should_run) {
+                return;
+        }
+
+        if (&test_h_unittest_setup) {
+                test_h_unittest_setup();
+        }
+
+        struct test_file_metadata *i = test_file_head;
+        int failed = 0, success = 0;
+        while (i) {
+                fprintf(stderr, "Running tests from %s:\n", i->name);
+                struct test_case_metadata *j = i->tests;
+                while (j) {
+                        fprintf(stderr, "\t%s ... ", j->name);
+                        j->failure.present = false;
+                        j->fn(j, i);
+                        if (j->failure.present) {
+                                fprintf(stderr, "failed (%s at %s:%d)\n", j->failure.message,
+                                        j->failure.file, j->failure.line);
+                                failed++;
+                        } else {
+                                fprintf(stderr, "passed\n");
+                                success++;
+                        }
+                        j = j->next;
+                }
+                fprintf(stderr, "\n");
+                i = i->next;
+        }
+        int total = failed + success;
+        fprintf(stderr, "Test results: passed %d/%d, failed %d/%d\n", success, total,
+                failed, total);
+        exit(failed == 0 ? EXIT_SUCCESS : EXIT_FAILURE);
+}
+
+#else
+
+#include <stdbool.h>
+
+#define TEST_CASE(name) static void __attribute__((unused)) __test_h_##name(void)
+
+#define TEST_EQUAL(a, b)                                                                 \
+        (void)(a);                                                                       \
+        (void)(b)
+#define TEST_TRUE(a) (void)(a)
+#define TEST_STREQUAL(a, b)                                                              \
+        (void)(a);                                                                       \
+        (void)(b)
+
+#endif