arena.h

#ifndef ARENA_H
#define ARENA_H

#include <assert.h>
#include <stdalign.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

struct ArenaBlock {
  uint8_t* data;
  size_t capacity;
  size_t used;
  struct ArenaBlock* next;
};

struct ArenaMark {
  struct ArenaBlock* block;
  size_t used;
};

struct Arena {
  struct ArenaBlock* first;
  struct ArenaBlock* current;
  size_t default_block_size;
};

void arena_init(struct Arena* arena, size_t default_block_size);
void* arena_alloc(struct Arena* arena, size_t size);
void arena_free(struct Arena* arena);

// Marks
struct ArenaMark arena_mark(struct Arena* arena);
void arena_reset(struct Arena* arena, struct ArenaMark mark);

static struct ArenaBlock* arena_block_create(size_t capacity);
static void arena_block_chain_free(struct ArenaBlock* block);

// Helpers
char* arena_strndup(struct Arena* arena, const char* s, size_t n);

static size_t align_up(size_t n, size_t align);

#endif  // ARENA_H

#ifdef ARENA_IMPLEMENTATION

void arena_init(struct Arena* arena, size_t default_block_size) {
  arena->first = NULL;
  arena->current = NULL;
  arena->default_block_size = default_block_size;
}

void* arena_alloc(struct Arena* arena, size_t size) {
  assert(arena != NULL);
  assert(size > 0);

  size_t align = alignof(max_align_t);
  assert(size <= SIZE_MAX - (align - 1));

  size_t aligned_size = size + align - 1;

  if (arena->first == NULL) {
    size_t capacity = arena->default_block_size;
    if (capacity < aligned_size) {
      capacity = aligned_size;
    }
    arena->first = arena_block_create(capacity);
    if (arena->first == NULL) {
      return NULL;
    }
    arena->current = arena->first;
  }

  struct ArenaBlock* block = arena->current;
  assert(block != NULL);

  size_t start = align_up(block->used, align);

  if (start > block->capacity || size > block->capacity - start) {
    assert(block->next == NULL);
    size_t new_capacity = arena->default_block_size;
    // NOTE: Block starts with used = 0, but check alignment anyways in case
    // that changes over time
    if (new_capacity < aligned_size) {
      new_capacity = aligned_size;
    }
    struct ArenaBlock* new_block = arena_block_create(new_capacity);
    if (new_block == NULL) return NULL;

    block->next = new_block;
    arena->current = new_block;
    block = new_block;
    start = align_up(block->used, align);

    assert(start <= block->capacity);
    assert(size <= block->capacity - start);
  }

  void* ptr = block->data + start;
  block->used = start + size;

  return ptr;
}

void arena_free(struct Arena* arena) {
  assert(arena != NULL);

  arena_block_chain_free(arena->first);

  arena->first = NULL;
  arena->current = NULL;
}

struct ArenaMark arena_mark(struct Arena* arena) {
  assert(arena != NULL);

  struct ArenaMark mark = {
      .block = NULL,
      .used = 0,
  };
  if (arena->current != NULL) {
    mark.block = arena->current;
    mark.used = arena->current->used;
  }
  return mark;
}

void arena_reset(struct Arena* arena, struct ArenaMark mark) {
  assert(arena != NULL);

  if (mark.block == NULL) {
    arena_block_chain_free(arena->first);
    arena->first = NULL;
    arena->current = NULL;
    return;
  }

  struct ArenaBlock* tail = mark.block->next;
  mark.block->next = NULL;
  arena->current = mark.block;
  mark.block->used = mark.used;
  arena_block_chain_free(tail);
}

static struct ArenaBlock* arena_block_create(size_t capacity) {
  size_t align = alignof(max_align_t);
  size_t header_size = align_up(sizeof(struct ArenaBlock), align);
  struct ArenaBlock* block = malloc(header_size + capacity);
  if (block == NULL) return NULL;

  block->data = (uint8_t*)block + header_size;

  block->capacity = capacity;
  block->used = 0;
  block->next = NULL;
  return block;
}

static void arena_block_chain_free(struct ArenaBlock* block) {
  while (block != NULL) {
    struct ArenaBlock* next = block->next;  // Save before prev is freed
    free(block);
    block = next;
  }
}

char* arena_strndup(struct Arena* arena, const char* s, size_t n) {
  if (n == SIZE_MAX) return NULL;
  if (s == NULL && n > 0) return NULL;

  n = strnlen(s, n);

  char* dst = arena_alloc(arena, n + 1);
  if (dst == NULL) return NULL;
  if (n > 0) memcpy(dst, s, n);
  dst[n] = '\0';
  return dst;
}

static size_t align_up(size_t n, size_t align) {
  return (n + align - 1) & ~(align - 1);
}

#endif