// code_heap.zig - Code heap management
// Manages the JIT-compiled code heap: allocation, block tracking,
// remembered sets for GC, scan flags, and mark bits.

const std = @import("std");
const builtin = @import("builtin");

const code_blocks_mod = @import("code_blocks.zig");
const free_list = @import("free_list.zig");
const layouts = @import("layouts.zig");
const mark_bits = @import("mark_bits.zig");
const segments = @import("segments.zig");
const write_barrier = @import("write_barrier.zig");

const Cell = layouts.Cell;
const CodeBlock = code_blocks_mod.CodeBlock;

pub const CodeHeap = struct {
    seg: ?*segments.Segment,
    free_list: ?*free_list.FreeListAllocator = null, // typed allocator for JIT
    safepoint_page: Cell,
    // Code heap address range (for callstack walking)
    code_start: Cell = 0,
    code_size: Cell = 0,
    // Source of truth for live code block addresses (kept sorted).
    // New inserts go to pending_blocks (O(1) append) and are merged
    // into the sorted list lazily via flushPending().
    all_blocks_sorted: std.ArrayList(Cell) = .empty,
    pending_blocks: std.ArrayListUnmanaged(Cell) = .empty,
    // Memory allocator used for code heap metadata (mark bits, hash maps, etc.)
    allocator: ?std.mem.Allocator = null,
    // Remembered sets for GC - track code blocks that may reference young objects
    remembered_sets: write_barrier.CodeHeapRememberedSets,
    // Scan flags for code heap GC (skip blocks without literals/code pointers)
    scan_literals: ?std.DynamicBitSet = null,
    scan_code_ptrs: ?std.DynamicBitSet = null,
    // Uninitialized blocks: code blocks allocated but not yet relocated.
    // Maps block address -> literals Cell for deferred initialization.
    uninitialized_blocks: std.AutoArrayHashMapUnmanaged(Cell, Cell) = .{},
    // Scratch map reused during compaction to avoid repeated allocations.
    uninitialized_blocks_scratch: std.AutoArrayHashMapUnmanaged(Cell, Cell) = .{},
    marks: ?*mark_bits.MarkBits = null,

    const Self = @This();

    pub fn allocate(self: *Self, size: Cell) ?*CodeBlock {
        const alloc = self.free_list orelse return null;
        const aligned_size = layouts.alignCell(size, layouts.data_alignment);
        if (alloc.allocate(aligned_size)) |addr| {
            const block: *CodeBlock = @ptrFromInt(addr);
            // Add to all_blocks for codeBlockForAddress to find
            self.addToAllBlocks(addr);
            return block;
        }
        return null;
    }

    fn addToAllBlocks(self: *Self, block_addr: Cell) void {
        const al = self.allocator orelse return;
        self.pending_blocks.append(al, block_addr) catch @panic("OOM");
    }

    // Merge pending inserts into the sorted list. Call before bulk iteration
    // or binary-search-dependent operations.
    pub fn flushPending(self: *Self) void {
        if (self.pending_blocks.items.len == 0) return;
        self.flushPendingSlow();
    }

    fn flushPendingSlow(self: *Self) void {
        const pending = self.pending_blocks.items;
        const al = self.allocator orelse return;

        // Sort pending, then merge into sorted list
        std.mem.sortUnstable(Cell, pending, {}, std.sort.asc(Cell));

        const old_len = self.all_blocks_sorted.items.len;
        const new_total = old_len + pending.len;
        self.all_blocks_sorted.ensureTotalCapacity(al, new_total) catch @panic("OOM");

        // Merge: append pending, then do a single merge pass in-place.
        // Since both halves are sorted, we merge from the end backwards.
        self.all_blocks_sorted.items.len = new_total;
        const items = self.all_blocks_sorted.items;

        var dst = new_total;
        var a = old_len; // end of sorted portion
        var b = pending.len; // end of pending portion

        while (b > 0) {
            if (a > 0 and items[a - 1] >= pending[b - 1]) {
                dst -= 1;
                a -= 1;
                items[dst] = items[a];
            } else {
                dst -= 1;
                b -= 1;
                items[dst] = pending[b];
            }
        }
        // Remaining items[0..a] are already in place at items[0..dst] since dst == a.

        // Deduplicate (pending may contain addresses already in sorted list)
        if (items.len > 1) {
            var write: usize = 1;
            for (items[1..]) |v| {
                if (v != items[write - 1]) {
                    items[write] = v;
                    write += 1;
                }
            }
            self.all_blocks_sorted.items.len = write;
        }

        self.pending_blocks.clearRetainingCapacity();
    }

    pub fn occupiedSpace(self: *const Self) Cell {
        if (self.free_list) |alloc| {
            return alloc.size - alloc.free_space;
        }
        return self.code_size; // If no allocator, assume all space is occupied
    }

    /// Returns the byte extent from code_start to the end of the last occupied
    /// block. Unlike occupiedSpace() (which sums non-free bytes), this accounts
    /// for free blocks interleaved among occupied ones. Needed by save-image
    /// because the Zig VM does not compact the code heap.
    pub fn codeHeapExtent(self: *const Self) Cell {
        const end = self.code_start + self.code_size;
        var current = self.code_start;
        var last_occupied_end: Cell = self.code_start;

        while (current < end) {
            const block: *const CodeBlock = @ptrFromInt(current);
            const block_size = block.size();
            if (block_size == 0) break;
            if (!block.isFree()) {
                last_occupied_end = current + block_size;
            }
            current += block_size;
        }
        return last_occupied_end - self.code_start;
    }

    pub fn writeBarrier(self: *Self, compiled: *CodeBlock) !void {
        try self.remembered_sets.ensureInitialized(self.code_start, self.code_size);
        try self.remembered_sets.writeBarrier(compiled);
    }

    pub fn ensureScanFlags(self: *Self, al: std.mem.Allocator) !void {
        if (self.scan_literals != null and self.scan_code_ptrs != null) return;
        const bit_count: usize = @intCast(self.code_size / layouts.data_alignment);
        self.scan_literals = try std.DynamicBitSet.initEmpty(al, bit_count);
        self.scan_code_ptrs = try std.DynamicBitSet.initEmpty(al, bit_count);
    }

    fn blockIndex(self: *const Self, block: *CodeBlock) usize {
        return @intCast((@intFromPtr(block) - self.code_start) / layouts.data_alignment);
    }

    fn blockIndexFromAddress(self: *const Self, block_addr: Cell) usize {
        return @intCast((block_addr - self.code_start) / layouts.data_alignment);
    }

    pub fn updateScanFlags(self: *Self, al: std.mem.Allocator, block: *CodeBlock) void {
        // Graceful degradation: without scan flags, all blocks are scanned (slower but correct)
        self.ensureScanFlags(al) catch return;
        const flags = code_blocks_mod.scanRelocationFlags(block);
        const idx = self.blockIndex(block);
        if (self.scan_literals) |*set| {
            if (flags.has_literals) set.set(idx) else set.unset(idx);
        }
        if (self.scan_code_ptrs) |*set| {
            if (flags.has_code_ptrs) set.set(idx) else set.unset(idx);
        }
    }

    pub fn putUninitializedBlock(self: *Self, al: std.mem.Allocator, block_addr: Cell, literals_cell: Cell) !void {
        try self.uninitialized_blocks.put(al, block_addr, literals_cell);
    }

    pub fn removeUninitializedBlock(self: *Self, block_addr: Cell) bool {
        return self.uninitialized_blocks.swapRemove(block_addr);
    }

    pub fn clearUninitializedBlocks(self: *Self) void {
        self.uninitialized_blocks.clearRetainingCapacity();
    }

    pub fn isBlockUninitialized(self: *const Self, block: *const CodeBlock) bool {
        return self.isUninitializedAddress(@intFromPtr(block));
    }

    pub fn isUninitializedAddress(self: *const Self, block_addr: Cell) bool {
        return self.uninitialized_blocks.contains(block_addr);
    }

    pub fn removeScanFlags(self: *Self, block: *CodeBlock) void {
        if (self.scan_literals == null or self.scan_code_ptrs == null) return;
        const idx = self.blockIndex(block);
        if (self.scan_literals) |*set| set.unset(idx);
        if (self.scan_code_ptrs) |*set| set.unset(idx);
    }

    pub fn removeScanFlagsByAddress(self: *Self, block_addr: Cell) void {
        if (self.scan_literals == null or self.scan_code_ptrs == null) return;
        const idx = self.blockIndexFromAddress(block_addr);
        if (self.scan_literals) |*set| set.unset(idx);
        if (self.scan_code_ptrs) |*set| set.unset(idx);
    }

    pub fn clearScanFlags(self: *Self) void {
        if (self.scan_literals) |*set| set.unmanaged.unsetAll();
        if (self.scan_code_ptrs) |*set| set.unmanaged.unsetAll();
    }

    pub fn blockHasLiterals(self: *const Self, block: *CodeBlock) bool {
        if (self.scan_literals) |set| return set.isSet(self.blockIndex(block));
        return true;
    }

    pub fn blockHasCodePointers(self: *const Self, block: *CodeBlock) bool {
        if (self.scan_code_ptrs) |set| return set.isSet(self.blockIndex(block));
        return true;
    }

    pub fn clearRememberedSets(self: *Self) void {
        self.remembered_sets.clear();
    }

    pub fn free(self: *Self, block: *CodeBlock) void {
        self.remembered_sets.removeCodeBlock(block);
        self.removeScanFlags(block);
        const block_addr = @intFromPtr(block);
        _ = self.removeUninitializedBlock(block_addr);

        self.removeFromAllBlocks(@intFromPtr(block));

        const size = block.size();
        block.markFree(size);
        if (self.free_list) |alloc| {
            alloc.free(@intFromPtr(block), size);
        }
    }

    fn removeFromAllBlocks(self: *Self, block_addr: Cell) void {
        // Check pending first (swap-remove is O(1))
        for (self.pending_blocks.items, 0..) |addr, i| {
            if (addr == block_addr) {
                _ = self.pending_blocks.swapRemove(i);
                return;
            }
        }
        // Fall back to sorted list: O(log n) search + O(n) shift
        const items = self.all_blocks_sorted.items;
        const pos = std.sort.lowerBound(Cell, items, block_addr, layouts.orderCell);
        if (pos < items.len and items[pos] == block_addr) {
            _ = self.all_blocks_sorted.orderedRemove(pos);
        }
    }

    // Batch remove addresses from all_blocks_sorted.
    pub fn batchRemoveFromAllBlocks(self: *Self, removes: []const Cell) void {
        if (removes.len == 0) return;

        if (!isNonDecreasing(removes)) {
            for (removes) |addr| {
                self.removeFromAllBlocks(addr);
            }
            return;
        }

        const items = self.all_blocks_sorted.items;
        var write_idx: usize = 0;
        var remove_idx: usize = 0;

        for (items) |addr| {
            while (remove_idx < removes.len and removes[remove_idx] < addr) : (remove_idx += 1) {}

            if (remove_idx < removes.len and removes[remove_idx] == addr) {
                // Skip duplicate remove entries for the same address.
                const removed_addr = addr;
                while (remove_idx < removes.len and removes[remove_idx] == removed_addr) : (remove_idx += 1) {}
                continue;
            }

            items[write_idx] = addr;
            write_idx += 1;
        }

        self.all_blocks_sorted.items.len = write_idx;
    }

    // Free a code block without removing from all_blocks.
    // Used for batch freeing where all_blocks is updated separately.
    pub fn freeBlockOnly(self: *Self, block: *CodeBlock) void {
        self.remembered_sets.removeCodeBlock(block);
        self.removeScanFlags(block);
        const block_addr = @intFromPtr(block);
        _ = self.removeUninitializedBlock(block_addr);

        const size = block.size();
        block.markFree(size);
        if (self.free_list) |alloc| {
            alloc.free(@intFromPtr(block), size);
        }
    }

    pub fn codeBlockForAddress(self: *Self, address: Cell) ?*CodeBlock {
        // Check sorted list first (binary search)
        const blocks = self.all_blocks_sorted.items;
        if (blocks.len > 0) {
            const ub = std.sort.upperBound(Cell, blocks, address, layouts.orderCell);
            if (ub > 0) {
                const block: *CodeBlock = @ptrFromInt(blocks[ub - 1]);
                const block_end = blocks[ub - 1] + block.size();
                if (address < block_end) return block;
            }
        }

        // Check pending blocks (linear scan — typically small)
        for (self.pending_blocks.items) |block_addr| {
            if (address >= block_addr) {
                const block: *CodeBlock = @ptrFromInt(block_addr);
                const block_end = block_addr + block.size();
                if (address < block_end) return block;
            }
        }

        return null;
    }

    // walking: choose the previous block start without requiring the address to
    // lie within the block extent.
    pub fn codeBlockForReturnAddress(self: *Self, address: Cell) ?*CodeBlock {
        const blocks = self.all_blocks_sorted.items;
        if (blocks.len > 0) {
            const ub = std.sort.upperBound(Cell, blocks, address, layouts.orderCell);
            if (ub > 0) {
                return @ptrFromInt(blocks[ub - 1]);
            }
        }

        for (self.pending_blocks.items) |block_addr| {
            if (address >= block_addr) {
                return @ptrFromInt(block_addr);
            }
        }

        return null;
    }

    // Get the predecessor frame (caller's frame) given the current frame top
    pub fn framePredecessor(self: *Self, frame_top: Cell) Cell {
        if (builtin.cpu.arch == .aarch64) {
            // ARM64: frame_top[0] contains the saved frame pointer (x29)
            // which points directly to the previous frame
            return @as(*const Cell, @ptrFromInt(frame_top)).*;
        } else {
            // x86-64: FRAME_RETURN_ADDRESS = 0, so return address is at frame_top
            // We compute the next frame by adding the frame size
            const FRAME_RETURN_ADDRESS: Cell = 0;
            const addr = @as(*const Cell, @ptrFromInt(frame_top + FRAME_RETURN_ADDRESS)).*;
            const block = self.codeBlockForAddress(addr) orelse {
                // Can't find code block - return minimum frame size
                return frame_top + CodeBlock.LEAF_FRAME_SIZE;
            };
            const frame_size = block.stackFrameSizeForAddress(addr);
            return frame_top + frame_size;
        }
    }

    // Verify that all live code blocks in the heap are present in all_blocks_sorted.
    // Catches missed inserts/removes or stale entries.
    pub fn verifyAllBlocksSet(self: *const Self) void {
        if (comptime builtin.mode != .Debug) return;
        if (self.code_start == 0 or self.code_size == 0) return;

        const code_end = self.code_start + self.code_size;
        var idx: usize = 0;
        var current = self.code_start;

        while (current < code_end) {
            const block: *const CodeBlock = @ptrFromInt(current);
            const block_size = codeBlockSize(current);
            if (block_size == 0) break;

            if (!block.isFree()) {
                std.debug.assert(idx < self.all_blocks_sorted.items.len);
                std.debug.assert(self.all_blocks_sorted.items[idx] == current);
                idx += 1;
            }
            current += block_size;
        }

        std.debug.assert(idx == self.all_blocks_sorted.items.len);
        if (self.all_blocks_sorted.items.len > 1) {
            for (1..self.all_blocks_sorted.items.len) |i| {
                std.debug.assert(self.all_blocks_sorted.items[i - 1] < self.all_blocks_sorted.items[i]);
            }
        }
    }

    // Initialize all_blocks_sorted by scanning the code heap.
    // This must be called after image loading/fixup
    pub fn initializeAllBlocksSet(self: *Self) !void {
        if (self.code_start == 0 or self.code_size == 0) return;

        const alloc = self.allocator orelse return;

        // First pass: count non-free blocks
        var count: usize = 0;
        var current = self.code_start;
        const code_end = self.code_start + self.code_size;

        while (current < code_end) {
            const block_size = codeBlockSize(current);
            if (block_size == 0) break;
            const block: *const CodeBlock = @ptrFromInt(current);
            if (!block.isFree()) {
                count += 1;
            }
            current += block_size;
        }

        // Rebuild sorted array (pending is stale after full scan).
        self.pending_blocks.clearRetainingCapacity();
        self.all_blocks_sorted.clearRetainingCapacity();
        try self.all_blocks_sorted.ensureTotalCapacity(alloc, count);
        // Second pass: populate. Iteration order is ascending address, so
        // all_blocks_sorted is naturally sorted without a separate sort call.
        current = self.code_start;
        while (current < code_end) {
            const block_size = codeBlockSize(current);
            if (block_size == 0) break;
            const block: *const CodeBlock = @ptrFromInt(current);
            if (!block.isFree()) {
                self.all_blocks_sorted.appendAssumeCapacity(current);
            }
            current += block_size;
        }

        self.verifyAllBlocksSet();
    }

    pub fn rebuildScanFlags(self: *Self, al: std.mem.Allocator) void {
        self.flushPending();
        self.ensureScanFlags(al) catch return;
        self.clearScanFlags();
        for (self.all_blocks_sorted.items) |block_addr| {
            const block: *CodeBlock = @ptrFromInt(block_addr);
            if (!block.isFree()) {
                self.updateScanFlags(al, block);
            }
        }
    }

    pub fn ensureMarks(self: *Self, al: std.mem.Allocator) !void {
        if (self.marks != null) return;
        if (self.code_start == 0 or self.code_size == 0) return;
        const marks = try al.create(mark_bits.MarkBits);
        errdefer al.destroy(marks);
        marks.* = try mark_bits.MarkBits.init(al, self.code_start, self.code_size);
        self.marks = marks;
        self.allocator = al;
    }

    pub fn clearMarks(self: *Self) void {
        if (self.marks) |marks| {
            marks.clearMarks();
        }
    }

    pub fn deinit(self: *Self) void {
        const alloc = self.allocator orelse return;
        self.pending_blocks.deinit(alloc);
        self.all_blocks_sorted.deinit(alloc);
        self.uninitialized_blocks.deinit(alloc);
        self.uninitialized_blocks_scratch.deinit(alloc);
        // Deinit remembered sets and scan flags
        self.remembered_sets.deinit();
        if (self.scan_literals) |*set| {
            set.deinit();
            self.scan_literals = null;
        }
        if (self.scan_code_ptrs) |*set| {
            set.deinit();
            self.scan_code_ptrs = null;
        }
        if (self.marks) |marks| {
            marks.deinit();
            if (self.allocator) |gc_alloc| {
                gc_alloc.destroy(marks);
            }
            self.marks = null;
        }
    }

    pub fn blockSizeAt(self: *const Self, addr: Cell) Cell {
        _ = self;
        return codeBlockSize(addr);
    }
};

// Compute block size at address, handling both code-block and free-list layouts.
// FreeListAllocator encodes size in the header (header = size | 1).
fn codeBlockSize(addr: Cell) Cell {
    const block: *const CodeBlock = @ptrFromInt(addr);
    if (!block.isFree()) {
        return block.size();
    }
    var size = block.size();
    if (size == 0) {
        // For free blocks, read size from header (FreeBlock encodes size in header)
        const free_block: *const free_list.FreeBlock = @ptrFromInt(addr);
        size = free_block.size();
    }
    return size;
}

fn isNonDecreasing(values: []const Cell) bool {
    if (values.len < 2) return true;
    for (1..values.len) |i| {
        if (values[i] < values[i - 1]) return false;
    }
    return true;
}
