#include "memstack.h"

#include <cassert.h>

#include <stdlib.h>
#include <string.h>

static bool is_pow2_or_0(size_t x) { return (x & (x - 1)) == 0; }

/// Align the given address to the next address that is a multiple of the
/// alignment. If the given address is already aligned, return the address.
static uint8_t* align(uint8_t* address, size_t alignment) {
  assert(is_pow2_or_0(alignment));
  const size_t mask = alignment - 1;
  return (uint8_t*)(((uintptr_t)address + mask) & ~mask);
}

bool memstack_make(memstack* stack, size_t capacity, void* memory) {
  assert(stack);
  assert(capacity >= 1);

  // Allocate memory if not user-provided.
  uint8_t* stack_memory = memory;
  if (!stack_memory) {
    stack_memory = calloc(1, capacity);
    if (stack_memory == nullptr) {
      return false;
    }
  }
  assert(stack_memory);

  stack->capacity  = capacity;
  stack->base      = stack_memory;
  stack->watermark = stack_memory;
  stack->owned     = (stack_memory != memory);
  stack->trap      = true;

  return true;
}

void memstack_del(memstack* stack) {
  assert(stack);

  if (stack->owned && (stack->base != nullptr)) {
    free(stack->base);
    stack->base  = nullptr;
    stack->owned = false;
  }

  stack->capacity  = 0;
  stack->watermark = stack->base;
}

void memstack_clear(memstack* stack) {
  assert(stack);

  stack->watermark = stack->base;
  memset(stack->base, 0, stack->capacity);
}

void* memstack_alloc(memstack* stack, size_t bytes) {
  assert(stack);

  if ((memstack_size(stack) + bytes) > stack->capacity) {
    if (stack->trap) {
      FAIL("memstack allocation failed, increase the stack's capacity.");
    }
    return nullptr; // Block does not fit in remaining memory.
  }

  // Allocate the block.
  uint8_t* const block = stack->watermark;
  stack->watermark += bytes;
  assert(memstack_size(stack) <= stack->capacity);

  return block;
}

void* memstack_alloc_aligned(memstack* stack, size_t bytes, size_t alignment) {
  assert(stack);

  uint8_t* const new_watermark = align(stack->watermark, alignment);
  assert(new_watermark >= stack->watermark);
  assert((size_t)(new_watermark - stack->base) <= stack->capacity);
  stack->watermark = new_watermark;

  return memstack_alloc(stack, bytes);
}

size_t memstack_size(const memstack* stack) {
  assert(stack);
  return stack->watermark - stack->base;
}

size_t memstack_capacity(const memstack* stack) {
  assert(stack);
  return stack->capacity;
}

void memstack_enable_traps(memstack* stack, bool enable) {
  assert(stack);
  stack->trap = enable;
}