#include <stdalign.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

#include "finalize.h"
#include "memregion.h"

enum {
    EXTRA_CAPACITY = 64,
};

union BlockHeader {
    struct Block
    {
        struct Block *next;
        char *free, *end;
    } block;
    max_align_t align;
};

static struct Block head, *tail = &head, *released;

static void
cleanup(void)
{
    for (struct Block *b = head.next, *next; b; b = next) {
        next = b->next;
        free(b);
    }
}

static void *
allocBlock(size_t capacity)
{
    static bool first = true;
    if (first) {
        first = false;
        finalizeRegister(cleanup);
    }

    size_t size = sizeof(union BlockHeader) + capacity;
    struct Block *b = malloc(size);
    if (!b) {
        fprintf(stderr, "allocBlock: out of memory\n");
        finalizeExit(1);
    }
    b->end = (char *)b + size;
    return b;
}

static size_t
roundUp(size_t a, size_t b)
{
    // will here be optimized to: (a + b - 1) & ~(b - 1);
    return (a + b - 1) / b * b;
}

void *
allocFromMemRegion(size_t numBytes)
{
    numBytes = roundUp(numBytes, alignof(max_align_t));

    while (tail->free + numBytes > tail->end) {
        if ((tail->next = released) != 0) {
            released = released->next;
        } else {
            tail->next = allocBlock(numBytes + EXTRA_CAPACITY);
        }
        tail = tail->next;
        tail->free = (char *)tail + sizeof(union BlockHeader);
        tail->next = 0;
    }
    tail->free += numBytes;
    return tail->free - numBytes;
}

void
releaseMemRegion(void)
{
    tail->next = released;
    released = head.next;
    head.next = 0;
    tail = &head;
}

void
printInfoMemRegion(void)
{
    printf("MemRegion:\n");
    printf("used blocks:\n");
    for (const struct Block *b = head.next; b; b = b->next) {
        size_t cap = b->end - (const char *)b - sizeof(union BlockHeader);
        size_t avail = b->end - b->free;
        printf("  block at address %p, cap = %zu, avail = %zu\n",
               (const void *)b, cap, avail);
    }
    printf("released blocks:\n");
    for (const struct Block *b = released; b; b = b->next) {
        size_t cap = b->end - (const char *)b - sizeof(union BlockHeader);
        size_t avail = b->end - b->free;
        printf("  block at address %p, cap = %zu, avail = %zu\n",
               (const void *)b, cap, avail);
    }
    printf("\n");
}