//! Register assignment.

//!

//! This pass assigns physical registers to SSA values. It runs after spilling

//! has ensured register pressure never exceeds available registers.

//!

//! The IL is not modified. Instead, a mapping from SSA registers to physical

//! registers is produced for use by instruction selection.

use std::lang::il;

use std::lang::alloc;

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

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

use std::lang::gen::regalloc::spill;

/// Maximum number of active register mappings.

const MAX_ACTIVE: u32 = 64;

/// Register mapping at a program point.

/// Maps SSA registers to physical registers.

pub record RegMap {

    /// SSA (virtual) registers that have mappings.

    virtRegs: *mut [u32],

    /// Physical register for each virtual register.

    physRegs: *mut [u8],

    /// Number of active mappings.

    n: u32,

}

/// Register assignment result, per function.

pub record AssignInfo {

    /// SSA register -> physical register mapping.

    assignments: *mut [?u8],

    /// Bitmask of used callee-saved registers.

    usedCalleeSaved: u32,

}

/// Per-instruction context for freeing and allocating register uses.

record InstrCtx {

    current: *mut RegMap,

    usedRegs: *mut bitset::Bitset,

    func: *il::Fn,

    live: *liveness::LiveInfo,

    blockIdx: u32,

    instrIdx: u32,

    allocatable: *[u8],

    calleeSaved: *[u8],

    assignments: *mut [?u8],

    spillInfo: *spill::SpillInfo,

}

/// Compute register assignment.

pub fn assign(

    func: *il::Fn,

    live: *liveness::LiveInfo,

    spillInfo: *spill::SpillInfo,

    config: *super::TargetConfig,

    arena: *mut alloc::Arena

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

    let maxReg = live.maxReg;

    let blockCount = func.blocks.len;

    let allocatable = config.allocatable;

    if maxReg == 0 or blockCount == 0 {

        return AssignInfo {

            assignments: &mut [],

            usedCalleeSaved: 0,

        };

    }

    // Allocate output structures.

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

    for i in 0..maxReg {

        assignments[i] = nil;

    }

    // Pre-assign function parameters to argument registers.

    // Cross-call params are NOT pre-assigned here; they will be allocated

    // to callee-saved registers by the normal path, and isel emits moves

    // from the arg register to the assigned register at function entry.

    for param, i in func.params {

        if i < config.argRegs.len {

            if not bitset::contains(&spillInfo.calleeClass, param.value.n) {

                assignments[param.value.n] = config.argRegs[i];

            }

        }

    }

    let blkEnd = try alloc::allocSlice(arena, @sizeOf(RegMap), @alignOf(RegMap), blockCount) as *mut [RegMap];

    for i in 0..blockCount {

        blkEnd[i] = try createRegMap(arena);

    }

    // Allocate used registers bitset (32 physical registers per 32-bit word).

    let usedRegsBits = try alloc::allocSlice(arena, @sizeOf(u32), @alignOf(u32), 1) as *mut [u32];

    let mut usedRegs = bitset::init(usedRegsBits);

    // Current register mapping.

    let mut current = try createRegMap(arena);

    // Phase 2: Linear scan allocation.

    for b in 0..blockCount {

        let block = &func.blocks[b];

        // Reset for new block.

        current.n = 0;

        bitset::clearAll(&mut usedRegs);

        // Initialize from predecessor if single predecessor.

        // Only copy values that are live-in to this block.

        if block.preds.len == 1 {

            let pred = &blkEnd[block.preds[0]];

            for k in 0..pred.n {

                if bitset::contains(&live.liveIn[b], pred.virtRegs[k]) {

                    current.virtRegs[current.n] = pred.virtRegs[k];

                    current.physRegs[current.n] = pred.physRegs[k];

                    current.n += 1;

                    bitset::set(&mut usedRegs, pred.physRegs[k] as u32);

                }

            }

        }

        // Mark all live-in values' registers as used.

        // This ensures we don't reuse registers for values that flow in

        // from predecessors, even at merge points with multiple predecessors.

        // Values that are live-in but have no assignment yet (e.g. callee-saved

        // function parameters not used before a phi block) are allocated now to

        // prevent conflicts with block parameters.

        let mut liveInIter = bitset::iter(&live.liveIn[b]);

        while let ssaReg = bitset::iterNext(&mut liveInIter) {

            let reg = il::Reg { n: ssaReg };

            if not spill::isSpilled(spillInfo, reg) {

                if let phys = assignments[ssaReg] {

                    bitset::set(&mut usedRegs, phys as u32);

                    // Also track in current mapping for block-end state.

                    if rmapFind(&current, ssaReg) == nil {

                        rmapSet(&mut current, ssaReg, phys);

                    }

                } else {

                    assignments[ssaReg] = rallocReg(&mut current, &mut usedRegs, ssaReg, allocatable, config.calleeSaved, spillInfo);

                }

            }

        }

        // Allocate block parameters.

        for p in block.params {

            if p.value.n < maxReg and not spill::isSpilled(spillInfo, p.value) {

                assignments[p.value.n] = rallocReg(&mut current, &mut usedRegs, p.value.n, allocatable, config.calleeSaved, spillInfo);

            }

        }

        // Process each instruction.

        for instr, i in block.instrs {

            let mut ctx = InstrCtx {

                current: &mut current,

                usedRegs: &mut usedRegs,

                func,

                live,

                blockIdx: b,

                instrIdx: i,

                allocatable,

                calleeSaved: config.calleeSaved,

                assignments,

                spillInfo,

            };

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

            // Allocate destination.

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

                if dst.n < maxReg and not spill::isSpilled(spillInfo, dst) {

                    assignments[dst.n] = rallocReg(&mut current, &mut usedRegs, dst.n, allocatable, config.calleeSaved, spillInfo);

                }

            }

        }

        // Save block-end state.

        rmapCopy(&mut blkEnd[b], &current);

    }

    // Compute bitmask of used callee-saved registers.

    let mut usedCalleeSaved: u32 = 0;

    for i in 0..maxReg {

        if let phys = assignments[i] {

            for saved, j in config.calleeSaved {

                if phys == saved {

                    usedCalleeSaved |= (1 << j);

                }

            }

        }

    }

    return AssignInfo {

        assignments,

        usedCalleeSaved,

    };

}

