//! 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;

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

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

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

/// Maximum number of active register mappings.

constant MAX_ACTIVE: u32 = 64;

/// Register mapping at a program point.

/// Maps SSA registers to physical registers.

export record RegMap: Copy {

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

    virtRegs: *mut [u32],

    /// Physical register for each virtual register.

    physRegs: *mut [gen::Reg],

    /// Number of active mappings.

    n: u32,

}

/// Register assignment result, per function.

export record AssignInfo: Copy {

    /// SSA register -> physical register mapping.

    assignments: *mut [?gen::Reg],

    /// Bitmask of used callee-saved registers.

    usedCalleeSaved: u32,

}

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

record InstrCtx: Copy {

    current: *mut RegMap,

    usedRegs: *mut bitset::Bitset,

    /// Last operand-use index for each register used in the current block.

    lastUse: *[u32],

    live: *liveness::LiveInfo,

    blockIdx: u32,

    instrIdx: u32,

    allocatable: *[gen::Reg],

    calleeSaved: *[gen::Reg],

    assignments: *mut [?gen::Reg],

    spillInfo: *spill::SpillInfo,

}

/// Context for recording the last operand-use index in a block.

record LastUseCtx: Copy {

    /// Per-register indices, shared by all blocks in the function.

    lastUse: *mut [u32],

    /// Index of the instruction whose operands are being recorded.

    index: u32,

}

/// Compute register assignment.

export 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(?gen::Reg), @alignOf(?gen::Reg), maxReg) as *mut [?gen::Reg];

    for i in 0..maxReg {

        set assignments[i] = nil;

    }

    // Reuse one last-use table for all blocks in the function.

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

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

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

            }

        }

    }

    // 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];

        // Record every operand before allocation. Only current-block operands

        // query this table, so every read is initialized by this scan.

        // Entries for other registers need not be cleared between blocks.

        for instr, i in block.instrs {

            let mut ctx = LastUseCtx { lastUse, index: i };

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

        }

        // Reset for new block.

        set current.n = 0;

        bitset::clearAll(&mut usedRegs);

        // 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::put(&mut usedRegs, *phys as u32);

                    rmapSet(&mut current, ssaReg, phys);

                } else {

                    set 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) {

                set 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,

                lastUse,

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

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

                }

            }

        }

    }

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

                    set 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(gen::Reg), @alignOf(gen::Reg), MAX_ACTIVE) as *mut [gen::Reg];

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

}

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

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

    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: gen::Reg) {

    assert rmap.n < MAX_ACTIVE, "rmapSet: register map overflow";

    set rmap.virtRegs[rmap.n] = virtReg;

    set rmap.physRegs[rmap.n] = physReg;

    set rmap.n += 1;

}

/// Remove a mapping from the register map and return its physical register.

fn rmapRemove(rmap: *mut RegMap, virtReg: u32) -> ?gen::Reg {

    for i in 0..rmap.n {

        if rmap.virtRegs[i] == virtReg {

            let phys = rmap.physRegs[i];

            // Swap with last and decrement.

            set rmap.n -= 1;

            if i < rmap.n {

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

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

            }

            return phys;

        }

    }

    return nil;

}

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

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

    for i in 0..pool.len {

        let r = pool[i];

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

            bitset::put(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: *[gen::Reg],

    calleeSaved: *[gen::Reg],

    spillInfo: *spill::SpillInfo

) -> gen::Reg {

    // 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";

}

/// Record the current index; forward traversal leaves the last operand use.

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

    let ctx = ctxPtr as *mut LastUseCtx;

    set ctx.lastUse[reg.n] = ctx.index;

}

/// Free operands with no later block use or live-out use, then allocate missing uses.

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

    let ctx = ctxPtr as *mut InstrCtx;

    if not (bitset::contains(&ctx.live.liveOut[ctx.blockIdx], reg.n) or ctx.lastUse[reg.n] > ctx.instrIdx) {

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

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

        }

    }

    assert reg.n < ctx.assignments.len, "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 {

        set ctx.assignments[reg.n] = rallocReg(

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

            ctx.calleeSaved, ctx.spillInfo

        );

    }

}