spectral_vm 0.1.6

/*
 * ═══════════════════════════════════════════════════════════════════════════
 * TECHNICAL MANIFEST: compiler.rs
 * SOVEREIGN SPECTRAL ROLE: Circuit Compilation & Optimization Pipeline
 * ═══════════════════════════════════════════════════════════════════════════
 *
 * COMPLEXITY: O(n log n) [Optimization] | FIELD: Goldilocks
 * DOMAIN: Hybrid (IR -> Spectral)
 *
 * ARCHITECTURAL INVARIANTS:
 * - SSA Form: Single assignment for all intermediate values
 * - Live Intervals: Computed for Linear Scan Allocation
 * - Spilling: Automatic spill to stack (Addr 1000+) when registers exhausted
 *
 * ═══════════════════════════════════════════════════════════════════════════
 */

use crate::vm::SpectralOp;
use std::collections::{BTreeMap, HashMap, HashSet};

/// SSA Value Identifier
pub type ValueId = usize;
/// Basic Block Identifier
pub type BlockId = usize;
/// Virtual Register or Stack Slot
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Location {
    Register(u64),
    Stack(u64), // Stack Offset (Memory Address = BASE + Offset)
}

const STACK_BASE: u64 = 1000;
const REG_COUNT: u64 = 28; // Reserve R28-R31 for Scratch/Control
const SCRATCH_REG_1: u64 = 28;
const SCRATCH_REG_2: u64 = 29;
const JMP_REG: u64 = 30;
const BR_REG: u64 = 31;

/// High-Level Spectral Intermediate Representation Operations
#[derive(Debug, Clone, PartialEq)]
pub enum IrOp {
    Const(u64),
    Add(ValueId, ValueId),
    Sub(ValueId, ValueId),
    Mul(ValueId, ValueId),
    Load(ValueId),
    Store(ValueId, ValueId),
    Beq(ValueId, ValueId, BlockId),
    Jmp(BlockId),
    Halt,
}

#[derive(Debug, Clone)]
pub struct Instruction {
    pub id: usize, // Global Instruction Index (for Liveness)
    pub result_id: Option<ValueId>,
    pub op: IrOp,
}

#[derive(Debug, Clone)]
pub struct BasicBlock {
    pub id: BlockId,
    pub instructions: Vec<Instruction>,
}

#[derive(Debug)]
pub struct SpectralIR {
    pub blocks: Vec<BasicBlock>,
    pub value_counter: usize,
    pub block_counter: usize,
    pub inst_counter: usize,
}

impl SpectralIR {
    pub fn new() -> Self {
        Self {
            blocks: vec![BasicBlock {
                id: 0,
                instructions: vec![],
            }],
            value_counter: 0,
            block_counter: 1,
            inst_counter: 0,
        }
    }

    pub fn next_value(&mut self) -> ValueId {
        let id = self.value_counter;
        self.value_counter += 1;
        id
    }

    pub fn next_block(&mut self) -> BlockId {
        let id = self.block_counter;
        self.block_counter += 1;
        self.blocks.push(BasicBlock {
            id,
            instructions: vec![],
        });
        id
    }

    pub fn append_inst(&mut self, block_id: BlockId, op: IrOp) -> Option<ValueId> {
        let result_id = match op {
            IrOp::Store(_, _) | IrOp::Jmp(_) | IrOp::Beq(_, _, _) | IrOp::Halt => None,
            _ => Some(self.next_value()),
        };

        let inst = Instruction {
            id: self.inst_counter,
            result_id,
            op,
        };
        self.inst_counter += 1;

        if let Some(block) = self.blocks.iter_mut().find(|b| b.id == block_id) {
            block.instructions.push(inst);
        }
        result_id
    }
}

/// Liveness Interval for a Value
#[derive(Debug, Clone)]
struct LiveInterval {
    value_id: ValueId,
    start: usize,
    end: usize,
}

pub struct CircuitCompiler {
    pub ir: SpectralIR,
    pub allocation: HashMap<ValueId, Location>,
    pub spill_count: usize,
}

impl CircuitCompiler {
    pub fn new() -> Self {
        Self {
            ir: SpectralIR::new(),
            allocation: HashMap::new(),
            spill_count: 0,
        }
    }

