//! This module converts one basic block's instrs from stack to register form using a simulated stack
//!
//! The output reg form instrs are structured as:
//! `dsts = op(imm, args)`
//!
//! Every instr's pushes allocates new register operands which are pushed on the stack and form the
//! output instr's `dsts`, the instr's pops pop operands from the stack and form the output instr's
//! `args`
//!
//! Each instr's `op` and `imm` is preserved except for some special cases
//!
//! `set_loc` is henceforth used to move data from one operand (local, arg, reg) to another

#![allow(non_snake_case)]

use super::{BasicBlock, Operand, RegInstr};
use crate::libqjs::parse::{ImmOperand, Instr, opcodes::OpCode};

/// Walk a basic block's instructions, converting each instr to register form
pub fn walk_block(
    block: &BasicBlock,
    code: &[Instr],
    stack: &mut Vec<Operand>,
    next_reg: &mut usize,
    out: &mut [RegInstr],
) {
    let mut fresh_reg = || {
        let r = *next_reg;
        *next_reg += 1;
        Operand::Reg(r)
    };

    for instr_idx in block.start..block.end {
        let instr = &code[instr_idx];
        let out_instr = &mut out[instr_idx];

        if try_handle_named_read(instr, stack, &mut fresh_reg, out_instr) {
            continue;
        }

        if try_handle_named_write(instr, stack, out_instr) {
            continue;
        }

        if try_handle_shuffle(instr, stack, out_instr) {
            continue;
        }

        handle_normal_case(instr, stack, &mut fresh_reg, out_instr);

        if instr.op == OpCode::DefineFieldRef {
            let prev = std::mem::replace(&mut out[instr_idx - 1], RegInstr::nop());
            resolve_location(&prev, &mut out[instr_idx]);
        }
    }
}

/// Handles a "named read" i.e. a read of a local or arg (and not a stack value)
///
/// We simulate pushing the operand to stack, and replace the instr with `SetLoc`, which
/// is our catch-all for a data-move instr
fn try_handle_named_read(
    instr: &Instr,
    stack: &mut Vec<Operand>,
    mut fresh_reg: impl FnMut() -> Operand,
    out: &mut RegInstr,
) -> bool {
    use OpCode::*;

    let arg = match instr.op {
        GetLoc0 => Operand::Local(0),
        GetLoc1 => Operand::Local(1),
        GetLoc2 => Operand::Local(2),
        GetLoc3 => Operand::Local(3),

        GetLoc | GetLoc8 | GetLocCheck => match instr.imm {
            ImmOperand::Loc(k) => Operand::Local(k as usize),
            _ => unreachable!(),
        },

        GetArg0 => Operand::Arg(0),
        GetArg1 => Operand::Arg(1),
        GetArg2 => Operand::Arg(2),
        GetArg3 => Operand::Arg(3),

        GetArg => match instr.imm {
            ImmOperand::Arg(k) => Operand::Arg(k as usize),
            _ => unreachable!(),
        },

        _ => return false,
    };

    let dst = fresh_reg();

    // Converted SetLoc, reading from arg, writing to dst
    *out = RegInstr {
        dsts: vec![dst],
        op: SetLoc,
        imm: ImmOperand::None,
        args: vec![arg],
    };

    // Simulate stack push
    stack.push(dst);
    true
}

