//! 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::il::published;

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;

/// Largest spill frame representable by signed byte offsets.

constant MAX_FRAME_SIZE: u32 = 0x7fffffff;

/// 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: 'scratch + Copy {

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

    slots: &'scratch [i32],

    /// Total spill frame size needed in bytes.

    frameSize: i32,

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

    calleeClass: &'scratch [u32],

    /// Maximum SSA register number.

    maxReg: u32,

}

/// Candidate buffer for spill decisions.

record Candidates: Copy {

    entries: [CostEntry; MAX_CANDIDATES],

    n: u32,

}

/// Entry for cost sorting.

record CostEntry: Copy {

    reg: u32,

    cost: u32,

}

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

export fn analyze 'input 'scratch (

    func: &published::Function 'input,

    live: &liveness::LiveInfo 'scratch,

    numRegs: u32,

    numCalleeSaved: u32,

    slotSize: u32,

    storage: &Session 'scratch

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

    return try analyzeBlocks(func.blocks, live, numRegs, numCalleeSaved, slotSize, storage);

}

/// Allocate spill storage and analyze pressure over a borrowed block table.

fn analyzeBlocks 'input 'scratch (

    blocks: &[published::Block 'input],

    live: &liveness::LiveInfo 'scratch,

    numRegs: u32,

    numCalleeSaved: u32,

    slotSize: u32,

    storage: &Session 'scratch,

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

    let maxReg = live.maxReg;

    if maxReg == 0 {

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

        return SpillInfo 'scratch {

            slots: try storage.fill(-1 as i32, 0),

            frameSize: 0,

            calleeClass: &calleeClass[..],

            maxReg: 0,

        };

    }

    // Allocate spill slots array.

    let slots = try storage.fill(-1 as i32, maxReg);

    // Allocate cost array.

    let costs = try storage.fill(SpillCost { defs: 0, uses: 0 }, maxReg);

    // Phase 1: Calculate spill costs.

    for b in 0..blocks.len {

        let block = &blocks[b];

        fillCosts(block.params, block.instructions, block.loopDepth, costs);

    }

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

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

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

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

    for b in 0..blocks.len {

        let block = &blocks[b];

        // Start with live-out set.

        bitset::copy(scratch, liveness::liveOutRow(live, b));

        try limitBlockPressure(block.instructions, scratch, spilled, costs, calleeClass, numRegs, numCalleeSaved);

    }

    let frameSize = try finishSpills(slots, spilled, calleeClass, numCalleeSaved, slotSize);

    return SpillInfo 'scratch { slots: &slots[..], frameSize, calleeClass: &calleeClass[..], maxReg };

}

/// Limit register pressure from the block's live-out set to its entry.

fn limitBlockPressure 'input (

    instructions: &[published::Instruction 'input],

    live: &mut [u32],

    spilled: &mut [u32],

    costs: &[SpillCost],

    calleeClass: &mut [u32],

    numRegs: u32,

    numCalleeSaved: u32,

) throws (alloc::AllocError) {

    // Walk instructions backwards.

    let mut i = instructions.len;

    while i > 0 {

        set i -= 1;

        let instr = &instructions[i];

        // Limit register pressure before processing this instruction.

        try limitPressure(live, spilled, costs, numRegs);

        // Enforce cross-call pressure at call sites.

        if published::isCall(instr) {

            try limitCrossCallPressure(

                live, spilled, costs,

                calleeClass, numCalleeSaved, published::destination(instr)

            );

        }

        // Remove definition from live set.

        if let dst = published::destination(instr) {

            bitset::clear(live, dst.n);

        }

        // Add uses to live set.

        // Fixed operands precede the instruction's variable argument groups.

        for reg in published::registers(instr) {

            bitset::put(live, reg.n);

        }

    }

    // Also limit pressure at block entry.

    try limitPressure(live, spilled, costs, numRegs);

}

/// Enforce the global register-class limit and assign bounded stack offsets.

fn finishSpills(

    slots: &mut [i32],

    spilled: &mut [u32],

    calleeClass: &mut [u32],

    numCalleeSaved: u32,

    slotSize: u32,

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

    // 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, calleeClass) {

                bitset::clear(calleeClass, reg);

                bitset::put(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, spilled) {

        if slotSize > MAX_FRAME_SIZE - frameSize as u32 {

            throw alloc::AllocError::OutOfMemory;

        }

        set slots[n] = frameSize;

        set frameSize += slotSize as i32;

    }

    return frameSize;

}

/// Add one block's register spill costs, weighted by loop depth.

fn fillCosts 'input (params: &[il::Param], instructions: &[published::Instruction 'input], loopDepth: u32, costs: &mut [SpillCost]) {

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

    let depth = MAX_LOOP_WEIGHT if loopDepth > MAX_LOOP_WEIGHT else loopDepth;

    let weight: u32 = 1 << depth;

    // Count block parameter definitions.

    for p in 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..instructions.len {

        let instr = &instructions[i];

        // Count definition.

        if let dst = published::destination(instr) {

            if dst.n < costs.len {

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

            }

        }

        // Count uses.

        // Fixed operands precede the instruction's variable argument groups.

        for reg in published::registers(instr) {

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

        }

    }

}

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

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

fn spillCheapest(

    c: &mut Candidates,

    excess: u32,

    source: &mut [u32],

    spilled: &mut [u32]

) {

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

fn collectCandidates(bs: &[u32], costs: &[SpillCost]) -> Candidates throws (alloc::AllocError) {

    let mut c = Candidates { entries: [CostEntry { reg: 0, cost: 0 }; MAX_CANDIDATES], n: 0 };

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

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

        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.

fn limitPressure(

    live: &mut [u32],

    spilled: &mut [u32],

    costs: &[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.

fn limitCrossCallPressure(

    live: &mut [u32],

    spilled: &mut [u32],

    costs: &[SpillCost],

    calleeClass: &mut [u32],

    numCalleeSaved: u32,

    callDst: ?il::Reg

) throws (alloc::AllocError) {

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

    let mut candidates: [CostEntry; MAX_CANDIDATES] = [CostEntry { reg: 0, cost: 0 }; MAX_CANDIDATES];

    let mut numCandidates: u32 = 0;

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

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

        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, live) {

        if not isCallDst(callDst, n) {

            bitset::put(calleeClass, n);

        }

    }

}

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

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

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

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

}

/// Check if a register is spilled.

export fn isSpilled 'scratch (info: &SpillInfo 'scratch, 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 'scratch (info: &SpillInfo 'scratch, reg: il::Reg) -> ?i32 {

    if isSpilled(info, reg) {

        return info.slots[reg.n];

    }

    return nil;

}