//! demutate: every static-key field write becomes a call to a function that
//! RETURNS a fresh object.
//!
//! There is no grouping and no analysis of what surrounds a write. Each one
//! stands alone, because its two arguments are the call's two arguments and
//! the enclosing code already holds both:
//!
//!     define_field @F [v, o] -> [n]     n = <define F>(o, v)
//!     put_field    @F [v, o]            n = <put F>(o, v) ; store n back
//!
//! `object` is left as it is. It already makes a fresh object, so it needs no
//! function; infer reads it as `Type::Empty`, the base every shape builds on.
//!
//! The lifted functions have no body. `CodeKind::DefineField` / `PutField`
//! name the operation and the field, and infer gives each a type rule, so ONE
//! function serves every write of that field in the program. Their ids run on
//! after the real functions.

use crate::{
    infer::{CodeKind, InferFn, Type},
    parse::{AtomRef, ImmOperand, JsValue, opcodes::OpCode},
    regify::{RegFn, RegInstr},
};

pub fn demutate(reg_fns: Vec<RegFn>) -> Vec<InferFn> {
    // The lifted functions take the ids after the real ones, so a write can
    // name the one it needs before the list of them is finished
    let first_lifted = reg_fns.len();
    let mut lifted = vec![];

    let mut fns = reg_fns
        .into_iter()
        .map(|f| rewrite_fn(f, first_lifted, &mut lifted))
        .collect::<Vec<InferFn>>();

    fns.extend(lifted.into_iter().map(|code| InferFn {
        name: lifted_name(&code),
        cpool: vec![],
        code,
        sigs: vec![],
        types: vec![],
        blocks: vec![],
    }));

    fns
}

/// One scan of the code. Everything that is not a static-key write is moved
/// across untouched; a write becomes a call, and a `put_field` a call and a
/// store.
fn rewrite_fn(mut reg_fn: RegFn, first_lifted: usize, lifted: &mut Vec<CodeKind>) -> InferFn {
    let code = std::mem::take(&mut reg_fn.code);
    let mut new = Vec::with_capacity(code.len());

    for instr in code {
        let Some(field) = static_key(&instr) else {
            new.push(instr); // Not a field write, so it passes through
            continue;
        };

        let define = instr.op == OpCode::DefineField;

        // Where the object lives says where its replacement goes. The
        // instruction that made it is already in `new`.
        let object = instr.args[1];
        let slot = (!define).then(|| frame_slot(&new, object)).flatten();

        let callee = intern_const(
            &mut reg_fn.cpool,
            first_lifted + intern_fn(lifted, instr.op, field),
        );

        let dst = match slot {
            // define_field hands the object back, so the call keeps its dst
            _ if define => instr.dsts[0],

            // regify does not rename a local, so the local IS the object's
            // value: defining it again is the store, and costs no instruction
            Some(Source::Loc(_)) => object,

            // Anything else needs a value of its own to be written back from,
            // and with no slot at all the write is lost -- `f().x = 1` does that
            _ => reg_fn.fresh_val(),
        };

        new.push(call(callee, dst, instr.args).starting(instr.block_start));

        if let Some(slot) = slot.filter(|s| !matches!(s, Source::Loc(_))) {
            new.push(store(slot, dst));
        }
    }

    InferFn {
        types: vec![Type::Unknown; reg_fn.num_vals()],
        blocks: reg_fn.blocks,
        name: reg_fn.name,
        cpool: reg_fn.cpool,
        code: CodeKind::Code(new),
        sigs: vec![],
    }
}

/// The field a write names, when the name is fixed at compile time. A computed
/// key is `define_array_el` / `put_array_el`, and is left alone.
fn static_key(i: &RegInstr) -> Option<AtomRef> {
    let is_write = matches!(i.op, OpCode::DefineField | OpCode::PutField) && i.args.len() == 2;

    match i.imm {
        ImmOperand::Atom(a) if is_write => Some(a),
        _ => None,
    }
}

/// The index of the function for this write, adding it if it is the first of
/// its kind. One function serves every write of one field by one opcode.
fn intern_fn(lifted: &mut Vec<CodeKind>, op: OpCode, field: AtomRef) -> usize {
    let wanted = match op {
        OpCode::DefineField => CodeKind::DefineField(field),
        _ => CodeKind::PutField(field),
    };

    let same = |c: &CodeKind| match (c, &wanted) {
        (CodeKind::DefineField(a), CodeKind::DefineField(b)) => a == b,
        (CodeKind::PutField(a), CodeKind::PutField(b)) => a == b,
        _ => false,
    };

    lifted.iter().position(same).unwrap_or_else(|| {
        lifted.push(wanted);
        lifted.len() - 1
    })
}

