#pragma once

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

/// Define a typed mempool 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];           \
  } POOL;

/// Create a new 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);                                                   \
  }

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

/// 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 < (sizeof((POOL)->blocks) / sizeof(__typeof__((POOL)->blocks[0]))); \
       ++i) {                                                                \
    if (!(POOL)->block_info[i].used) {                                       \
      continue;                                                              \
    }                                                                        \
    __typeof__((POOL)->blocks[0])* ITER = &(POOL)->blocks[i];                \
    (void)ITER;                                                              \
    BODY;                                                                    \
  }

typedef struct BlockInfo {
  bool used;
} BlockInfo;

typedef struct mempool {
  size_t     block_size_bytes;
  size_t     num_blocks;
  size_t     next_free_block;
  bool       full;
  BlockInfo* block_info;
  uint8_t*   blocks;
} mempool;

/// Create a pool allocator from user-provided memory.
/// `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.
void mempool_make_(
    mempool*, BlockInfo*, void* blocks, size_t num_blocks,
    size_t block_size_bytes);

/// Allocate a new block.
/// Return 0 if there is no memory left.
void* mempool_alloc_(mempool*);

/// Free the block.
/// The block pointer is conveniently set to 0.
void mempool_free_(mempool*, void** block_ptr);

/// Return the ith block.
/// The block must have been allocated.
void* mempool_get_block_(const mempool*, size_t block_index);

/// Get the index to the given block.
size_t mempool_get_block_index_(const mempool*, const void* block);