// arena — the arena allocator, in pure Flash.
//
// `ArenaAllocator` wraps a child allocator and hands out a real `Allocator`
// built from a Flash-implemented v-table, so everything allocated through
// it is released together by one `deinit`. Chunks of backing memory come
// from the child allocator and are chained in a list; allocation bumps a
// cursor inside the newest chunk. `free` rewinds the cursor only when the
// freed memory is the most recent allocation — freeing anything older is a
// no-op until `deinit`, and `resize` can grow in place only at the end of
// the newest chunk. This mirrors the observable contract of Zig's
// `std.heap.ArenaAllocator`, single-threaded.

use "base"
use "list"

pub const ArenaAllocator = struct {
    child base.Allocator,
    head ?*mut Node,

    const Self = #This()

    // One chunk of backing memory. `end` is the bump cursor: bytes below
    // it are handed out, bytes above are available. The newest chunk is
    // the head of the list; older chunks are never allocated from again.
    const Node = struct {
        data []mut u8,
        end usize,
        next ?*mut Node,
    }

    const vtable base.Allocator.VTable = .{ .alloc = allocImpl, .resize = resizeImpl, .remap = remapImpl, .free = freeImpl }

    // No allocation until the first request.
    pub fn init(child base.Allocator) Self {
        return .{ .child = child, .head = null }
    }

    // Release every chunk. The values allocated through the arena are
    // invalid afterwards.
    pub fn deinit(self Self) void {
        var it = self.head
        while it |node| {
            next := node.next
            self.child.free(node.data)
            self.child.destroy(node)
            it = next
        }
    }

    // The allocator handle. The arena must outlive every use of it.
    pub fn allocator(self *mut Self) base.Allocator {
        return .{ .ptr = self, .vtable = &vtable }
    }

    // Chain a fresh chunk of `size` bytes in front of the list. On failure
    // the arena is untouched.
    fn newChunk(self *mut Self, size usize) !*mut Node {
        node := try self.child.create(Node)
        errdefer self.child.destroy(node)
        data := try self.child.alloc(u8, size)
        node.* = .{ .data = data, .end = 0, .next = self.head }
        self.head = node
        return node
    }

    // The cursor advanced to the next address satisfying `alignment`,
    // expressed as an index into `data`.
    fn alignedIndex(node *mut Node, alignment base.Alignment) usize {
        addr := #intFromPtr(node.data.ptr)
        return alignment.forward(addr + node.end) - addr
    }

    fn allocImpl(ctx *mut anyopaque, len usize, alignment base.Alignment, ret_addr usize) ?[*]mut u8 {
        _ = ret_addr
        const arena *mut Self = #ptrCast(#alignCast(ctx))
        if arena.head |node| {
            idx := alignedIndex(node, alignment)
            if idx + len <= node.data.len {
                node.end = idx + len
                return node.data[idx..].ptr
            }
        }
        // The newest chunk is full (or there is none): chain a bigger one.
        // Doubling keeps the chunk count logarithmic; the slack guarantees
        // the aligned request fits whatever address the child returns.
        var size usize = 4096
        if arena.head |node| {
            if node.data.len * 2 > size {
                size = node.data.len * 2
            }
        }
        need := len + alignment.toByteUnits() - 1
        if size < need {
            size = need
        }
        node := arena.newChunk(size) catch return null
        idx := alignedIndex(node, alignment)
        node.end = idx + len
        return node.data[idx..].ptr
    }

    fn resizeImpl(ctx *mut anyopaque, memory []mut u8, alignment base.Alignment, new_len usize, ret_addr usize) bool {
        _ = alignment
        _ = ret_addr
        const arena *mut Self = #ptrCast(#alignCast(ctx))
        node := arena.head.?
        addr := #intFromPtr(node.data.ptr)
        if addr + node.end != #intFromPtr(memory.ptr) + memory.len {
            // Not the most recent allocation: it cannot move its end, but
            // shrinking in place is always fine.
            return new_len <= memory.len
        }
        if new_len <= memory.len {
            node.end -= memory.len - new_len
            return true
        }
        if node.end + (new_len - memory.len) <= node.data.len {
            node.end += new_len - memory.len
            return true
        }
        return false
    }

    fn remapImpl(ctx *mut anyopaque, memory []mut u8, alignment base.Alignment, new_len usize, ret_addr usize) ?[*]mut u8 {
        if resizeImpl(ctx, memory, alignment, new_len, ret_addr) {
            return memory.ptr
        }
        return null
    }

    fn freeImpl(ctx *mut anyopaque, memory []mut u8, alignment base.Alignment, ret_addr usize) void {
        _ = alignment
        _ = ret_addr
        const arena *mut Self = #ptrCast(#alignCast(ctx))
        node := arena.head.?
        addr := #intFromPtr(node.data.ptr)
        if addr + node.end == #intFromPtr(memory.ptr) + memory.len {
            node.end -= memory.len
        }
    }
}