/// Handles a "named write" i.e. a write of a local or arg (and not a stack value)
///
/// We simulate popping the operand off stack, and replace the instr with `SetLoc`, which
/// is our catch-all for a data-move instr
///
/// A distinction needs to be made between `Set*` and `Put*` since `Set*` instrs leave
/// the value-to-write on the stack and `Put*` pops it off
fn try_handle_named_write(instr: &Instr, stack: &mut Vec<Operand>, out: &mut RegInstr) -> bool {
    use OpCode::*;

    let (dst, preserved) = match instr.op {
        PutLoc0 => (Operand::Local(0), false),
        PutLoc1 => (Operand::Local(1), false),
        PutLoc2 => (Operand::Local(2), false),
        PutLoc3 => (Operand::Local(3), false),

        PutLoc | PutLoc8 | PutLocCheck | PutLocCheckInit => match instr.imm {
            ImmOperand::Loc(k) => (Operand::Local(k as usize), false),
            _ => unreachable!(),
        },

        PutArg0 => (Operand::Arg(0), false),
        PutArg1 => (Operand::Arg(1), false),
        PutArg2 => (Operand::Arg(2), false),
        PutArg3 => (Operand::Arg(3), false),

        SetLoc0 => (Operand::Local(0), true),
        SetLoc1 => (Operand::Local(1), true),
        SetLoc2 => (Operand::Local(2), true),
        SetLoc3 => (Operand::Local(3), true),

        SetLoc | SetLoc8 => match instr.imm {
            ImmOperand::Loc(k) => (Operand::Local(k as usize), true),
            _ => unreachable!(),
        },

        SetArg0 => (Operand::Arg(0), true),
        SetArg1 => (Operand::Arg(1), true),
        SetArg2 => (Operand::Arg(2), true),
        SetArg3 => (Operand::Arg(3), true),

        _ => return false,
    };

    // Simulate stack pop or peek, depending on instr
    let arg = if preserved {
        stack.last().unwrap().clone()
    } else {
        stack.pop().unwrap()
    };

    // Converted SetLoc, reading from arg, writing to dst
    *out = RegInstr {
        dsts: vec![dst],
        op: SetLoc,
        imm: ImmOperand::None,
        args: vec![arg],
    };

    true
}

/// Handles an instr that only moves stack values around: it pops `n` values, and pushes some
/// of them again. No value is made, so `out_instr` becomes `nop`
fn try_handle_shuffle(instr: &Instr, stack: &mut Vec<Operand>, out_instr: &mut RegInstr) -> bool {
    let Some((n, pushes)) = shuffle(instr.op) else {
        return false;
    };

    let popped = stack.split_off(stack.len() - n);
    stack.extend(pushes.iter().map(|&i| popped[i]));

    *out_instr = RegInstr::nop();
    true
}

/// The shuffles: how many values an instr pops, and which of them it pushes, in stack order
/// (bottom first). The numbers index the popped values, bottom first
fn shuffle(op: OpCode) -> Option<(usize, &'static [usize])> {
    use OpCode::*;

    Some(match op {
        Drop => (1, &[]),                 // a ->
        Nip => (2, &[1]),                 // a b -> b
        Nip1 => (3, &[1, 2]),             // a b c -> b c
        Dup => (1, &[0, 0]),              // a -> a a
        Dup1 => (2, &[0, 0, 1]),          // a b -> a a b
        Dup2 => (2, &[0, 1, 0, 1]),       // a b -> a b a b
        Dup3 => (3, &[0, 1, 2, 0, 1, 2]), // a b c -> a b c a b c
        Insert2 => (2, &[1, 0, 1]),       // obj a -> a obj a
        Insert3 => (3, &[2, 0, 1, 2]),    // obj prop a -> a obj prop a
        Insert4 => (4, &[3, 0, 1, 2, 3]), // this obj prop a -> a this obj prop a
        Perm3 => (3, &[1, 0, 2]),         // obj a b -> a obj b
        Perm4 => (4, &[2, 0, 1, 3]),      // obj prop a b -> a obj prop b
        Perm5 => (5, &[3, 0, 1, 2, 4]),   // this obj prop a b -> a this obj prop b
        Swap => (2, &[1, 0]),             // a b -> b a
        Swap2 => (4, &[2, 3, 0, 1]),      // a b c d -> c d a b
        Rot3l => (3, &[1, 2, 0]),         // x a b -> a b x
        Rot3r => (3, &[2, 0, 1]),         // a b x -> x a b
        Rot4l => (4, &[1, 2, 3, 0]),      // x a b c -> a b c x
        Rot5l => (5, &[1, 2, 3, 4, 0]),   // x a b c d -> a b c d x
        // Binds [[HomeObject]], which only `super` reads. `super` is not supported
        SetHomeObject => (2, &[0, 1]),
        // Makes the key a property key. A key is kept as it is (see infer)
        ToPropkey2 => (2, &[0, 1]),
        _ => return None,
    })
}

