import "@ffi/c"

AllocError :: enum {
    out_of_memory
}

#! allocator context and vtable are borrowed and must outlive allocations made through them.
Allocator :: struct {
    context ?@mut anyopaque
    vtable @AllocatorVTable
}

AllocatorVTable :: struct {
    alloc @proc(context ?@mut anyopaque, size usize, alignment usize) ?*mut u8
    realloc @proc(context ?@mut anyopaque, memory ?*mut u8, old_size usize, new_size usize, alignment usize) ?*mut u8
    free @proc(context ?@mut anyopaque, memory ?*mut u8, size usize, alignment usize) void
}

#! returned storage belongs to the caller; release it with raw_free using the same allocator, size, and alignment.
raw_alloc proc(
    allocator Allocator,
    size usize,
    alignment usize,
) ?*mut u8 {
    return allocator.vtable.alloc(allocator.context, size, alignment)
}

#! resize caller-owned storage with its original allocator and alignment.
#! on allocation failure with nonzero `new_size`, original storage remains caller-owned.
#! release returned storage with raw_free using the same allocator and alignment, with `new_size` as its size.
raw_realloc proc(
    allocator Allocator,
    memory ?*mut u8,
    old_size usize,
    new_size usize,
    alignment usize,
) ?*mut u8 {
    return allocator.vtable.realloc(allocator.context, memory, old_size, new_size, alignment)
}

#! release caller-owned storage with its original allocator, size, and alignment.
raw_free proc(
    allocator Allocator,
    memory ?*mut u8,
    size usize,
    alignment usize,
) void {
    allocator.vtable.free(allocator.context, memory, size, alignment)
}

eql proc($T type, left, right []T) bool {
    if (left.len != right.len) return false
    for (0..left.len) |i| {
        if (left[i] != right[i]) return false
    }
    return true
}

#! assign `value` to each slice element, preserving its typed value rather than repeating a byte.
#! an empty `destination` is a no-op.
set proc($T type, destination []mut T, value T) void {
    for (0..destination.len) |i| destination[i] = value
}

#! allocate memory for a slice of type `T` with `count` elements.
#! returned storage belongs to the caller; release it with free using the same allocator.
alloc proc($T type, allocator Allocator, count usize) []mut T ! AllocError {
    if (count == 0) return &[]

    element_size :: sizeof!(T)
    if (element_size == 0) return empty_slice(T, count)

    if count > divtrunc!(maxval!(usize), element_size) {
        return .out_of_memory
    }

    memory :: raw_alloc(allocator, count * element_size, alignof!(T))
    if memory |bytes| {
        pointer *mut T :: ptrcast!(bytes)
        return pointer[..count]
    }

    return .out_of_memory
}

#! reallocate memory for a slice of type `T` with `new_count` elements.
#! reallocating with `new_count == 0` will free the memory and return an empty slice.
#! note: memory must be reallocated with the same allocator that was used to allocate it.
#! the result belongs to the caller; release it with free using that allocator.
#! on failure the original slice remains valid and caller-owned.
realloc proc(
    $T type,
    allocator Allocator,
    memory []mut T,
    new_count usize,
) []mut T ! AllocError {
    if (new_count == memory.len) return memory

    if new_count == 0 {
        free(allocator, memory)
        return &[]
    }

    element_size :: sizeof!(T)
    if (element_size == 0) return empty_slice(T, new_count)
    if (new_count > divtrunc!(maxval!(usize), element_size)) return .out_of_memory

    old_memory ?*mut u8 := null
    old_size := 0
    if memory.len != 0 {
        old_memory = ptrcast!(memory.ptr)
        old_size = memory.len * element_size
    }

    resized :: raw_realloc(
        allocator,
        old_memory,
        old_size,
        new_count * element_size,
        alignof!(T),
    )

    if resized |bytes| {
        pointer *mut T :: ptrcast!(bytes)
        return pointer[..new_count]
    }

    return .out_of_memory
}

#! free memory allocated for a slice of type `T`.
#! note: memory must be freed with the same allocator that was used to allocate it.
free proc($T type, allocator Allocator, memory []T) void {
    if (memory.len == 0 or sizeof!(T) == 0) return
    raw_free(
        allocator,
        ptrcast!(constcast!(memory).ptr),
        memory.len * sizeof!(T),
        alignof!(T),
    )
}

#! create a logically nonempty slice of zero-sized elements without allocating.
#! precondition: sizeof!(T) must be 0; no element storage is accessed.
@hide:file
empty_slice proc($T type, count usize) []mut T {
    pointer *mut T :: ptrcast!((&empty_storage).ptr)
    return pointer[..count]
}

@hide:file
empty_storage [1]mut u64 := [0]

@hide:file
malloc_alignment usize :: 16 # note: aarch64-macos libc malloc alignment assumption.

@hide:file
power_of_two proc(value usize) bool {
    if (value == 0) return false
    current := value
    while current > 1 {
        half :: divtrunc!(current, 2)
        if (half * 2 != current) return false
        current = half
    }
    return true
}

@hide:file
c_alloc proc(_ ?@mut anyopaque, size usize, alignment usize) ?*mut u8 {
    if (power_of_two(alignment) == false) return null
    if (alignment <= malloc_alignment) return ptrcast!(c.malloc(c_ulong(size)))

    memory [1]mut ?*mut anyopaque := [null]
    status :: c.posix_memalign((&memory).ptr, c_ulong(alignment), c_ulong(size))
    if (status != 0) return null

    return ptrcast!(memory[0])
}

@hide:file
c_realloc proc(
    _ ?@mut anyopaque,
    memory ?*mut u8,
    old_size usize,
    new_size usize,
    alignment usize,
) ?*mut u8 {
    if (power_of_two(alignment) == false) return null

    if new_size == 0 {
        c.free(memory)
        return null
    }

    if memory |old_memory| {
        if alignment <= malloc_alignment {
            return ptrcast!(c.realloc(old_memory, c_ulong(new_size)))
        }

        new_memory :: c_alloc(null, new_size, alignment)
        if new_memory |new_bytes| {
            copy_size := old_size
            if (new_size < copy_size) copy_size = new_size
            memcopy!(new_bytes[..copy_size], old_memory[..copy_size])
            c.free(old_memory)
        }
        return new_memory
    }

    return c_alloc(null, new_size, alignment)
}

@hide:file
c_free proc(_ ?@mut anyopaque, memory ?*mut u8, _ usize, _ usize) void {
    c.free(memory)
}

@hide:file
c_vtable AllocatorVTable :: AllocatorVTable {
    alloc = c_alloc,
    realloc = c_realloc,
    free = c_free,
}

c_allocator Allocator :: Allocator {
    context = null,
    vtable = &c_vtable,
}