/*
 * Block-based Memory Allocator.
 *
 * Clients should use the macros to define and use allocators. They make the API
 * type-safe.
 *
 * Like a pool/block-based allocator, this allocator stores data in fixed-size
 * blocks. However, this allocator also supports allocation of contiguous chunks
 * of a variable number of blocks.
 *
 * Chunk 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 typed memory allocator backed by a statically-allocated array.
#define DEF_MEM(MEM, TYPE, NUM_BLOCKS)                        \
  typedef struct MEM {                                        \
    Memory mem;                                               \
    Chunk  chunks[NUM_BLOCKS];                                \
    TYPE   blocks[NUM_BLOCKS];                                \
    /* For uniformity with the dynamically-allocated pool. */ \
    TYPE* object;                                             \
  } MEM;

/// Define a typed memory allocator backed by a dynamically-allocated array.
#define DEF_MEM_DYN(MEM, TYPE)                                               \
  typedef struct MEM {                                                       \
    Memory mem;                                                              \
    Chunk* chunks;                                                           \
    TYPE*  blocks;                                                           \
    /* Does not point anywhere. Storing a pointer here so that we can recall \
     * the type. */                                                          \
    TYPE* object;                                                            \
  } MEM;

/// Initialize a statically-backed memory allocator.
#define mem_make(MEM)                                                       \
  {                                                                         \
    assert(MEM);                                                            \
    const size_t block_size = sizeof((MEM)->blocks[0]);                     \
    const size_t num_blocks = sizeof((MEM)->blocks) / block_size;           \
    mem_make_(                                                              \
        &(MEM)->mem, (MEM)->chunks, (MEM)->blocks, num_blocks, block_size); \
  }

/// Initialize a dynamically-backed memory allocator.
#define mem_make_dyn(MEM, num_blocks, block_size) \
  mem_make_(&(MEM)->mem, 0, 0, num_blocks, block_size)

/// Destroy the allocator.
///
/// If the allocator is dynamically-backed, then this function frees the
/// underlying memory.
#define mem_del(MEM) mem_del_(&(MEM)->mem)

/// Clear the allocator.
///
/// This function frees all of the allocator's blocks. The resulting allocator
/// is as if it were newly created.
#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.
/// New chunks are conveniently zeroed out.
#define mem_alloc(MEM, num_blocks) mem_alloc_(&(MEM)->mem, num_blocks)

/// Free the chunk.
/// The chunk pointer is conveniently set to 0.
#define mem_free(MEM, CHUNK) mem_free_(&(MEM)->mem, (void**)CHUNK)

/// Return a pointer to a chunk given the chunk's handle.
/// The chunk must have been allocated.
#define mem_get_chunk(MEM, HANDLE) \
  ((__typeof__((MEM)->object[0])*)mem_get_chunk_(&(MEM)->mem, HANDLE))

/// Get the handle to the given chunk.
#define mem_get_chunk_handle(MEM, CHUNK_PTR) \
  mem_get_chunk_handle_(&(MEM)->mem, CHUNK_PTR)

/// Return the total capacity of the allocator in bytes.
#define mem_capacity(MEM) mem_capacity_(&(MEM)->mem)

/// Iterate over the used chunks of the allocator.
///
/// The caller can use 'i' as the index of the current chunk.
///
/// It is valid to mem_free() the chunk at each step of the iteration.
#define mem_foreach(MEM, ITER, BODY)                                \
  size_t i = 0;                                                     \
  do {                                                              \
    if ((MEM)->mem.chunks[i].used) {                                \
      __typeof__((MEM)->object[0])* ITER =                          \
          &(((__typeof__((MEM)->object[0])*)(MEM)->mem.blocks))[i]; \
      (void)ITER;                                                   \
      BODY;                                                         \
    }                                                               \
    i = (MEM)->chunks[i].next;                                      \
  } while (i);

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

/// Chunk information.
///
/// Every chunk represents a contiguous array of some number of blocks. The
/// allocator begins as one big unused chunk.
///
/// Allocation looks for a free chunk large enough to hold to requested number
/// of blocks. If the free chunk is larger than the requested chunk size, then
/// the requested chunk is carved out of the larger block.
///
/// Deallocation frees the chunk back and merges it with free neighbouring
/// chunks. Two free chunks are never contiguous in memory.
///
/// 'next' and 'prev' always point to a valid chunk (e.g., 0). Allocation stops
/// looking for free chunks when it loops over.
typedef struct Chunk {
  size_t num_blocks;
  size_t prev;
  size_t next;
  bool   used;
} Chunk;

typedef struct Memory {
  size_t   block_size_bytes;
  size_t   num_blocks;
  size_t   next_free_chunk;
  bool     dynamic; /// True if blocks and chunks are dynamically-allocated.
  Chunk*   chunks;  /// Array of chunk information.
  uint8_t* blocks;  /// Array of blocks;
} Memory;

/// Create a memory allocator.
///
/// 'chunks' and 'blocks' may be user-provided (statically-backed allocator) or
/// null (dynamically-backed allocator).
/// - If null, the allocator malloc()s the memory for them.
/// - If given:
///   - `chunks` must be at least `num_blocks` chunks.
///   - `blocks` must be at least `num_blocks` * `block_size_bytes` bytes.
///
/// All blocks are zeroed out for convenience.
bool mem_make_(
    Memory* mem, Chunk* chunks, void* blocks, size_t num_blocks,
    size_t block_size_bytes);

void   mem_del_(Memory*);
void   mem_clear_(Memory*);
void*  mem_alloc_(Memory*, size_t num_blocks);
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*);