//! Spilling pass - determines which SSA values need stack slots.

//!

//! This pass analyzes register pressure and determines which values must be

//! spilled to memory. It does not modify the IL.

//!

//! The instruction selector uses this information to emit load/store

//! instructions for spilled values.

//!

//! # Algorithm

//!

//! 1. Cost calculation: for each value, compute spill cost:

//!

//!        cost = (defs + uses) * 2^min(loopDepth, 10)

//!

//!    Values used in inner loops are expensive to spill.

//!

//! 2. Pressure analysis: walk instructions backwards, tracking live set.

//!

//!    When `|live| > numRegs`:

//!

//!    * Sort live values by cost (ascending).

//!    * Mark lowest-cost values as spilled until pressure is acceptable.

//!

//! 3. Slot assignment: assign stack offsets to spilled values.

//!

//! # Output

//!

//! Spill info containing:

//!

//! * Spill slot assignments: SSA reg -> stack offset, `-1` if not spilled.

//! * Total frame size needed for spills.

use std::lang::il;

use std::lang::alloc;

use std::lang::gen::bitset;

use std::lang::gen::regalloc::liveness;

/// Maximum number of candidates for spill sorting.

constant MAX_CANDIDATES: u32 = 256;

/// Maximum loop depth for cost weighting (2^10 = 1024).

constant MAX_LOOP_WEIGHT: u32 = 10;

/// Spill cost for a single SSA register.

record SpillCost: Copy {

    /// Number of definitions (weighted by loop depth).

    defs: u32,

    /// Number of uses (weighted by loop depth).

    uses: u32,

}

/// Spill decision for a function.

export record SpillInfo: Copy {

    /// SSA register mapped to stack slot offset. `-1` means not spilled.

    slots: *[i32],

    /// Total spill frame size needed in bytes.

    frameSize: i32,

    /// Values that must be allocated in callee-saved registers.

    calleeClass: bitset::Bitset,

    /// Maximum SSA register number.

    maxReg: u32,

}

/// Candidate buffer for spill decisions.

record Candidates: Copy {

    entries: [CostEntry; 256],

    n: u32,

}

/// Entry for cost sorting.

record CostEntry: Copy {

    reg: u32,

    cost: u32,

}

/// Context for counting register uses.

record CountCtx: Copy {

    costs: *unsafe mut [SpillCost],

    weight: u32,

}

/// Analyze a function and determine which values need spill slots.