test "an allocation goes through the arena" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    s := try a.alloc(u8, 4)
    try base.expectEqual(4, s.len)
    s[0] = 'x'
    try base.expectEqual('x', s[0])
}

test "deinit frees every chunk" {
    var arena = ArenaAllocator.init(base.testAlloc)
    a := arena.allocator()
    _ = try a.alloc(u8, 5000)
    _ = try a.alloc(u8, 100)
    arena.deinit()
}

test "free rewinds only the newest allocation" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    s1 := try a.alloc(u8, 16)
    p1 := #intFromPtr(s1.ptr)
    a.free(s1)
    s2 := try a.alloc(u8, 16)
    try base.expectEqual(p1, #intFromPtr(s2.ptr))
}

test "freeing an older allocation is a no-op until deinit" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    s1 := try a.alloc(u8, 8)
    s2 := try a.alloc(u8, 8)
    a.free(s1)
    s3 := try a.alloc(u8, 8)
    try base.expect(#intFromPtr(s3.ptr) > #intFromPtr(s2.ptr))
}

test "the newest allocation grows and shrinks in place" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    s := try a.alloc(u8, 8)
    grown := try a.realloc(s, 16)
    try base.expectEqual(#intFromPtr(s.ptr), #intFromPtr(grown.ptr))
    shrunk := try a.realloc(grown, 4)
    try base.expectEqual(#intFromPtr(s.ptr), #intFromPtr(shrunk.ptr))
}

test "allocations honor the requested alignment" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    _ = try a.alloc(u8, 3)
    w := try a.alloc(u64, 2)
    try base.expectEqual(0, #intFromPtr(w.ptr) % 8)
}

// The failing-allocator sweep body: `checkAllAllocationFailures` fails
// each child allocation in turn — both the node and the data of every
// `newChunk` — and each induced OOM must reach the caller as an error,
// with whatever was already chained still released by `deinit`.
fn arenaSweep(alloc base.Allocator) !void {
    var arena = ArenaAllocator.init(alloc)
    defer arena.deinit()
    a := arena.allocator()
    s := try a.alloc(u8, 100)
    s[0] = 'x'
    big := try a.alloc(u8, 5000)
    big[4999] = 'y'
    grown := try a.realloc(big, 9000)
    try base.expectEqual('y', grown[4999])
    try base.expectEqual('x', s[0])
}

test "arena allocation survives failure at every chunk allocation point" {
    try base.checkAllAllocationFailures(base.testAlloc, arenaSweep, .{})
}

test "a failed first chunk leaves the arena untouched" {
    var failing = base.FailingAllocator.init(base.testAlloc, .{ .fail_index = 0 })
    fa := failing.allocator()
    var arena = ArenaAllocator.init(fa)
    defer arena.deinit()
    a := arena.allocator()
    try base.expectError(error.OutOfMemory, a.alloc(u8, 8))
    try base.expect(arena.head == null)
}

test "a failed chunk chain leaves the arena consistent" {
    // Child allocations 0 and 1 build the first chunk; allocation 2 is the
    // second chunk's node, and 3 — its data — fails, driving newChunk's
    // errdefer. The live chunk must be untouched and still serving.
    var failing = base.FailingAllocator.init(base.testAlloc, .{ .fail_index = 3 })
    fa := failing.allocator()
    var arena = ArenaAllocator.init(fa)
    defer arena.deinit()
    a := arena.allocator()
    s := try a.alloc(u8, 8)
    s[0] = 'x'
    head_before := #intFromPtr(arena.head.?)
    end_before := arena.head.?.end
    try base.expectError(error.OutOfMemory, a.alloc(u8, 5000))
    try base.expectEqual(head_before, #intFromPtr(arena.head.?))
    try base.expectEqual(end_before, arena.head.?.end)
    t := try a.alloc(u8, 4)
    t[0] = 'y'
    try base.expectEqual('x', s[0])
}

test "a list grows through the arena" {
    var arena = ArenaAllocator.init(base.testAlloc)
    defer arena.deinit()
    a := arena.allocator()
    var xs list.List(u8) = .empty
    var v u8 = 0
    while v < 100 {
        try xs.append(a, v)
        v += 1
    }
    try base.expectEqual(100, xs.items.len)
    try base.expectEqual(99, xs.items[99])
}