    /// Linear Scan Register Allocation
    pub fn allocate_registers(&mut self) {
        // 1. Compute Liveness Intervals
        let mut intervals: HashMap<ValueId, LiveInterval> = HashMap::new();

        for block in &self.ir.blocks {
            for inst in &block.instructions {
                // Definition
                if let Some(res) = inst.result_id {
                    intervals.entry(res).or_insert(LiveInterval {
                        value_id: res,
                        start: inst.id,
                        end: inst.id,
                    });
                }

                // Usage (Extend End)
                let mut touch = |vid: ValueId| {
                    if let Some(interval) = intervals.get_mut(&vid) {
                        if inst.id > interval.end {
                            interval.end = inst.id;
                        }
                    }
                };

                match inst.op {
                    IrOp::Add(a, b) | IrOp::Sub(a, b) | IrOp::Mul(a, b) => {
                        touch(a);
                        touch(b);
                    }
                    IrOp::Store(a, b) => {
                        touch(a);
                        touch(b);
                    }
                    IrOp::Load(a) => {
                        touch(a);
                    }
                    IrOp::Beq(a, b, _) => {
                        touch(a);
                        touch(b);
                    }
                    _ => {}
                }
            }
        }

        // 2. Sort Intervals by Start
        let mut sorted_intervals: Vec<LiveInterval> = intervals.into_values().collect();
        sorted_intervals.sort_by_key(|i| i.start);

        // 3. Linear Scan
        let mut free_regs: Vec<u64> = (1..=REG_COUNT).rev().collect(); // R1..R28
        let mut active: Vec<LiveInterval> = Vec::new(); // Sorted by End? No, we just scan.
        let mut val_loc: HashMap<ValueId, Location> = HashMap::new();
        let mut next_stack_slot = 0;

        for interval in sorted_intervals {
            // Expire Old Intervals
            active.retain(|act| {
                if act.end < interval.start {
                    // Release Register
                    if let Some(Location::Register(r)) = val_loc.get(&act.value_id) {
                        free_regs.push(*r);
                    }
                    false // Remove from active
                } else {
                    true
                }
            });

            // Allocate
            if let Some(reg) = free_regs.pop() {
                // Assign Register
                val_loc.insert(interval.value_id, Location::Register(reg));
                active.push(interval);
            } else {
                // Spill logic
                let mut victim_idx = usize::MAX;
                let mut max_end = interval.end;
                let mut found_better = false;

                for (idx, act) in active.iter().enumerate() {
                    if act.end > max_end {
                        max_end = act.end;
                        victim_idx = idx;
                        found_better = true;
                    }
                }

                if found_better {
                    // Spill Active Victim
                    let victim_val = active[victim_idx].value_id;
                    let loc = val_loc[&victim_val];
                    let reg = match loc {
                        Location::Register(r) => r,
                        _ => panic!("Active not in reg"),
                    };

                    // Assign Stack to Victim
                    val_loc.insert(victim_val, Location::Stack(next_stack_slot));
                    next_stack_slot += 1;
                    self.spill_count += 1;

                    // Steal Register for Current
                    val_loc.insert(interval.value_id, Location::Register(reg));

                    // Update Active list
                    active.remove(victim_idx);
                    active.push(interval);
                } else {
                    // Spill Current
                    val_loc.insert(interval.value_id, Location::Stack(next_stack_slot));
                    next_stack_slot += 1;
                    self.spill_count += 1;
                }
            }
        }

        self.allocation = val_loc;
    }