/// Create an empty register map.

fn createRegMap(arena: *mut alloc::Arena) -> RegMap throws (alloc::AllocError) {

    let virtRegs = try alloc::allocSlice(arena, @sizeOf(u32), @alignOf(u32), MAX_ACTIVE) as *mut [u32];

    let physRegs = try alloc::allocSlice(arena, @sizeOf(u8), @alignOf(u8), MAX_ACTIVE) as *mut [u8];

    return RegMap { virtRegs, physRegs, n: 0 };

}

/// Copy a register map.

fn rmapCopy(dst: *mut RegMap, src: *RegMap) {

    dst.n = src.n;

    for i in 0..src.n {

        dst.virtRegs[i] = src.virtRegs[i];

        dst.physRegs[i] = src.physRegs[i];

    }

}

/// Find physical register for a virtual register in RegMap.

fn rmapFind(rmap: *RegMap, virtReg: u32) -> ?u8 {

    for i in 0..rmap.n {

        if rmap.virtRegs[i] == virtReg {

            return rmap.physRegs[i];

        }

    }

    return nil;

}

/// Add a mapping to the register map.

fn rmapSet(rmap: *mut RegMap, virtReg: u32, physReg: u8) {

    if rmap.n >= MAX_ACTIVE {

        panic "rmapSet: register map overflow";

    }

    rmap.virtRegs[rmap.n] = virtReg;

    rmap.physRegs[rmap.n] = physReg;

    rmap.n += 1;

}

/// Remove a mapping from register map.

fn rmapRemove(rmap: *mut RegMap, virtReg: u32) {

    for i in 0..rmap.n {

        if rmap.virtRegs[i] == virtReg {

            // Swap with last and decrement.

            rmap.n -= 1;

            if i < rmap.n {

                rmap.virtRegs[i] = rmap.virtRegs[rmap.n];

                rmap.physRegs[i] = rmap.physRegs[rmap.n];

            }

            return;

        }

    }

    panic "rmapRemove: register not found";

}

/// Find first free register in pool, allocate it, return it.

fn findFreeInPool(usedRegs: *mut bitset::Bitset, current: *mut RegMap, ssaReg: u32, pool: *[u8]) -> ?u8 {

    for i in 0..pool.len {

        let r = pool[i];

        if not bitset::contains(usedRegs, r as u32) {

            bitset::set(usedRegs, r as u32);

            rmapSet(current, ssaReg, r);

            return r;

        }

    }

    return nil;

}

/// Allocate a physical register for an SSA register.

/// Cross-call values are steered to callee-saved registers.

fn rallocReg(

    current: *mut RegMap,

    usedRegs: *mut bitset::Bitset,

    ssaReg: u32,

    allocatable: *[u8],

    calleeSaved: *[u8],

    spillInfo: *spill::SpillInfo

) -> u8 {

    // Check if already assigned.

    if let phys = rmapFind(current, ssaReg) {

        return phys;

    }

    let crossCall = bitset::contains(&spillInfo.calleeClass, ssaReg);

    // Allocate from appropriate pool. Cross-call values must use callee-saved

    // registers since they are live across function calls.

    if crossCall {

        if let r = findFreeInPool(usedRegs, current, ssaReg, calleeSaved) {

            return r;

        }

        panic "rallocReg: no callee-saved register for cross-call value";

    }

    if let r = findFreeInPool(usedRegs, current, ssaReg, allocatable) {

        return r;

    }

    panic "rallocReg: no free register, spilling fault";

}

/// Callback for [`il::forEachReg`]: free last uses and allocate missing uses.

fn processInstrRegCb(reg: il::Reg, ctxPtr: *mut opaque) {

    let ctx = ctxPtr as *mut InstrCtx;

    if liveness::isLastUse(ctx.live, ctx.func, ctx.blockIdx, ctx.instrIdx, reg) {

        if let phys = rmapFind(ctx.current, reg.n) {

            bitset::clear(ctx.usedRegs, phys as u32);

            rmapRemove(ctx.current, reg.n);

        }

    }

    if reg.n >= ctx.assignments.len {

        panic "processInstrRegCb: register out of bounds";

    }

    if spill::isSpilled(ctx.spillInfo, reg) {

        return; // Spilled values don't get physical registers.

    }

    if ctx.assignments[reg.n] == nil {

        ctx.assignments[reg.n] = rallocReg(

            ctx.current, ctx.usedRegs, reg.n, ctx.allocatable,

            ctx.calleeSaved, ctx.spillInfo

        );

    }

}