/*
 * Block/Pool Allocator.
 *
 * Clients should use the macros to define and use pools. They make the API
 * type-safe.
 *
 * The pool allocator works on one big chunk of memory, which can be statically
 * or dynamically allocated. This chunk is divided into fixed-sized blocks.
 * Allocation/deallocation happens with block granularity.
 *
 * Block information is stored in a separate array so that client data is
 * contiguous in the main pool of memory and better cached.
 */
#pragma once

#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>

/// Define a statically-allocated, typed pool 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];                             \
    /* For uniformity with the dynamically-allocated pool. */ \
    TYPE* object;                                             \
  } POOL;

/// Define a dynamically-allocated, typed pool.
#define DEF_MEMPOOL_DYN(POOL, TYPE)                                          \
  typedef struct POOL {                                                      \
    mempool pool;                                                            \
    /* Does not point anywhere. Storing a pointer here so that we can recall \
     * the type. */                                                          \
    TYPE* object;                                                            \
  } POOL;

/// Initialize a statically-allocated 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);                                                   \
  }

/// Initialize a dynamically-allocated pool.
#define mempool_make_dyn(POOL, num_blocks, block_size) \
  mempool_make_(&(POOL)->pool, 0, 0, num_blocks, block_size)

/// Destroy the pool.
///
/// If the pool is dynamically allocated, then this function frees its memory.
#define mempool_del(POOL) mempool_del_(&(POOL)->pool)

/// Clear the pool.
///
/// This function frees all of the pool's blocks. The resulting pool is as if it
/// were newly created.
#define mempool_clear(POOL) mempool_clear_(&(POOL)->pool)

/// Allocate a new block.
/// Return 0 if there is no memory left.
/// New blocks are conveniently zeroed out.
#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)->object[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)

/// Return the total capacity of the mempool in bytes.
#define mempool_capacity(POOL) mempool_capacity_(&(POOL)->pool)

/// Iterate over the used blocks of the pool.
///
/// The caller can use 'i' as the index of the current block.
///
/// It is valid to mempool_free() the object at each step of the iteration.
#define mempool_foreach(POOL, ITER, BODY)                            \
  for (size_t i = 0; i < (POOL)->pool.num_blocks; ++i) {             \
    if (!(POOL)->pool.block_info[i].used) {                          \
      continue;                                                      \
    }                                                                \
    __typeof__((POOL)->object[0])* ITER =                            \
        &(((__typeof__((POOL)->object[0])*)(POOL)->pool.blocks))[i]; \
    (void)ITER;                                                      \
    BODY;                                                            \
  }

// -----------------------------------------------------------------------------

typedef struct BlockInfo {
  size_t next; /// For free blocks, points to the next free block.
  bool   used;
} BlockInfo;

typedef struct mempool {
  size_t     block_size_bytes;
  size_t     num_blocks;
  size_t     head;    /// Points to the first block in the free list.
  bool       dynamic; /// True if blocks and info are dynamically-allocated.
  BlockInfo* block_info;
  uint8_t*   blocks;
} mempool;

/// Create a pool allocator.
///
/// 'BlockInfo' and 'blocks' may be user-provided (static pool) or null (dynamic
/// pool).
/// - If null, the pool malloc()s the memory for them.
/// - If given:
///   - `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.
bool mempool_make_(
    mempool*, BlockInfo*, void* blocks, size_t num_blocks,
    size_t block_size_bytes);

void   mempool_del_(mempool*);
void   mempool_clear_(mempool*);
void*  mempool_alloc_(mempool*);
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*);