    pub fn codegen(&mut self) -> Vec<u64> {
        self.allocate_registers();

        let mut block_ops_map: BTreeMap<BlockId, Vec<u64>> = BTreeMap::new();

        for block in &self.ir.blocks {
            let mut b_ops = Vec::new();

            for inst in &block.instructions {
                let opers_out = &mut b_ops;

                let mut load_operand = |vid: ValueId, scratch_reg: u64| -> u64 {
                    match *self.allocation.get(&vid).expect("Alloc missing") {
                        Location::Register(r) => r,
                        Location::Stack(offset) => {
                            // Emit S_MUL Scratch, 0 (Clear)
                            opers_out.extend_from_slice(&[
                                SpectralOp::S_MUL as u64,
                                scratch_reg,
                                0,
                            ]);
                            // Emit S_ADDI Scratch, offset+BASE
                            let addr = STACK_BASE + offset;
                            opers_out.extend_from_slice(&[
                                SpectralOp::S_ADDI as u64,
                                scratch_reg,
                                addr,
                            ]);
                            // Emit S_LOAD Scratch, 0 (Addr in Scratch, Val -> Scratch)
                            opers_out.extend_from_slice(&[
                                SpectralOp::S_LOAD as u64,
                                scratch_reg,
                                0,
                            ]);
                            scratch_reg
                        }
                    }
                };

                let get_target_reg = |vid: ValueId| -> u64 {
                    match *self.allocation.get(&vid).expect("Alloc missing") {
                        Location::Register(r) => r,
                        Location::Stack(_) => SCRATCH_REG_1, // Use Scratch for result computation
                    }
                };

                let commit_result = |vid: ValueId, target_reg: u64, out: &mut Vec<u64>| {
                    if let Location::Stack(offset) =
                        *self.allocation.get(&vid).expect("Alloc missing")
                    {
                        // Store target_reg -> Stack[offset]
                        // Need Address in another register. Use SCRATCH_REG_2.

                        // 1. Clear Addr Reg
                        out.extend_from_slice(&[SpectralOp::S_MUL as u64, SCRATCH_REG_2, 0]);
                        // 2. Set Addr
                        let addr = STACK_BASE + offset;
                        out.extend_from_slice(&[SpectralOp::S_ADDI as u64, SCRATCH_REG_2, addr]);
                        // 3. Store (Addr=R_2, Val=Target)
                        out.extend_from_slice(&[
                            SpectralOp::S_STORE as u64,
                            SCRATCH_REG_2,
                            target_reg,
                        ]);
                    }
                };

                match &inst.op {
                    IrOp::Const(c) => {
                        let res = inst.result_id.expect("Const operation must have result ID");
                        let target = get_target_reg(res);
                        // S_MUL Target, 0
                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, target, 0]);
                        // S_ADDI Target, c
                        b_ops.extend_from_slice(&[SpectralOp::S_ADDI as u64, target, *c]);
                        commit_result(res, target, &mut b_ops);
                    }
                    IrOp::Add(a, b) => {
                        let res = inst.result_id.expect("Add operation must have result ID");
                        let r_a = load_operand(*a, SCRATCH_REG_1);
                        let r_b = load_operand(*b, SCRATCH_REG_2);

                        let target = get_target_reg(res);

                        if target == r_a {
                            // S_ADD Target, B (Target+=B)
                            b_ops.extend_from_slice(&[SpectralOp::S_ADD as u64, target, r_b]);
                        } else {
                            // Zero Target
                            b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, target, 0]);
                            b_ops.extend_from_slice(&[SpectralOp::S_ADD as u64, target, r_a]);
                            b_ops.extend_from_slice(&[SpectralOp::S_ADD as u64, target, r_b]);
                        }
                        commit_result(res, target, &mut b_ops);
                    }
                    IrOp::Mul(a, b) => {
                        let res = inst.result_id.expect("Instruction must have result ID");
                        let r_a = load_operand(*a, SCRATCH_REG_1);
                        let r_b = load_operand(*b, SCRATCH_REG_2);
                        let target = get_target_reg(res);

                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, target, 0]);
                        b_ops.extend_from_slice(&[SpectralOp::S_ADD as u64, target, r_a]);
                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, target, r_b]);
                        commit_result(res, target, &mut b_ops);
                    }
                    IrOp::Load(addr) => {
                        let res = inst.result_id.expect("Instruction must have result ID");
                        let r_addr = load_operand(*addr, SCRATCH_REG_2);
                        let target = get_target_reg(res);

                        // Target = Move Addr
                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, target, 0]);
                        b_ops.extend_from_slice(&[SpectralOp::S_ADD as u64, target, r_addr]);
                        // Load (Destructive: Target = Mem[Target])
                        b_ops.extend_from_slice(&[SpectralOp::S_LOAD as u64, target, 0]);
                        commit_result(res, target, &mut b_ops);
                    }
                    IrOp::Store(addr, val) => {
                        let r_addr = load_operand(*addr, SCRATCH_REG_1);
                        let r_val = load_operand(*val, SCRATCH_REG_2);
                        b_ops.extend_from_slice(&[SpectralOp::S_STORE as u64, r_addr, r_val]);
                    }
                    IrOp::Beq(a, b, target_block) => {
                        let r_a = load_operand(*a, SCRATCH_REG_1);
                        let r_b = load_operand(*b, SCRATCH_REG_2);

                        // Set BR_REG (31)
                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, BR_REG, 0]);
                        b_ops.extend_from_slice(&[
                            SpectralOp::S_ADDI as u64,
                            BR_REG,
                            *target_block as u64,
                        ]);
                        b_ops.extend_from_slice(&[SpectralOp::S_BEQ as u64, r_a, r_b]);
                    }
                    IrOp::Jmp(target_block) => {
                        // Set JMP_REG (30)
                        b_ops.extend_from_slice(&[SpectralOp::S_MUL as u64, JMP_REG, 0]);
                        b_ops.extend_from_slice(&[
                            SpectralOp::S_ADDI as u64,
                            JMP_REG,
                            *target_block as u64,
                        ]);
                        b_ops.extend_from_slice(&[SpectralOp::S_JMP as u64, JMP_REG, 0]);
                    }
                    IrOp::Halt => {
                        b_ops.extend_from_slice(&[SpectralOp::S_HALT as u64, 0, 0]);
                    }
                    _ => {}
                }
            }
            block_ops_map.insert(block.id, b_ops);
        }

        // Pass 2: Patching & Concatenation
        let mut final_ops = Vec::new();
        let mut offsets = HashMap::new();
        let mut cursor = 0;

        for (id, ops_vec) in &block_ops_map {
            offsets.insert(*id, cursor);
            cursor += ops_vec.len() as u64;
        }

        for (_, ops_vec) in block_ops_map {
            let mut i = 0;
            while i < ops_vec.len() {
                let op = ops_vec[i];
                let arg1 = ops_vec[i + 1];
                let arg2 = ops_vec[i + 2];

                // Patch Jumps (S_ADDI to R30/R31)
                if op == SpectralOp::S_ADDI as u64 {
                    if arg1 == BR_REG || arg1 == JMP_REG {
                        // Arg2 is BlockId
                        let abs_addr = offsets.get(&(arg2 as usize)).expect("Target block missing");
                        final_ops.push(op);
                        final_ops.push(arg1);
                        final_ops.push(*abs_addr);
                    } else {
                        final_ops.push(op);
                        final_ops.push(arg1);
                        final_ops.push(arg2);
                    }
                } else {
                    final_ops.push(op);
                    final_ops.push(arg1);
                    final_ops.push(arg2);
                }
                i += 3;
            }
        }

        final_ops
    }

    pub fn compile_fibonacci(&mut self, n: u64) {
        // Blocks: 0: Init, 1: LoopHeader, 2: LoopBody, 3: Exit
        let init_block = 0;
        let loop_header = self.ir.next_block(); // 1
        let loop_body = self.ir.next_block(); // 2
        let exit_block = self.ir.next_block(); // 3

        // --- Block 0: Init ---
        let val_a = self.ir.append_inst(init_block, IrOp::Const(0)).expect("IR append should succeed");
        let val_b = self.ir.append_inst(init_block, IrOp::Const(1)).expect("IR append should succeed");
        let val_i = self.ir.append_inst(init_block, IrOp::Const(1)).expect("IR append should succeed");
        let val_n = self.ir.append_inst(init_block, IrOp::Const(n)).expect("IR append should succeed");

        let addr_a = self.ir.append_inst(init_block, IrOp::Const(100)).expect("IR append should succeed");
        let addr_b = self.ir.append_inst(init_block, IrOp::Const(101)).expect("IR append should succeed");
        let addr_i = self.ir.append_inst(init_block, IrOp::Const(102)).expect("IR append should succeed");
        let addr_n = self.ir.append_inst(init_block, IrOp::Const(103)).expect("IR append should succeed");

        self.ir.append_inst(init_block, IrOp::Store(addr_a, val_a));
        self.ir.append_inst(init_block, IrOp::Store(addr_b, val_b));
        self.ir.append_inst(init_block, IrOp::Store(addr_i, val_i));
        self.ir.append_inst(init_block, IrOp::Store(addr_n, val_n));
        self.ir.append_inst(init_block, IrOp::Jmp(loop_header));

        // --- Block 1: Loop Header ---
        let l_addr_i = self.ir.append_inst(loop_header, IrOp::Const(102)).expect("IR construction should succeed");
        let l_addr_n = self.ir.append_inst(loop_header, IrOp::Const(103)).expect("IR construction should succeed");

        let curl_i = self
            .ir
            .append_inst(loop_header, IrOp::Load(l_addr_i))
            .expect("IR construction should succeed");
        let curl_n = self
            .ir
            .append_inst(loop_header, IrOp::Load(l_addr_n))
            .expect("IR construction should succeed");
        self.ir
            .append_inst(loop_header, IrOp::Beq(curl_i, curl_n, exit_block));
        self.ir.append_inst(loop_header, IrOp::Jmp(loop_body));

        // --- Block 2: Loop Body ---
        let l2_addr_a = self.ir.append_inst(loop_body, IrOp::Const(100)).expect("IR construction should succeed");
        let l2_addr_b = self.ir.append_inst(loop_body, IrOp::Const(101)).expect("IR construction should succeed");

        let cur_a = self
            .ir
            .append_inst(loop_body, IrOp::Load(l2_addr_a))
            .expect("IR construction should succeed");
        let cur_b = self
            .ir
            .append_inst(loop_body, IrOp::Load(l2_addr_b))
            .expect("IR construction should succeed");

        let t = self
            .ir
            .append_inst(loop_body, IrOp::Add(cur_a, cur_b))
            .expect("IR construction should succeed");

        let l2_st_addr_a = self.ir.append_inst(loop_body, IrOp::Const(100)).expect("IR construction should succeed");
        let l2_st_addr_b = self.ir.append_inst(loop_body, IrOp::Const(101)).expect("IR construction should succeed");

        self.ir
            .append_inst(loop_body, IrOp::Store(l2_st_addr_a, cur_b));
        self.ir.append_inst(loop_body, IrOp::Store(l2_st_addr_b, t));

        let l2_addr_i = self.ir.append_inst(loop_body, IrOp::Const(102)).expect("IR construction should succeed");
        let cur_i_2 = self
            .ir
            .append_inst(loop_body, IrOp::Load(l2_addr_i))
            .expect("IR construction should succeed");
        let one = self.ir.append_inst(loop_body, IrOp::Const(1)).expect("IR construction should succeed");
        let next_i = self
            .ir
            .append_inst(loop_body, IrOp::Add(cur_i_2, one))
            .expect("IR construction should succeed");

        let l2_st_addr_i = self.ir.append_inst(loop_body, IrOp::Const(102)).expect("IR construction should succeed");
        self.ir
            .append_inst(loop_body, IrOp::Store(l2_st_addr_i, next_i));
        self.ir.append_inst(loop_body, IrOp::Jmp(loop_header));

        // --- Block 3: Exit ---
        self.ir.append_inst(exit_block, IrOp::Halt);
    }

    pub fn optimize_dce(&mut self) {
        // Simple DCE
        let mut used = HashSet::new();
        loop {
            let start_len = used.len();
            for block in &self.ir.blocks {
                for inst in &block.instructions {
                    let is_critical = match inst.op {
                        IrOp::Store(_, _) | IrOp::Beq(_, _, _) | IrOp::Jmp(_) | IrOp::Halt => true,
                        _ => false,
                    };
                    let result_used = if let Some(res) = inst.result_id {
                        used.contains(&res)
                    } else {
                        false
                    };
                    if is_critical || result_used {
                        match inst.op {
                            IrOp::Add(a, b) | IrOp::Sub(a, b) | IrOp::Mul(a, b) => {
                                used.insert(a);
                                used.insert(b);
                            }
                            IrOp::Store(a, b) => {
                                used.insert(a);
                                used.insert(b);
                            }
                            IrOp::Load(a) => {
                                used.insert(a);
                            }
                            IrOp::Beq(a, b, _) => {
                                used.insert(a);
                                used.insert(b);
                            }
                            _ => {}
                        }
                    }
                }
            }
            if used.len() == start_len {
                break;
            }
        }
        for block in &mut self.ir.blocks {
            block.instructions.retain(|inst| {
                if let Some(res) = inst.result_id {
                    used.contains(&res)
                } else {
                    true
                }
            });
        }
    }
}