#include "mem.h"

#include "test.h"

#define NUM_BLOCKS 10

DEF_MEM(test_mem, int, NUM_BLOCKS)

static int count(test_mem* mem) {
  int count = 0;
  mem_foreach(mem, n, { count++; });
  return count;
}

static int sum(test_mem* mem) {
  int sum = 0;
  mem_foreach(mem, n, { sum += *n; });
  return sum;
}

// Create a statically-backed allocator.
TEST_CASE(mem_create) {
  test_mem mem;
  mem_make(&mem);
}

// Create a dynamically-backed allocator.
TEST_CASE(mem_create_dyn) {
  DEF_MEM_DYN(dyn_mem, int);

  dyn_mem mem;
  mem_make_dyn(&mem, NUM_BLOCKS, sizeof(int));
}

// Allocate N chunks of 1 block each.
TEST_CASE(mem_fully_allocate) {
  test_mem mem;
  mem_make(&mem);

  for (int i = 0; i < NUM_BLOCKS; ++i) {
    const int* block = mem_alloc(&mem, 1);
    TEST_TRUE(block != 0);
  }
}

// Allocate N chunks of 1 block each, then free them.
TEST_CASE(mem_fill_then_free) {
  test_mem mem;
  mem_make(&mem);

  int* blocks[NUM_BLOCKS] = {0};
  for (int i = 0; i < NUM_BLOCKS; i++) {
    blocks[i] = mem_alloc(&mem, 1);
    TEST_TRUE(blocks[i] != 0);
  }

  for (int i = 0; i < NUM_BLOCKS; i++) {
    mem_free(&mem, &blocks[i]);
    TEST_EQUAL(blocks[i], 0); // Pointer should be set to 0 on free.
  }

  TEST_EQUAL(count(&mem), 0);
}

// Attempt to allocate blocks past the maximum allocator size.
// The allocator should handle the failed allocations gracefully.
TEST_CASE(mem_allocate_beyond_max_size) {
  test_mem mem;
  mem_make(&mem);

  // Fully allocate the mem.
  for (int i = 0; i < NUM_BLOCKS; ++i) {
    TEST_TRUE(mem_alloc(&mem, 1) != 0);
  }

  // Past the end.
  for (int i = 0; i < NUM_BLOCKS; ++i) {
    TEST_EQUAL(mem_alloc(&mem, 1), 0);
  }
}

// Free blocks should always remain zeroed out.
// This tests the invariant right after creating the allocator.
TEST_CASE(mem_zero_free_blocks_after_creation) {
  test_mem mem;
  mem_make(&mem);

  const int zero = 0;
  for (int i = 0; i < NUM_BLOCKS; ++i) {
    const int* block = (const int*)(mem.blocks) + i;
    TEST_EQUAL(memcmp(block, &zero, sizeof(int)), 0);
  }
}

// Free blocks should always remain zeroed out.
// This tests the invariant after freeing a block.
TEST_CASE(mem_zero_free_block_after_free) {
  test_mem mem;
  mem_make(&mem);

  int* val = mem_alloc(&mem, 1);
  TEST_TRUE(val != 0);
  *val = 177;

  int* old_val = val;
  mem_free(&mem, &val);    // val pointer is set to 0.
  TEST_EQUAL(*old_val, 0); // Block is zeroed out after free.
}

// Traverse an empty allocator.
TEST_CASE(mem_traverse_empty) {
  test_mem mem;
  mem_make(&mem);

  TEST_EQUAL(count(&mem), 0);
}

// Traverse a partially full allocator.
TEST_CASE(mem_traverse_partially_full) {
  const int N = NUM_BLOCKS / 2;

  test_mem mem;
  mem_make(&mem);

  for (int i = 0; i < N; ++i) {
    int* val = mem_alloc(&mem, 1);
    TEST_TRUE(val != 0);
    *val = i + 1;
  }

  TEST_EQUAL(sum(&mem), (N) * (N + 1) / 2);
}

// Traverse a full allocator.
TEST_CASE(mem_traverse_full) {
  test_mem mem;
  mem_make(&mem);

  for (int i = 0; i < NUM_BLOCKS; ++i) {
    int* val = mem_alloc(&mem, 1);
    TEST_TRUE(val != 0);
    *val = i + 1;
  }

  TEST_EQUAL(sum(&mem), (NUM_BLOCKS) * (NUM_BLOCKS + 1) / 2);
}

// Get the ith (allocated) chunk.
TEST_CASE(mem_get_block) {
  test_mem mem;
  mem_make(&mem);

  for (int i = 0; i < NUM_BLOCKS; ++i) {
    int* block = mem_alloc(&mem, 1);
    TEST_TRUE(block != 0);
    *block = i;
    TEST_EQUAL(mem_get_chunk_handle(&mem, block), (size_t)i);
  }

  for (int i = 0; i < NUM_BLOCKS; ++i) {
    TEST_EQUAL(*mem_get_chunk(&mem, i), i);
  }
}

// Test merging.
// 1. Allocate chunks of variable sizes.
// 2. Free them in a different order.
// 3. Then we should be able to allocate 1 chunk of N blocks.
TEST_CASE(mem_fragmentation) {
  test_mem mem;
  mem_make(&mem);

  int* blocks[NUM_BLOCKS] = {0};
  int  next_block         = 0;

#define ALLOC(num_blocks)                           \
  blocks[next_block] = mem_alloc(&mem, num_blocks); \
  TEST_TRUE(blocks[next_block] != 0);               \
  next_block++;

#define FREE(block_idx) mem_free(&mem, &blocks[block_idx])

  // 5 total allocations of variable chunk sizes.
  ALLOC(2); // 2;  idx = 0
  ALLOC(3); // 5;  idx = 1
  ALLOC(1); // 6;  idx = 2
  ALLOC(3); // 9;  idx = 3
  ALLOC(1); // 10; idx = 4

  // Free the 5 allocations in a different order.
  FREE(1);
  FREE(3);
  FREE(4);
  FREE(2);
  FREE(0);

  // Should be able to allocate 1 chunk of N blocks.
  const void* chunk = mem_alloc(&mem, NUM_BLOCKS);
  TEST_TRUE(chunk != 0);
}

// Clear and re-use an allocator.
TEST_CASE(mem_clear_then_reuse) {
  test_mem mem;
  mem_make(&mem);

  // Allocate chunks, contents not important.
  for (int i = 0; i < NUM_BLOCKS; ++i) {
    int* chunk = mem_alloc(&mem, 1);
    TEST_TRUE(chunk != 0);
  }

  mem_clear(&mem);

  // Allocate chunks and assign values 0..N.
  for (int i = 0; i < NUM_BLOCKS; ++i) {
    int* chunk = mem_alloc(&mem, 1);
    TEST_TRUE(chunk != 0);
    *chunk = i + 1;
  }

  TEST_EQUAL(sum(&mem), NUM_BLOCKS * (NUM_BLOCKS + 1) / 2);
}

// Stress test.
//
// 1. Allocate the mem, either fully or partially. If fully, attempt to
//    allocate some items past the end.
//
// 2. Free all allocated items in some random order.

int main() { return 0; }