//! 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.

/// Size must be positive. Alignment must be a positive power of two.

/// Failure leaves the allocation offset unchanged.

/// The caller must keep the storage live and must not reclaim live allocations.

export unsafe fn alloc(arena: *unsafe mut Arena, size: u32, alignment: u32) -> *mut opaque throws (AllocError) {

    assert alignment > 0;

    assert size > 0;

    assert (alignment & (alignment - 1)) == 0;

    let capacity = arena.data.len;

    if arena.offset >= capacity {

        throw AllocError::OutOfMemory;

    }

    let address = (&arena.data[0]) as u64;

    let mask = (alignment - 1) as u64;

    let remainder = ((address & mask) + (arena.offset as u64 & mask)) & mask;

    let padding = ((alignment as u64 - remainder) & mask) as u32;

    let available = capacity - arena.offset;

    if padding > available or size > available - padding {

        throw AllocError::OutOfMemory;

    }

    let aligned = arena.offset + padding;

    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.

/// The caller must keep the storage live and must commit each written prefix

/// before another allocation can use that prefix.

export unsafe fn remainingBuf(arena: *unsafe 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 [];

    }

    let bytes = try allocationSize(size, count);

    let ptr = try alloc(&mut *arena, bytes, 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(&mut *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(&mut *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 bytes = try allocationSize(size, count);

    let raw = try allocRaw(arena, bytes, alignment);

    return @sliceOf(raw, count);

}

/// Raw storage provider for allocation sessions.

export trait Alloc {

    /// Reserve one uninitialized object.

    unsafe fn (&mut Alloc) reserve(size: u32, alignment: u32) -> *unsafe mut opaque throws (AllocError);

    /// Reserve an uninitialized slice.

    unsafe fn (&mut Alloc) reserveSlice(itemSize: u32, itemAlignment: u32, count: u32) -> *unsafe mut [opaque] throws (AllocError);

}

instance Alloc for Arena {

    /// Reserve one uninitialized object in the arena.

    unsafe fn (arena: &mut Arena) reserve(size: u32, alignment: u32) -> *unsafe mut opaque throws (AllocError) {

        return try allocRaw(arena, size, alignment);

    }

    /// Reserve an uninitialized slice in the arena.

    unsafe fn (arena: &mut Arena) reserveSlice(size: u32, alignment: u32, count: u32) -> *unsafe mut [opaque] throws (AllocError) {

        return try allocRawSlice(arena, size, alignment, count);

    }

}

/// Compute a slice allocation size without unsigned multiplication overflow.

fn allocationSize(size: u32, count: u32) -> u32 throws (AllocError) {

    if size <> 0 and count > 4294967295 / size {

        throw AllocError::OutOfMemory;

    }

    return size * count;

}