1 files changed, 9 insertions, 0 deletions

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