1 files changed, 9 insertions, 2 deletions

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);