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

#include <webassembly.h>

typedef uint8_t  u8;
typedef uint32_t u32;
typedef uint64_t u64;
typedef int32_t  i32;
typedef int64_t  i64;
typedef float    f32;
typedef double   f64;

// First pass at memory allocation. We work with arenas that pull memory
// a page at a time. Right now just waste the ends of pages and don't allow
// allocations greater than a page.
//@OPTIMIZE: There's wasted memory between _heap_base to the end of the first page.

#define page_size 65536

extern void _grow();

void *allocate_page() {
  int page = __builtin_wasm_grow_memory(1);
  _grow(); // Updates the wrapper arrays so console_log works.
  return (uint8_t*)0 + (page * page_size);
}

typedef struct PageFooter {
  struct PageFooter *prev;
} PageFooter;

typedef struct {
  void *page;
  size_t used;
} Arena;

static void* free_pages = NULL;

void *arena_push_size(Arena* arena, u32 size) {
  if (size > page_size - sizeof(PageFooter)) {
    console_error("size > maximum allocatable");
  }
  if (!arena->page || arena->used + size > page_size) {
    // Allocate a new page
    void *new_page = NULL;
    if (free_pages != NULL) {
      new_page = free_pages;
      free_pages = *(void**)new_page;
    } else {
      new_page = allocate_page();
    }
    PageFooter *footer = (PageFooter*)((u8*)new_page + (page_size - sizeof(PageFooter)));
    footer->prev = arena->page;
    arena->page = new_page;
    arena->used = 0;
  }
  if (size > (page_size - sizeof(PageFooter)) - arena->used) {
    console_error("No room for some reason!");
  }
  void *result = arena->page + arena->used;
  arena->used += size;
  return result;
}

void arena_free(Arena* arena) {
  void *page = arena->page;
  while(page) {
    PageFooter *footer = (PageFooter*)((u8*)page + (page_size - sizeof(PageFooter)));
    *(void**)page = free_pages;
    free_pages = page;
    page = footer->prev;
  }
  arena->page = NULL;
  arena->used = 0;
}

export void test() {
  console_log("Running Tests\n");
  Arena arena;
  Arena arena_two;
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  void *mem = arena_push_size(&arena, 1024);
  console_log("push 1024: %i", (int)mem);
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  mem = arena_push_size(&arena, 1024);
  console_log("push 1024: %i", (int)mem);
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  console_log("free arena");
  arena_free(&arena);
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  mem = arena_push_size(&arena, 1024);
  console_log("push 1024: %i", (int)mem);
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  for (int i=0; i<12; i++) {
    mem = arena_push_size(&arena, page_size-100);
    console_log("push %i: %i", page_size-100, (int)mem);
    console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  }
  console_log("Arena2{base: %i, used: %i}", (int)arena_two.page, arena_two.used);
  mem = arena_push_size(&arena_two, 1024);
  console_log("push 1024 to arena 2: %i", (int)mem);
  console_log("Arena2{base: %i, used: %i}", (int)arena_two.page, arena_two.used);
  console_log("free arena 1");
  arena_free(&arena);
  console_log("Arena{base: %i, used: %i}", (int)arena.page, arena.used);
  mem = arena_push_size(&arena_two, 65436);
  console_log("push 65436 to arena 2: %i", (int)mem);
  console_log("Arena2{base: %i, used: %i}", (int)arena_two.page, arena_two.used);
}