/// The cpool slot holding function `id`, adding it if it is not there yet. A
/// call names its function the way every other reference to one does.
fn intern_const(cpool: &mut Vec<JsValue>, id: usize) -> u32 {
    let at = cpool
        .iter()
        .position(|c| matches!(c, JsValue::FnById(f) if *f == id))
        .unwrap_or_else(|| {
            cpool.push(JsValue::FnById(id));
            cpool.len() - 1
        });

    at as u32
}

/// The frame slot an object lives in, from whichever instruction defined it.
fn frame_slot(code: &[RegInstr], object: usize) -> Option<Source> {
    code.iter()
        .find(|i| i.dsts.contains(&object))
        .and_then(frame_load)
}

/// A frame slot, and the two instructions that read and write it.
#[derive(Clone, Copy, PartialEq, Eq)]
enum Source {
    Arg(u16),
    Loc(u16),
    VarRef(u16),
}

// ── emitted instructions ───────────────────────────────────────────────────

/// A call to a function fixed at rewrite time. The callee is named by the
/// immediate rather than pushed, so the call fits where the write stood.
fn call(callee: u32, dst: usize, args: Vec<usize>) -> RegInstr {
    RegInstr::new(OpCode::Call, ImmOperand::Const(callee), vec![dst], args)
}

fn store(slot: Source, value: usize) -> RegInstr {
    let (op, imm) = match slot {
        Source::Arg(k) => (OpCode::PutArg, ImmOperand::Arg(k)),
        Source::Loc(k) => (OpCode::PutLoc, ImmOperand::Loc(k)),
        Source::VarRef(k) => (OpCode::PutVarRef, ImmOperand::VarRef(k)),
    };

    RegInstr::new(op, imm, vec![], vec![value])
}

fn lifted_name(code: &CodeKind) -> String {
    match code {
        CodeKind::DefineField(a) => format!("<define {a:?}>"),
        CodeKind::PutField(a) => format!("<put {a:?}>"),
        CodeKind::Code(_) => unreachable!("a lifted function has no body"),
    }
}

// ── operand helpers ────────────────────────────────────────────────────────

/// The frame slot an instruction reads. Only plain one-value loads count;
/// `get_loc0_loc1` pushes two.
fn frame_load(i: &RegInstr) -> Option<Source> {
    let slot = |imm| match imm {
        ImmOperand::Loc(k) | ImmOperand::Arg(k) | ImmOperand::VarRef(k) => k,
        _ => unreachable!("a frame load without a slot operand"),
    };

    match i.op {
        OpCode::GetArg0 => Some(Source::Arg(0)),
        OpCode::GetArg1 => Some(Source::Arg(1)),
        OpCode::GetArg2 => Some(Source::Arg(2)),
        OpCode::GetArg3 => Some(Source::Arg(3)),
        OpCode::GetArg => Some(Source::Arg(slot(i.imm))),

        OpCode::GetLoc0 => Some(Source::Loc(0)),
        OpCode::GetLoc1 => Some(Source::Loc(1)),
        OpCode::GetLoc2 => Some(Source::Loc(2)),
        OpCode::GetLoc3 => Some(Source::Loc(3)),
        // Every local op names the same value, so any of them identifies it
        OpCode::GetLoc
        | OpCode::GetLoc8
        | OpCode::GetLocCheck
        | OpCode::PutLoc
        | OpCode::PutLoc8
        | OpCode::PutLocCheck
        | OpCode::PutLocCheckInit
        | OpCode::SetLoc
        | OpCode::SetLoc8
        | OpCode::SetLocUninitialized => Some(Source::Loc(slot(i.imm))),

        OpCode::PutLoc0 | OpCode::SetLoc0 => Some(Source::Loc(0)),
        OpCode::PutLoc1 | OpCode::SetLoc1 => Some(Source::Loc(1)),
        OpCode::PutLoc2 | OpCode::SetLoc2 => Some(Source::Loc(2)),
        OpCode::PutLoc3 | OpCode::SetLoc3 => Some(Source::Loc(3)),

        OpCode::GetVarRef0 => Some(Source::VarRef(0)),
        OpCode::GetVarRef1 => Some(Source::VarRef(1)),
        OpCode::GetVarRef2 => Some(Source::VarRef(2)),
        OpCode::GetVarRef3 => Some(Source::VarRef(3)),
        OpCode::GetVarRef | OpCode::GetVarRefCheck => Some(Source::VarRef(slot(i.imm))),

        _ => None,
    }
}