6 files changed, 561 insertions, 0 deletions

diff --git a/mempool/CMakeLists.txt b/mempool/CMakeLists.txt
new file mode 100644
index 0000000..1ff4ff1
--- /dev/null
+++ b/mempool/CMakeLists.txt
@@ -0,0 +1,23 @@
+cmake_minimum_required(VERSION 3.16)
+
+project(mempool)
+
+# Library
+
+add_library(mempool
+  src/mempool.c)
+
+target_include_directories(mempool PUBLIC
+  include)
+
+target_compile_options(mempool PRIVATE -Wall -Wextra)
+
+# Test
+
+add_executable(mempool_test
+  test/mempool_test.c)
+
+target_link_libraries(mempool_test
+  mempool)
+
+target_compile_options(mempool_test PRIVATE -DUNIT_TEST -DNDEBUG -Wall -Wextra)
diff --git a/mempool/README.md b/mempool/README.md
new file mode 100644
index 0000000..ed2935e
--- /dev/null
+++ b/mempool/README.md
@@ -0,0 +1,20 @@
+# Mempool
+
+A memory pool implementation.
+
+Each block in the pool is tagged with a boolean value that indicates whether the
+block is free or in use. The allocator otherwise maintains little additional
+information. Therefore, some operations have linear-time complexity.
+Specifically:
+
+- Allocating a block scans the pool for the next free block in linear time.
+- Freeing a block runs in constant time.
+- Iterating over the pool's used blocks is linear over the number of blocks in
+  the pool, as opposed to just the number of used blocks.
+
+The allocator's internal data is also stored separately from the main array of
+blocks so that the block data remains tightly packed.
+
+For convenience of programming and debugging, free blocks in the mempool always
+remain zeroed out. If your data structures are valid when zeroed out, then you
+do not need to explicitly initialize them.
diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
new file mode 100644
index 0000000..41d56e5
--- /dev/null
+++ b/mempool/include/mempool.h
@@ -0,0 +1,91 @@
+#pragma once
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+
+/// Define a typed mempool of the given number of blocks.
+#define DEF_MEMPOOL(POOL, TYPE, NUM_BLOCKS)                                    \
+  typedef struct POOL {                                                        \
+    mempool pool;                                                              \
+    BlockInfo block_info[NUM_BLOCKS];                                          \
+    TYPE blocks[NUM_BLOCKS];                                                   \
+  } POOL;
+
+/// Create a new pool.
+#define mempool_make(POOL)                                                     \
+  {                                                                            \
+    assert(POOL);                                                              \
+    const size_t block_size = sizeof((POOL)->blocks[0]);                       \
+    const size_t num_blocks = sizeof((POOL)->blocks) / block_size;             \
+    mempool_make_(&(POOL)->pool, (POOL)->block_info, (POOL)->blocks,           \
+                  num_blocks, block_size);                                     \
+  }
+
+/// Allocate a new block.
+/// Return 0 if there is no memory left.
+#define mempool_alloc(POOL) mempool_alloc_(&(POOL)->pool)
+
+/// Free the block.
+/// The block pointer is conveniently set to 0.
+#define mempool_free(POOL, BLOCK_PTR)                                          \
+  assert(*BLOCK_PTR);                                                          \
+  mempool_free_(&(POOL)->pool, (void**)BLOCK_PTR)
+
+/// Return the ith block.
+/// The block must have been allocated.
+#define mempool_get_block(POOL, INDEX)                                         \
+  ((typeof((POOL)->blocks[0])*)mempool_get_block_(&(POOL)->pool, INDEX))
+
+/// Get the index to the given block.
+#define mempool_get_block_index(POOL, BLOCK_PTR)                               \
+  mempool_get_block_index_(&(POOL)->pool, BLOCK_PTR)
+
+/// Iterate over the used blocks of the pool.
+#define mempool_foreach(POOL, ITER, BODY)                                      \
+  for (size_t i = 0;                                                           \
+       i < (sizeof((POOL)->blocks) / sizeof(typeof((POOL)->blocks[0])));       \
+       ++i) {                                                                  \
+    if (!(POOL)->block_info[i].used) {                                         \
+      continue;                                                                \
+    }                                                                          \
+    typeof((POOL)->blocks[0])* ITER = &(POOL)->blocks[i];                      \
+    (void)ITER;                                                                \
+    BODY;                                                                      \
+  }
+
+typedef struct BlockInfo {
+  bool used;
+} BlockInfo;
+
+typedef struct mempool {
+  size_t block_size_bytes;
+  size_t num_blocks;
+  size_t next_free_block;
+  bool full;
+  BlockInfo* block_info;
+  uint8_t* blocks;
+} mempool;
+
+/// Create a pool allocator from user-provided memory.
+/// `BlockInfo` must hold at least `num_blocks` entries.
+/// `blocks` must be at least `num_blocks` * `block_size_bytes` bytes.
+/// All blocks are zeroed out for convenience.
+void mempool_make_(mempool*, BlockInfo*, void* blocks, size_t num_blocks,
+                   size_t block_size_bytes);
+
+/// Allocate a new block.
+/// Return 0 if there is no memory left.
+void* mempool_alloc_(mempool*);
+
+/// Free the block.
+/// The block pointer is conveniently set to 0.
+void mempool_free_(mempool*, void** block_ptr);
+
+/// Return the ith block.
+/// The block must have been allocated.
+void* mempool_get_block_(const mempool*, size_t block_index);
+
+/// Get the index to the given block.
+size_t mempool_get_block_index_(const mempool*, const void* block);
diff --git a/mempool/src/mempool.c b/mempool/src/mempool.c
new file mode 100644
index 0000000..b7e8705
--- /dev/null
+++ b/mempool/src/mempool.c
@@ -0,0 +1,85 @@
+#include "mempool.h"
+
+#include <string.h>
+
+static inline size_t min(size_t a, size_t b) { return a < b ? a : b; }
+
+void mempool_make_(mempool* pool, BlockInfo* block_info, void* blocks,
+                   size_t num_blocks, size_t block_size_bytes) {
+  assert(pool);
+  assert(block_info);
+  assert(blocks);
+  assert(num_blocks >= 1);
+  pool->block_size_bytes = block_size_bytes;
+  pool->num_blocks = num_blocks;
+  pool->next_free_block = 0;
+  pool->full = false;
+  pool->block_info = block_info;
+  pool->blocks = blocks;
+  memset(blocks, 0, num_blocks * block_size_bytes);
+}
+
+void* mempool_alloc_(mempool* pool) {
+  assert(pool);
+
+  if (pool->full) {
+    return 0;
+  }
+
+  // Allocate the block.
+  void* block = &pool->blocks[pool->next_free_block * pool->block_size_bytes];
+  pool->block_info[pool->next_free_block].used = true;
+
+  // Search for the next free block. If it does not exist, flag the pool full.
+  pool->full = true;
+  for (size_t i = 1; i < pool->num_blocks && i != 0; i++) {
+    pool->next_free_block = (pool->next_free_block + 1) % pool->num_blocks;
+    if (!pool->block_info[pool->next_free_block].used) {
+      pool->full = false;
+      break;
+    }
+  }
+
+  return block;
+}
+
+void mempool_free_(mempool* pool, void** block_ptr) {
+  assert(pool);
+  assert(block_ptr);
+
+  memset(*block_ptr, 0, pool->block_size_bytes);
+
+  const size_t block_index =
+      ((uint8_t*)*block_ptr - pool->blocks) / pool->block_size_bytes;
+  assert(block_index < pool->num_blocks);
+
+  // Disallow double-frees.
+  assert(pool->block_info[block_index].used);
+
+  pool->block_info[block_index].used = false;
+  if (pool->full) {
+    pool->next_free_block = block_index;
+    pool->full = false;
+  } else {
+    // Prefer to allocate blocks towards the start of the pool. This way, blocks
+    // should cluster around this area and the pool should offer better memory
+    // locality for used blocks.
+    pool->next_free_block = min(pool->next_free_block, block_index);
+  }
+
+  *block_ptr = 0;
+}
+
+void* mempool_get_block_(const mempool* pool, size_t block_index) {
+  assert(pool);
+  assert(block_index < pool->num_blocks);
+  assert(pool->block_info[block_index].used);
+  return pool->blocks + block_index * pool->block_size_bytes;
+}
+
+size_t mempool_get_block_index_(const mempool* pool, const void* block) {
+  assert(pool);
+  const size_t block_byte_index = (const uint8_t*)block - pool->blocks;
+  assert(block_byte_index % pool->block_size_bytes == 0);
+  return block_byte_index / pool->block_size_bytes;
+}
diff --git a/mempool/test/mempool_test.c b/mempool/test/mempool_test.c
new file mode 100644
index 0000000..d257922
--- /dev/null
+++ b/mempool/test/mempool_test.c
@@ -0,0 +1,157 @@
+#include "mempool.h"
+
+#include "test.h"
+
+#define NUM_BLOCKS 10
+
+DEF_MEMPOOL(test_pool, int, NUM_BLOCKS);
+
+static int count(test_pool* pool) {
+  int count = 0;
+  mempool_foreach(pool, n, { count++; });
+  return count;
+}
+
+static int sum(test_pool* pool) {
+  int sum = 0;
+  mempool_foreach(pool, n, { sum += *n; });
+  return sum;
+}
+
+// Create a pool.
+TEST_CASE(mempool_create) {
+  test_pool pool;
+  mempool_make(&pool);
+}
+
+// Allocate all N blocks.
+TEST_CASE(mempool_allocate_until_full) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    const int* block = mempool_alloc(&pool);
+    TEST_TRUE(block != 0);
+  }
+}
+
+// Allocate all N blocks, then free them.
+TEST_CASE(mempool_fill_then_free) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  int* blocks[NUM_BLOCKS] = {0};
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    blocks[i] = mempool_alloc(&pool);
+    TEST_TRUE(blocks[i] != 0);
+  }
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    mempool_free(&pool, &blocks[i]);
+    TEST_EQUAL(blocks[i], 0); // Pointer should be set to 0 on free.
+  }
+
+  TEST_EQUAL(count(&pool), 0);
+}
+
+// Attempt to allocate blocks past the maximum pool size.
+// The pool should handle the failed allocations gracefully.
+TEST_CASE(mempool_allocate_beyond_max_size) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  // Fully allocate the pool.
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    TEST_TRUE(mempool_alloc(&pool) != 0);
+  }
+
+  // Past the end.
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    TEST_EQUAL(mempool_alloc(&pool), 0);
+  }
+}
+
+// Free blocks should always remain zeroed out.
+// This tests the invariant right after creating the pool.
+TEST_CASE(mempool_zero_free_blocks_after_creation) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  const int zero = 0;
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    const int* block = (const int*)(pool.blocks) + i;
+    TEST_EQUAL(memcmp(block, &zero, sizeof(int)), 0);
+  }
+}
+
+// Free blocks should always remain zeroed out.
+// This tests the invariant after freeing a block.
+TEST_CASE(mempool_zero_free_block_after_free) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  int* val = mempool_alloc(&pool);
+  TEST_TRUE(val != 0);
+  *val = 177;
+
+  int* old_val = val;
+  mempool_free(&pool, &val); // val pointer is set to 0.
+  TEST_EQUAL(*old_val, 0);   // Block is zeroed out after free.
+}
+
+// Traverse an empty pool.
+TEST_CASE(mempool_traverse_empty) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  TEST_EQUAL(count(&pool), 0);
+}
+
+// Traverse a partially full pool.
+TEST_CASE(mempool_traverse_partially_full) {
+  const int N = NUM_BLOCKS / 2;
+
+  test_pool pool;
+  mempool_make(&pool);
+
+  for (int i = 0; i < N; ++i) {
+    int* val = mempool_alloc(&pool);
+    TEST_TRUE(val != 0);
+    *val = i + 1;
+  }
+
+  TEST_EQUAL(sum(&pool), N * (N + 1) / 2);
+}
+
+// Traverse a full pool.
+TEST_CASE(mempool_traverse_full) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    int* val = mempool_alloc(&pool);
+    TEST_TRUE(val != 0);
+    *val = i + 1;
+  }
+
+  TEST_EQUAL(sum(&pool), NUM_BLOCKS * (NUM_BLOCKS + 1) / 2);
+}
+
+// Get the ith (allocated) block.
+TEST_CASE(mempool_get_block) {
+  test_pool pool;
+  mempool_make(&pool);
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    int* block = mempool_alloc(&pool);
+    TEST_TRUE(block != 0);
+    *block = i;
+    TEST_EQUAL(mempool_get_block_index(&pool, block), (size_t)i);
+  }
+
+  for (int i = 0; i < NUM_BLOCKS; ++i) {
+    TEST_EQUAL(*mempool_get_block(&pool, i), i);
+  }
+}
+
+int main() { return 0; }
diff --git a/mempool/test/test.h b/mempool/test/test.h
new file mode 100644
index 0000000..fd8dc22
--- /dev/null
+++ b/mempool/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