//! Bump allocator for compiler data structures.
//!
//! Provides a simple arena-style allocator that allocates from a contiguous
//! byte buffer. Memory is never freed individually - the entire arena is
//! reset at once. This is ideal for compiler passes where all allocations
//! have the same lifetime.
@test mod tests;

/// Error thrown by allocator.
export union AllocError: Copy {
    /// Allocator is out of memory.
    OutOfMemory,
}

/// Bump allocator backed by a byte slice.
///
/// Allocations are made by advancing an offset pointer. Individual allocations
/// cannot be freed; instead, the entire arena is reset at once via [`reset`].
export record Arena {
    /// Backing storage.
    data: *mut [u8],
    /// Current allocation offset in bytes.
    offset: u32,
}

/// Create a new arena backed by the given byte slice.
export fn new(data: *mut [u8]) -> Arena {
    return Arena { data, offset: 0 };
}

/// Allocate `size` bytes with the given alignment.
///
/// Returns an opaque pointer to the allocated memory. Throws `AllocError` if
/// the arena is exhausted. The caller is responsible for casting to the
/// appropriate type and initializing the memory.
export fn alloc(arena: &mut Arena, size: u32, alignment: u32) -> *mut opaque throws (AllocError) {
    assert alignment > 0;
    assert size > 0;

    let aligned64 = (arena.offset as u64 + alignment as u64 - 1) & ~(alignment as u64 - 1);
    if aligned64 > arena.data.len as u64 or size as u64 > arena.data.len as u64 - aligned64 {
        throw AllocError::OutOfMemory;
    }
    let aligned = aligned64 as u32;
    let newOffset = aligned + size;

    let base: *mut u8 = &mut arena.data[aligned];
    set arena.offset = newOffset;

    return base as *mut opaque;
}

/// Reset the arena, allowing all memory to be reused.
///
/// Does not zero the memory.
export fn reset(arena: &mut Arena) {
    set arena.offset = 0;
}

/// Save the current arena state for later restoration.
export fn save(arena: &Arena) -> u32 {
    return arena.offset;
}

/// Restore the arena to a previously saved state, reclaiming all
/// allocations made since that point.
export fn restore(arena: &mut Arena, savedOffset: u32) {
    set arena.offset = savedOffset;
}

/// Returns the number of bytes currently allocated.
export fn used(arena: &Arena) -> u32 {
    return arena.offset;
}

/// Returns the number of bytes remaining in the arena.
export fn remaining(arena: &Arena) -> u32 {
    return arena.data.len as u32 - arena.offset;
}

/// Returns the remaining buffer as a mutable slice.
export fn remainingBuf(arena: &mut Arena) -> *mut [u8] {
    return &mut arena.data[arena.offset..];
}

/// Commits `size` bytes of allocation, advancing the offset.
/// Use after writing to the buffer returned by [`remainingBuf`].
export fn commit(arena: &mut Arena, size: u32) {
    set arena.offset += size;
}

/// Allocate a slice of `count` elements, each of `size` bytes with given alignment.
///
/// Returns a type-erased slice that should be cast to the appropriate `*[T]`.
/// The slice length is set to `count` (element count, not bytes).
/// Throws `AllocError` if the arena is exhausted.
export unsafe fn allocSlice(arena: &mut Arena, size: u32, alignment: u32, count: u32) -> *mut [opaque] throws (AllocError) {
    if count == 0 {
        return &mut [];
    }
    if size > 0xffffffff / count {
        throw AllocError::OutOfMemory;
    }
    let ptr = try alloc(arena, size * count, alignment);

    return @sliceOf(ptr, count);
}

/// Generic allocator interface.
///
/// Bundles an allocation function with an opaque context pointer so that
/// any allocation strategy (arena, free-list, mmap, pool) can be used
/// through a uniform interface. The `func` field is called with the context
/// pointer, a byte size and an alignment, and must return a pointer to
/// the allocated memory or panic on failure.
export record Allocator: Copy {
    /// Allocation function. Returns a pointer to `size` bytes
    /// aligned to `alignment`, or panics on failure.
    func: unsafe fn(*unsafe mut opaque, u32, u32) -> *mut opaque,
    /// Opaque context pointer passed to `func`.
    ctx: *unsafe mut opaque,
}

/// Create an allocator whose arena must outlive all uses of the allocator.
export unsafe fn arenaAllocator(arena: &mut Arena) -> Allocator {
    return Allocator {
        func: arenaAllocFn,
        ctx: (arena as *unsafe mut Arena) as *unsafe mut opaque,
    };
}

/// Arena allocation function conforming to the `Allocator` interface.
unsafe fn arenaAllocFn(ctx: *unsafe mut opaque, size: u32, alignment: u32) -> *mut opaque {
    let arena = ctx as *unsafe mut Arena;
    return try! alloc(arena, size, alignment);
}

/// Allocate raw storage that remains valid until the arena is reset.
export unsafe fn allocRaw(arena: &mut Arena, size: u32, alignment: u32) -> *unsafe mut opaque throws (AllocError) {
    let owner = try alloc(arena, size, alignment);
    let byte = owner as *mut u8;
    let raw: *unsafe mut u8 = &mut *byte;
    return raw as *unsafe mut opaque;
}

/// Allocate a raw slice that remains valid until the arena is reset.
export unsafe fn allocRawSlice(arena: &mut Arena, size: u32, alignment: u32, count: u32) -> *unsafe mut [opaque] throws (AllocError) {
    if count == 0 {
        return &mut [];
    }
    let raw = try allocRaw(arena, size * count, alignment);
    return @sliceOf(raw, count);
}