// Allocators in C.
#include <assert.h>
#include <stdalign.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// Allocation errors.
typedef enum {
    Error_None = 0,
    Error_OutOfMemory,
    Error_SizeOverflow,
} Error;

// Allocation result.
typedef struct {
    void* ptr;
    Error err;
} AllocResult;

// Allocator interface.
struct _Allocator {
    AllocResult (*alloc)(void* self, size_t size, size_t align);
    AllocResult (*realloc)(void* self, void* ptr, size_t oldSize, size_t newSize, size_t align);
    void (*free)(void* self, void* ptr, size_t size, size_t align);
};

typedef struct {
    const struct _Allocator* m;
    void* self;
} Allocator;

// Helper to prevent integer overflow during N-item allocation.
static inline size_t calcSize(size_t size, size_t count) {
    if (count > 0 && size > SIZE_MAX / count) {
        return 0;
    }
    return size * count;
}

// Allocates an item of type T.
// `AllocResult Alloc[T](Allocator a, T)`
#define Alloc(a, T) \
    ((a).m->alloc((a).self, sizeof(T), alignof(T)))

// Frees an item allocated with Alloc.
// Only accepts typed pointers, not void*.
// `void Free[T](Allocator a, T* ptr)`
#define Free(a, ptr) \
    ((a).m->free((a).self, (ptr), sizeof(*(ptr)), alignof(typeof(*(ptr)))))

// Allocates n items of type T.
// `AllocResult AllocN[T](Allocator a, T, size_t n)`
#define AllocN(a, T, n) \
    ((a).m->alloc((a).self, calcSize(sizeof(T), (n)), alignof(T)))

// Frees n items allocated with AllocN.
// Only accepts typed pointers, not void*.
// `void FreeN[T](Allocator a, T* ptr, size_t n)`
#define FreeN(a, ptr, n)               \
    ((a).m->free(                      \
        (a).self, (ptr),               \
        calcSize(sizeof(*(ptr)), (n)), \
        alignof(typeof(*(ptr)))))

// The libc allocator wrapper.
// Ignores alignment and treats zero-size allocations as errors.
AllocResult Libc_Alloc(void* self, size_t size, size_t align) {
    (void)self;
    (void)align;

    if (size == 0) return (AllocResult){NULL, Error_SizeOverflow};
    void* ptr = malloc(size);
    if (!ptr) return (AllocResult){NULL, Error_OutOfMemory};
    return (AllocResult){ptr, Error_None};
}

AllocResult Libc_Realloc(void* self, void* ptr, size_t oldSize,
                         size_t newSize, size_t align) {
    (void)self;
    (void)oldSize;
    (void)align;

    if (newSize == 0) {
        free(ptr);
        return (AllocResult){NULL, Error_SizeOverflow};
    }

    void* newPtr = realloc(ptr, newSize);
    if (!newPtr) return (AllocResult){ptr, Error_OutOfMemory};
    return (AllocResult){newPtr, Error_None};
}

void Libc_Free(void* self, void* ptr, size_t size, size_t align) {
    (void)self;
    (void)size;
    (void)align;
    free(ptr);
}

Allocator LibcAllocator(void) {
    static const struct _Allocator mtab = {
        .alloc = Libc_Alloc,
        .realloc = Libc_Realloc,
        .free = Libc_Free,
    };
    return (Allocator){.m = &mtab, .self = NULL};
}

// A simple arena allocator.
// Doesn't support reallocation.
typedef struct {
    uint8_t* buf;
    size_t cap;
    size_t offset;
} Arena;

Arena NewArena(uint8_t* buf, size_t cap) {
    return (Arena){.buf = buf, .cap = cap, .offset = 0};
}

static AllocResult Arena_Alloc(void* self, size_t size, size_t align) {
    Arena* arena = (Arena*)self;

    // 1. Calculate the alignment padding.
    if (size == 0) return (AllocResult){NULL, Error_SizeOverflow};
    uintptr_t currentPtr = (uintptr_t)arena->buf + arena->offset;
    uintptr_t alignedPtr = (currentPtr + (align - 1)) & ~(align - 1);
    size_t newOffset = (alignedPtr - (uintptr_t)arena->buf) + size;

    // 2. Check for errors.
    if (newOffset < arena->offset) {
        return (AllocResult){NULL, Error_SizeOverflow};
    }
    if (newOffset > arena->cap) {
        return (AllocResult){NULL, Error_OutOfMemory};
    }

    // 3. Commit the allocation.
    arena->offset = newOffset;
    return (AllocResult){(void*)alignedPtr, Error_None};
}

static void Arena_Free(void* self, void* ptr, size_t size, size_t align) {
    // Individual deallocations are no-ops.
    (void)self;
    (void)ptr;
    (void)size;
    (void)align;
}

static void Arena_Reset(Arena* arena) {
    arena->offset = 0;
}

Allocator Arena_Allocator(Arena* arena) {
    static const struct _Allocator mtab = {
        .alloc = Arena_Alloc,
        .free = Arena_Free,
    };
    return (Allocator){.m = &mtab, .self = arena};
}

void exampleLibc(void) {
    Allocator allocator = LibcAllocator();

    {
        // Allocate a single integer.
        AllocResult res = Alloc(allocator, int64_t);
        if (res.err != Error_None) {
            printf("Error: %d\n", res.err);
            return;
        }

        int64_t* x = res.ptr;
        *x = 42;

        Free(allocator, x);
    }

    {
        // Allocate an array of integers.
        size_t n = 100;
        AllocResult res = AllocN(allocator, int64_t, n);
        if (res.err != Error_None) {
            printf("Error: %d\n", res.err);
            return;
        }

        int64_t* arr = res.ptr;
        for (size_t i = 0; i < n; i++) {
            arr[i] = i + 1;
        }

        FreeN(allocator, arr, n);
    }
}

void exampleArena(void) {
    uint8_t buf[1024];
    Arena arena = NewArena(buf, sizeof(buf));
    Allocator allocator = Arena_Allocator(&arena);

    {
        // Allocate a single integer.
        AllocResult res = Alloc(allocator, int64_t);
        if (res.err != Error_None) {
            printf("Error: %d\n", res.err);
            return;
        }

        int64_t* x = res.ptr;
        *x = 42;
    }

    {
        // Allocate an array of integers.
        size_t n = 100;
        AllocResult res = AllocN(allocator, int64_t, n);
        if (res.err != Error_None) {
            printf("Error: %d\n", res.err);
            return;
        }

        int64_t* arr = res.ptr;
        for (size_t i = 0; i < n; i++) {
            arr[i] = i + 1;
        }
    }

    Arena_Reset(&arena);
}

int main(void) {
    exampleLibc();
    exampleArena();
}