export unsafe fn analyze(

    func: *unsafe il::Fn,

    live: &liveness::LiveInfo,

    numRegs: u32,

    numCalleeSaved: u32,

    slotSize: u32,

    arena: &mut alloc::Arena

) -> SpillInfo throws (alloc::AllocError) {

    let maxReg = live.maxReg;

    if maxReg == 0 {

        let calleeClass = try bitset::allocate(arena, 0);

        return SpillInfo {

            slots: &mut [],

            frameSize: 0,

            calleeClass,

            maxReg: 0,

        };

    }

    // Allocate spill slots array.

    let slots = try alloc::allocSlice(arena, @sizeOf(i32), @alignOf(i32), maxReg) as *mut [i32];

    for i in 0..maxReg {

        set slots[i] = -1;

    }

    // Allocate cost array.

    let costs = try alloc::allocRawSlice(arena, @sizeOf(SpillCost), @alignOf(SpillCost), maxReg) as *unsafe mut [SpillCost];

    for i in 0..maxReg {

        set costs[i] = SpillCost { defs: 0, uses: 0 };

    }

    // Phase 1: Calculate spill costs.

    fillCosts(func, costs);

    // Phase 2: Find values that exceed register pressure.

    let mut spilled = try bitset::allocate(arena, maxReg);

    let mut calleeClass = try bitset::allocate(arena, maxReg);

    let mut scratch = try bitset::allocate(arena, maxReg);

    for b in 0..func.blocks.len {

        let block = &func.blocks[b];

        // Start with live-out set.

        bitset::copy(&mut scratch, &live.liveOut[b]);

        // Walk instructions backwards.

        let mut i = block.instrs.len;

        while i > 0 {

            set i -= 1;

            let instr = block.instrs[i];

            // Limit register pressure before processing this instruction.

            try limitPressure(&mut scratch, &mut spilled, costs, numRegs);

            // Enforce cross-call pressure at call sites.

            if il::isCall(instr) {

                try limitCrossCallPressure(

                    &mut scratch, &mut spilled, costs,

                    &mut calleeClass, numCalleeSaved, il::instrDst(instr)

                );

            }

            // Remove definition from live set.

            if let dst = il::instrDst(instr) {

                bitset::clear(&mut scratch, dst.n);

            }

            // Add uses to live set.

            il::forEachReg(instr, addRegToSetCallback, &mut scratch as &mut opaque);

        }

        // Also limit pressure at block entry.

        try limitPressure(&mut scratch, &mut spilled, costs, numRegs);

    }

    // Phase 3: Enforce global callee-class limit.

    // The per-call-site limit may leave the callee-class set larger than

    // `numCalleeSaved` when different call sites keep different subsets.

    // Spill excess values to guarantee the assignment phase always finds

    // a callee-saved register for cross-call values.

    let calleeCount = bitset::count(&calleeClass);

    if calleeCount > numCalleeSaved {

        let mut it = bitset::iter(&calleeClass);

        for _ in 0..(calleeCount - numCalleeSaved) {

            if let reg = bitset::iterNext(&mut it) {

                bitset::clear(&mut calleeClass, reg);

                bitset::put(&mut spilled, reg);

            } else {

                panic "spill: count > 0 but no set bits found";

            }

        }

    }

    // Phase 4: Assign stack slots to spilled values.

    let mut frameSize: i32 = 0;

    let mut it = bitset::iter(&spilled);

    while let n = bitset::iterNext(&mut it) {

        set slots[n] = frameSize;

        set frameSize += slotSize as i32;

    }

    return SpillInfo { slots, frameSize, calleeClass, maxReg };

}

/// Calculate spill costs for all registers, weighted by loop depth.

unsafe fn fillCosts(func: *unsafe il::Fn, costs: *unsafe mut [SpillCost]) {

    for b in 0..func.blocks.len {

        let block = &func.blocks[b];

        // Exponential weight for loop depth, capped to avoid overflow.

        let depth = MAX_LOOP_WEIGHT if block.loopDepth > MAX_LOOP_WEIGHT else block.loopDepth;

        let weight: u32 = 1 << depth;

        // Count block parameter definitions.

        for p in block.params {

            if p.value.n < costs.len {

                set costs[p.value.n].defs = costs[p.value.n].defs + weight;

            }

        }

        // Count instruction defs and uses.

        for i in 0..block.instrs.len {

            let instr = block.instrs[i];

            // Count definition.

            if let dst = il::instrDst(instr) {

                if dst.n < costs.len {

                    set costs[dst.n].defs = costs[dst.n].defs + weight;

                }

            }

            // Count uses.

            let mut ctx = CountCtx { costs, weight };

            il::forEachReg(instr, countRegUseCallback, &mut ctx as &mut opaque);

        }

    }

}

/// Sort candidates by cost (ascending) using insertion sort, then spill

/// the cheapest `excess` values: set in `spilled`, clear in `source`.

unsafe fn spillCheapest(

    c: &mut Candidates,

    excess: u32,

    source: &mut bitset::Bitset,

    spilled: &mut bitset::Bitset

) {

    // Insertion sort ascending by cost.

    for i in 1..c.n {

        let key = c.entries[i];

        let mut j: u32 = i;

        while j > 0 and c.entries[j - 1].cost > key.cost {

            set c.entries[j] = c.entries[j - 1];

            set j -= 1;

        }

        set c.entries[j] = key;

    }

    let toSpill = c.n if excess > c.n else excess;

    for i in 0..toSpill {

        bitset::put(spilled, c.entries[i].reg);

        bitset::clear(source, c.entries[i].reg);

    }

}

/// Collect all values from a bitset into a candidates buffer with their costs.

