// pool.h - Fixed block size, free list based, memory pool
//
// Note: block_s must be >= sizeof(void*) and aligned to maintain proper
// pointer alignment for returned blocks.
//
// Thread safety: Not thread-safe. External synchronization required.
// Memory safety: pool_free validates pointer range, alignment, and detects
//   double-free via free-list walk (O(n) in number of free blocks).
//   Invalid frees are silently ignored.

#ifndef POOL_H
#define POOL_H

#include "type.h"

typedef struct pool_t {
  u32 block_n; // Block capacity
  u32 block_i; // Block count
  u32 block_s; // Block size in bytes
  void *data;   // Blocks
  void *free;   // Free chunk
} pool_t;

pool_t *pool_init(u32, u32);
void pool_fini(pool_t **);
void *pool_malloc(pool_t *);
void pool_free(pool_t *, void *);

u32 pool_capacity(const pool_t *);
u32 pool_size(const pool_t *);

#endif // POOL_H

#ifdef POOL_IMPLEMENTATION // IMPLEMENTATION

#include <stddef.h>
#include <assert.h>
#include <stdlib.h>

static void pool_free_list(uintptr_t *, u32, u32);

// pool_init()
//
// Initializes a memory pool and returns a pointer to the memory pool
//
// block_s - The size of a block in bytes
// block_n - The number of initial blocks

pool_t *pool_init(u32 block_s, u32 block_n) {
  if (block_n == 0)
    return NULL;
  if (block_s < sizeof(void *) || block_s % sizeof(void *) != 0)
    return NULL;

  pool_t *pool = (pool_t *)malloc(sizeof(pool_t));
  if (!pool)
    return NULL;

  // Set size and count
  pool->block_n = block_n;
  pool->block_s = block_s;
  pool->block_i = 0;

  if (block_n != 0 && block_s > SIZE_MAX / (size_t)block_n) {
    free(pool);
    return NULL;
  }
  pool->data = calloc(block_n, block_s);
  if (!pool->data) {
    free(pool);
    return NULL;
  }

  pool->free = pool->data;
  pool_free_list((uintptr_t *)pool->data, block_n, block_s);

  return pool;
}

// pool_malloc()
//
// Return a block of free memory
//
// pool - A pointer to the pool object

void *pool_malloc(pool_t *pool) {
  if (!pool)
    return NULL;
  uintptr_t *free = (uintptr_t *)pool->free;

  if (free) {
    pool->free = (uintptr_t *)*free;
    pool->block_i++;
  }

  return free;
}

// pool_free()
//
// Free the memory at the pointer p
//
// pool - A pointer to the pool object

void pool_free(pool_t *pool, void *p) {
  if (!pool || !p || !pool->data)
    return;
  // Validate pointer is within pool range
  u8 *base = (u8 *)pool->data;
  u8 *ptr = (u8 *)p;
  if ((u64)pool->block_n > UINT64_MAX / (u64)pool->block_s)
    return;
  u64 total = (u64)pool->block_n * pool->block_s;
  if (total > SIZE_MAX)
    return;
  u8 *end = base + (size_t)total;
  if (ptr < base || ptr >= end)
    return;
  // Validate pointer is aligned to block size
  if ((ptr - base) % pool->block_s != 0)
    return;
  // Check for double-free: walk free list to detect already-freed block
  {
    uintptr_t *cur = (uintptr_t *)pool->free;
    while (cur) {
      if ((uintptr_t)cur == (uintptr_t)p)
        return;
      cur = (uintptr_t *)*cur;
    }
  }
  // Convert the pointer into a usable type
  uintptr_t *free = (uintptr_t *)p;

  // Set the value at the address to the current front of the free list
  *free = (uintptr_t)pool->free;

  // Set the pool's free head to the newly freed pointer
  pool->free = (void *)free;

  if (pool->block_i > 0)
    pool->block_i--;
}

u32 pool_capacity(const pool_t *pool) { if (!pool) return 0; return pool->block_n; }

u32 pool_size(const pool_t *pool) { if (!pool) return 0; return pool->block_i; }

static void pool_free_list(uintptr_t *block, u32 block_n,
                             u32 block_s) {
  if (block_n == 0)
    return;

  u8 *ptr = (u8 *)block;
  for (u32 i = 0; i < block_n - 1; i++) {
    uintptr_t *tmp = (uintptr_t *)ptr;
    *tmp = (uintptr_t)ptr + block_s;
    ptr = ptr + block_s;
  }

  *(uintptr_t *)ptr = (uintptr_t)0;
}

void pool_fini(pool_t **pool) {
  if (pool && *pool) {
    free((*pool)->data);
    free(*pool);
    *pool = NULL;
  }
}

#endif // POOL_IMPLEMENTATION