/*
 * 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)                                \
  {                                                                      \
    size_t i = (POOL)->pool.used;                                        \
    do {                                                                 \
      if ((POOL)->pool.block_info[i].used) {                             \
        __typeof__((POOL)->object[0])* ITER =                            \
            &(((__typeof__((POOL)->object[0])*)(POOL)->pool.blocks))[i]; \
        (void)ITER;                                                      \
        BODY;                                                            \
      }                                                                  \
      const size_t next = (POOL)->pool.block_info[i].next_used;          \
      if (next == i) {                                                   \
        break;                                                           \
      }                                                                  \
      i = next;                                                          \
    } while (true);                                                      \
  }

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

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

/// Memory pool.
///
/// 'head' and 'used' always points to a valid block (e.g., 0).
/// The implementation must check whether the head of the lists are used/free.
/// For example, iteration must stop when it finds the first unused block
/// (BlockInfo.used == 0).
typedef struct mempool {
  size_t     block_size_bytes;
  size_t     num_blocks;
  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.
  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*);