unsafe fn collectCandidates(bs: &bitset::Bitset, costs: *unsafe [SpillCost]) -> Candidates throws (alloc::AllocError) {

    let mut c = Candidates { entries: undefined, n: 0 };

    let mut it = bitset::iter(bs);

    while let reg = bitset::iterNext(&mut it) {

        if c.n == MAX_CANDIDATES {

            throw alloc::AllocError::OutOfMemory;

        }

        if reg < costs.len {

            set c.entries[c.n] = CostEntry { reg, cost: costs[reg].defs + costs[reg].uses };

            set c.n += 1;

        }

    }

    return c;

}

/// Limit register pressure by marking low-cost values as spilled.

unsafe fn limitPressure(

    live: &mut bitset::Bitset,

    spilled: &mut bitset::Bitset,

    costs: *unsafe [SpillCost],

    numRegs: u32

) throws (alloc::AllocError) {

    let liveCount = bitset::count(live);

    if liveCount <= numRegs {

        return;

    }

    let mut c = try collectCandidates(live, costs);

    spillCheapest(&mut c, liveCount - numRegs, live, spilled);

}

/// Check whether an SSA register is the destination of a call instruction.

fn isCallDst(callDst: ?il::Reg, n: u32) -> bool {

    if let d = callDst {

        return d.n == n;

    }

    return false;

}

/// Limit cross-call pressure by spilling values that exceed callee-saved capacity.

///

/// At a call site, every live value must survive the call in a callee-saved

/// register. If the count exceeds `numCalleeSaved`, spill the cheapest

/// crossing values.

unsafe fn limitCrossCallPressure(

    live: &mut bitset::Bitset,

    spilled: &mut bitset::Bitset,

    costs: *unsafe [SpillCost],

    calleeClass: &mut bitset::Bitset,

    numCalleeSaved: u32,

    callDst: ?il::Reg

) throws (alloc::AllocError) {

    // Collect crossing candidates: live values excluding the call destination.

    let mut candidates: [CostEntry; 256] = undefined;

    let mut numCandidates: u32 = 0;

    let mut it = bitset::iter(live);

    while let n = bitset::iterNext(&mut it) {

        if not isCallDst(callDst, n) and n < costs.len {

            if numCandidates == MAX_CANDIDATES {

                throw alloc::AllocError::OutOfMemory;

            }

            set candidates[numCandidates] = CostEntry {

                reg: n,

                cost: costs[n].defs + costs[n].uses,

            };

            set numCandidates += 1;

        }

    }

    // Spill cheapest candidates if crossing count exceeds callee-saved capacity.

    if numCandidates > numCalleeSaved {

        let mut c = Candidates { entries: candidates, n: numCandidates };

        spillCheapest(&mut c, numCandidates - numCalleeSaved, live, spilled);

    }

    // Mark crossing values that remain after spilling as callee-saved class.

    let mut it2 = bitset::iter(live);

    while let n = bitset::iterNext(&mut it2) {

        if not isCallDst(callDst, n) {

            bitset::put(calleeClass, n);

        }

    }

}

/// Callback for [`il::forEachReg`]: increments use count for register.

unsafe fn countRegUseCallback(reg: il::Reg, ctxPtr: &mut opaque) {

    let weight = (ctxPtr as &mut CountCtx).weight;

    countRegUse(reg, &mut (ctxPtr as &mut CountCtx).costs[..], weight);

}

/// Add the block weight to the register use count.

fn countRegUse(reg: il::Reg, costs: &mut [SpillCost], weight: u32) {

    assert reg.n < costs.len, "countRegUseCallback: register out of bounds";

    set costs[reg.n].uses += weight;

}

/// Callback for [`il::forEachReg`]: adds register to live set.

unsafe fn addRegToSetCallback(reg: il::Reg, ctx: &mut opaque) {

    bitset::put(ctx as &mut bitset::Bitset, reg.n);

}

/// Check if a register is spilled.

export fn isSpilled(info: &SpillInfo, reg: il::Reg) -> bool {

    if reg.n >= info.maxReg {

        return false;

    }

    return info.slots[reg.n] >= 0;

}

/// Get spill slot offset for a register, or `nil` if not spilled.

export fn spillSlot(info: &SpillInfo, reg: il::Reg) -> ?i32 {

    if isSpilled(info, reg) {

        return info.slots[reg.n];

    }

    return nil;

}