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.

---
 mem/CMakeLists.txt  |  1 +
 mem/include/mem.h   | 11 +++++++++--
 mem/src/mem.c       | 12 ++++++++++++
 mem/test/mem_test.c |  1 +
 4 files changed, 23 insertions(+), 2 deletions(-)

(limited to 'mem')

diff --git a/mem/CMakeLists.txt b/mem/CMakeLists.txt
index 233d2be..e4b28c3 100644
--- a/mem/CMakeLists.txt
+++ b/mem/CMakeLists.txt
@@ -11,6 +11,7 @@ target_include_directories(mem PUBLIC
   include)
 
 target_link_libraries(mem
+  cassert
   list)
 
 target_compile_options(mem PRIVATE -Wall -Wextra)
diff --git a/mem/include/mem.h b/mem/include/mem.h
index bcff39f..892ea4f 100644
--- a/mem/include/mem.h
+++ b/mem/include/mem.h
@@ -66,8 +66,10 @@
 #define mem_clear(MEM) mem_clear_(&(MEM)->mem)
 
 /// Allocate a new chunk of N blocks.
-/// Return a pointer to the first block of the chunk, or 0 if there is no memory
-/// left.
+/// Return a pointer to the first block of the chunk.
+/// When there is no space left in the allocator, allocation can either trap
+/// (default) or gracefully return 0. Call mem_enable_traps() to toggle this
+/// behaviour.
 /// New chunks are conveniently zeroed out.
 #define mem_alloc(MEM, num_blocks) mem_alloc_(&(MEM)->mem, num_blocks)
 
@@ -87,6 +89,9 @@
 /// Return the total capacity of the allocator in bytes.
 #define mem_capacity(MEM) mem_capacity_(&(MEM)->mem)
 
+/// Set whether to trap when attempting to allocate beyond capacity.
+#define mem_enable_traps(MEM, enable) mem_enable_traps_(&(MEM)->mem, enable)
+
 /// Iterate over the used chunks of the allocator.
 ///
 /// The caller can use 'i' as the index of the current chunk.
@@ -134,6 +139,7 @@ typedef struct Memory {
   size_t   num_blocks;
   size_t   next_free_chunk;
   bool     dynamic; /// True if blocks and chunks are dynamically-allocated.
+  bool     trap;    /// Whether to trap when allocating beyond capacity.
   Chunk*   chunks;  /// Array of chunk information.
   uint8_t* blocks;  /// Array of blocks;
 } Memory;
@@ -159,3 +165,4 @@ void   mem_free_(Memory*, void** chunk_ptr);
 void*  mem_get_chunk_(const Memory*, size_t chunk_handle);
 size_t mem_get_chunk_handle_(const Memory*, const void* chunk);
 size_t mem_capacity_(const Memory*);
+void   mem_enable_traps_(Memory*, bool);
diff --git a/mem/src/mem.c b/mem/src/mem.c
index 056d947..2904035 100644
--- a/mem/src/mem.c
+++ b/mem/src/mem.c
@@ -1,5 +1,7 @@
 #include "mem.h"
 
+#include <cassert.h>
+
 #include <stdlib.h>
 #include <string.h>
 
@@ -13,6 +15,7 @@ bool mem_make_(
   mem->block_size_bytes = block_size_bytes;
   mem->num_blocks       = num_blocks;
   mem->next_free_chunk  = 0;
+  mem->trap             = true;
 
   // Allocate chunks and blocks if necessary and zero them out.
   if (!chunks) {
@@ -107,6 +110,10 @@ void* mem_alloc_(Memory* mem, size_t num_blocks) {
     mem->next_free_chunk = mem->chunks[chunk_idx].next;
     return &mem->blocks[chunk_idx * mem->block_size_bytes];
   } else {
+    if (mem->trap) {
+      FAIL("Memory allocation failed, increase the allocator's capacity or "
+           "avoid fragmentation.");
+    }
     return 0; // Large-enough free chunk not found.
   }
 }
@@ -186,3 +193,8 @@ size_t mem_capacity_(const Memory* mem) {
   assert(mem);
   return mem->num_blocks * mem->block_size_bytes;
 }
+
+void mem_enable_traps_(Memory* mem, bool enable) {
+  assert(mem);
+  mem->trap = enable;
+}
diff --git a/mem/test/mem_test.c b/mem/test/mem_test.c
index 2f242c3..d3c43b9 100644
--- a/mem/test/mem_test.c
+++ b/mem/test/mem_test.c
@@ -67,6 +67,7 @@ TEST_CASE(mem_fill_then_free) {
 TEST_CASE(mem_allocate_beyond_max_size) {
   test_mem mem;
   mem_make(&mem);
+  mem_enable_traps(&mem, false);
 
   // Fully allocate the mem.
   for (int i = 0; i < NUM_BLOCKS; ++i) {
-- 
cgit v1.2.3