From e153be0be2fb8df6656292daab3fa59963c76737 Mon Sep 17 00:00:00 2001
From: 3gg <3gg@shellblade.net>
Date: Tue, 13 Feb 2024 17:51:51 -0800
Subject: Let memory allocators trap by default when attempting to allocate
 beyond capacity.

---
 mempool/CMakeLists.txt      |  3 +++
 mempool/include/mempool.h   |  9 +++++++++
 mempool/src/mempool.c       | 11 +++++++++++
 mempool/test/mempool_test.c |  1 +
 4 files changed, 24 insertions(+)

(limited to 'mempool')

diff --git a/mempool/CMakeLists.txt b/mempool/CMakeLists.txt
index fe3e2a5..8c9dd30 100644
--- a/mempool/CMakeLists.txt
+++ b/mempool/CMakeLists.txt
@@ -10,6 +10,9 @@ add_library(mempool
 target_include_directories(mempool PUBLIC
   include)
 
+target_link_libraries(mempool PRIVATE
+  cassert)
+
 target_compile_options(mempool PRIVATE -Wall -Wextra)
 
 # Test
diff --git a/mempool/include/mempool.h b/mempool/include/mempool.h
index bd4d4dd..de9ea4f 100644
--- a/mempool/include/mempool.h
+++ b/mempool/include/mempool.h
@@ -65,6 +65,9 @@
 
 /// Allocate a new block.
 /// Return 0 if there is no memory left.
+/// When there is no space left in the pool, allocation can either trap
+/// (default) or gracefully return 0. Call mem_enable_traps() to toggle this
+/// behaviour.
 /// New blocks are conveniently zeroed out.
 #define mempool_alloc(POOL) mempool_alloc_(&(POOL)->pool)
 
@@ -86,6 +89,10 @@
 /// Return the total capacity of the mempool in bytes.
 #define mempool_capacity(POOL) mempool_capacity_(&(POOL)->pool)
 
+/// Set whether to trap when attempting to allocate beyond capacity.
+#define mempool_enable_traps(POOL, enable) \
+  mempool_enable_traps_(&(POOL)->pool, enable)
+
 /// Iterate over the used blocks of the pool.
 ///
 /// The caller can use 'i' as the index of the current block.
@@ -129,6 +136,7 @@ typedef struct mempool {
   size_t     head;    /// Points to the first block in the free list.
   size_t     used;    /// Points to the first block in the used list.
   bool       dynamic; /// True if blocks and info are dynamically-allocated.
+  bool       trap;    /// Whether to trap when allocating beyond capacity.
   BlockInfo* block_info;
   uint8_t*   blocks;
 } mempool;
@@ -154,3 +162,4 @@ void   mempool_free_(mempool*, void** block_ptr);
 void*  mempool_get_block_(const mempool*, size_t block_index);
 size_t mempool_get_block_index_(const mempool*, const void* block);
 size_t mempool_capacity_(const mempool*);
+void   mempool_enable_traps_(mempool*, bool);
diff --git a/mempool/src/mempool.c b/mempool/src/mempool.c
index 1100dad..b09038b 100644
--- a/mempool/src/mempool.c
+++ b/mempool/src/mempool.c
@@ -1,5 +1,7 @@
 #include "mempool.h"
 
+#include <cassert.h>
+
 #include <stdlib.h>
 #include <string.h>
 
@@ -24,6 +26,7 @@ bool mempool_make_(
   pool->num_blocks       = num_blocks;
   pool->head             = 0;
   pool->used             = 0;
+  pool->trap             = true;
 
   // Initialize blocks and block info.
   if (!block_info) {
@@ -74,6 +77,9 @@ void* mempool_alloc_(mempool* pool) {
 
   BlockInfo* head = &pool->block_info[pool->head];
   if (head->used) {
+    if (pool->trap) {
+      FAIL("mempool allocation failed, increase the pool's capacity.");
+    }
     return 0; // Pool is full.
   }
 
@@ -134,3 +140,8 @@ size_t mempool_capacity_(const mempool* pool) {
   assert(pool);
   return pool->num_blocks * pool->block_size_bytes;
 }
+
+void mempool_enable_traps_(mempool* pool, bool enable) {
+  assert(pool);
+  pool->trap = enable;
+}
diff --git a/mempool/test/mempool_test.c b/mempool/test/mempool_test.c
index d5ed1ea..6c48a2a 100644
--- a/mempool/test/mempool_test.c
+++ b/mempool/test/mempool_test.c
@@ -67,6 +67,7 @@ TEST_CASE(mempool_fill_then_free) {
 TEST_CASE(mempool_allocate_beyond_max_size) {
   test_pool pool;
   mempool_make(&pool);
+  mempool_enable_traps(&pool, false);
 
   // Fully allocate the pool.
   for (int i = 0; i < NUM_BLOCKS; ++i) {
-- 
cgit v1.2.3