mod basic_block;
mod walk;

use crate::{
    flatten::FlatFn,
    libqjs::parse::{ImmOperand, JsValue, opcodes::OpCode},
};

pub use basic_block::BasicBlock;
use basic_block::find_basic_blocks;
use walk::walk_block;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Operand {
    Local(usize),
    Arg(usize),
    Reg(usize),
}

#[derive(Debug, Clone)]
pub struct RegInstr {
    pub dsts: Vec<Operand>,
    pub op: OpCode,
    pub imm: ImmOperand,
    pub args: Vec<Operand>,
}

impl RegInstr {
    fn nop() -> Self {
        Self {
            dsts: vec![],
            op: OpCode::Nop,
            imm: ImmOperand::None,
            args: vec![],
        }
    }
}

#[derive(Debug)]
pub struct RegFn {
    pub arg_count: usize,
    pub local_count: usize,
    pub reg_count: usize,
    pub cpool: Vec<JsValue>,
    pub code: Vec<RegInstr>,
    pub blocks: Vec<BasicBlock>,
}

pub fn regify(flat_fns: Vec<FlatFn>) -> Vec<RegFn> {
    flat_fns.into_iter().map(RegFn::from_flat).collect()
}

impl RegFn {
    /// The locals and args that a closure shares: the locations of its `define_field_ref`s
    /// (see transform/closure.rs)
    pub fn captured_slots(&self) -> Vec<Operand> {
        let refs = self.code.iter().filter(|i| i.op == OpCode::DefineFieldRef);
        let mut slots: Vec<Operand> = refs
            .map(|i| i.args[0])
            .filter(|slot| !matches!(slot, Operand::Reg(_)))
            .collect();

        slots.sort_unstable();
        slots.dedup();
        slots
    }

    /// Joins: the stack that a block gets at entry is new regs, made when the walk first gets
    /// to the block. Each predecessor copies its own stack into them, in its postfix. So every
    /// block has one entry stack, whichever way control comes in (this does the job of phi
    /// nodes in SSA)
    ///
    /// A block with two successors copies into the entry regs of both, before its jump. The
    /// copies for the edge that is not taken do no harm: a block's entry regs are only read
    /// in that block, and every way into it writes them first
    fn from_flat(f: FlatFn) -> Self {
        let mut blocks = find_basic_blocks(&f.code);
        let mut out = vec![RegInstr::nop(); f.code.len()];
        let mut next_reg = 0;

        let mut entry_stacks: Vec<Option<Vec<Operand>>> = vec![None; blocks.len()];
        entry_stacks[0] = Some(vec![]);

        // Walk blocks in CFG order, starting from block 0. A block that the walk never
        // gets to is dead code, and stays `nop`
        let mut worklist = vec![0];

        while let Some(b) = worklist.pop() {
            let mut stack = entry_stacks[b].clone().unwrap();
            walk_block(&blocks[b], &f.code, &mut stack, &mut next_reg, &mut out);

            for s in blocks[b].succ.clone() {
                let entry = entry_stacks[s].get_or_insert_with(|| {
                    worklist.push(s);
                    stack.iter().map(|_| fresh_reg(&mut next_reg)).collect()
                });

                assert_eq!(
                    entry.len(),
                    stack.len(),
                    "a join of stacks of different depths"
                );

                let copies = stack
                    .iter()
                    .zip(entry.iter())
                    .filter(|(src, dst)| src != dst);

                for (&src, &dst) in copies {
                    blocks[b].postfix.push(RegInstr {
                        dsts: vec![dst],
                        op: OpCode::SetLoc,
                        imm: ImmOperand::None,
                        args: vec![src],
                    });
                }
            }

            // The copies run one after the other, so no copy may write a reg that a later
            // copy reads. A block's stack only holds its own entry regs and the regs it made,
            // so that could only be a block that jumps to itself and moves its stack around
            let postfix = &blocks[b].postfix;
            assert!(
                postfix
                    .iter()
                    .all(|c| postfix.iter().all(|d| c.args[0] != d.dsts[0])),
                "the copies of a join overlap"
            );
        }

        mark_puts(&mut out);

        Self {
            arg_count: f.arg_count,
            local_count: f.local_count,
            reg_count: next_reg,
            cpool: f.cpool,
            code: out,
            blocks,
        }
    }
}

fn fresh_reg(next_reg: &mut usize) -> Operand {
    let r = *next_reg;
    *next_reg += 1;
    Operand::Reg(r)
}

/// A define_field into an object that no instr of this fn made (a class field on `this`, a
/// static field on the class) is a write like any other: a put_field. So a define_field is
/// always a write into an object before it goes anywhere, and its field is never absent
fn mark_puts(code: &mut [RegInstr]) {
    let made: Vec<Operand> = (code.iter())
        .filter(|i| matches!(i.op, OpCode::Object | OpCode::Fclosure | OpCode::Fclosure8))
        .map(|i| i.dsts[0])
        .collect();

    for instr in code {
        if instr.op == OpCode::DefineField && !made.contains(&instr.args[1]) {
            instr.op = OpCode::PutField;
        }
    }
}