/// Handle a normal instruction with simulated stack pushes/pops
fn handle_normal_case(
    instr: &Instr,
    stack: &mut Vec<Operand>,
    mut fresh_reg: impl FnMut() -> Operand,
    out_instr: &mut RegInstr,
) {
    use OpCode::*;

    // op and imm carried as-is
    *out_instr = RegInstr {
        dsts: vec![],
        op: instr.op,
        imm: instr.imm,
        args: vec![],
    };

    // Some instrs have extra pops not reflected in info.n_pop
    let extra_pops = match instr.imm {
        ImmOperand::Npop(npop) => npop as usize,
        ImmOperand::NpopU16 { npop, .. } => npop as usize,
        _ => 0,
    };

    let info = instr.op.info();
    let n_pop = info.n_pop as usize + extra_pops;
    let n_push = info.n_push as usize;

    // Simulate this instr's pops, top first
    for _ in 0..n_pop {
        let x = stack.pop().unwrap();
        out_instr.args.push(x);
    }

    if matches!(instr.op, Call | CallMethod | CallConstructor) {
        out_instr.args = call_operands(instr.op, std::mem::take(&mut out_instr.args));
    }

    // These leave the object they wrote on the stack: the same value, not a new one
    if matches!(instr.op, DefineField | DefineFieldRef) {
        stack.push(out_instr.args[1]);
        return;
    }

    // Simulate this instr's pushes
    for _ in 0..n_push {
        let x = fresh_reg();
        stack.push(x);
        out_instr.dsts.push(x);
    }
}

/// Put a call's popped operands in the order of the callee's params: the callee (which is its
/// env), `this`, then the args in source order
///
/// `call_method` has `this` at the bottom of its frame, below the callee. `call_constructor`
/// has the callee at the bottom and `this` above it (see transform/new.rs). `call` has `this`
/// on top, where transform pushed it
fn call_operands(op: OpCode, mut popped: Vec<Operand>) -> Vec<Operand> {
    let (callee, this) = match op {
        OpCode::CallMethod => {
            let this = popped.pop().unwrap();
            (popped.pop().unwrap(), this)
        }
        OpCode::CallConstructor => {
            let callee = popped.pop().unwrap();
            (callee, popped.pop().unwrap())
        }
        _ => {
            let callee = popped.pop().unwrap();
            (callee, popped.remove(0))
        }
    };

    // What remains are the args, last arg first
    popped.reverse();

    [callee, this].into_iter().chain(popped).collect()
}

/// `define_field_ref x` defines `x` as the location that its value was read from (see
/// transform/closure.rs), and the instr before it is that read. Its value arg becomes the
/// location:
/// - a local or arg: the field is that binding
/// - the reg of an object `o`, for a read `get_field x` of `o`: the field is `o.x`, which is
///   a location itself (a capture that the owner gives to its own child)
///
/// The read itself is then not used, so it becomes `nop`
fn resolve_location(prev: &RegInstr, out_instr: &mut RegInstr) {
    let value = out_instr.args[0];
    assert_eq!(
        prev.dsts,
        [value],
        "define_field_ref must follow the read of its value"
    );

    out_instr.args[0] = match prev.op {
        OpCode::SetLoc => prev.args[0],
        OpCode::GetField => {
            assert!(prev.imm.atom() == out_instr.imm.atom());
            prev.args[0]
        }
        op => panic!("define_field_ref after {op:?}, which reads no location"),
    };
}