neo-decompiler 0.10.1

use std::collections::BTreeMap;

use crate::decompiler::ir::Literal;
use crate::instruction::{Instruction, OpCode};

use super::super::effects;
use super::super::form::{SsaExpr, SsaStmt};
use super::super::variable::SsaVariable;
use super::lowering::{binary_op_for, literal_for_push, mnemonic, syscall_effect, unary_op_for};
use super::state::{fresh_var, is_unknown, reverse_top, slot_name_for, unknown_var, SlotState};
use super::SsaBuilder;

impl SsaBuilder<'_> {
    /// Apply a single instruction's stack effect or transformation.
    pub(super) fn apply_instruction(
        &self,
        instr: &Instruction,
        stack: &mut Vec<SsaVariable>,
        slots: &mut SlotState,
        stmts: &mut Vec<SsaStmt>,
        uses: &mut Vec<(SsaVariable, usize)>,
        versions: &mut BTreeMap<String, usize>,
    ) {
        let op = instr.opcode;

        if effects::is_stack_reorder(op) {
            self.apply_reorder(op, stack, stmts, versions);
            return;
        }
        if effects::is_stack_special(op) {
            self.apply_special(instr, stack, uses);
            return;
        }

        let (pop, push) = effects::stack_effect(op);
        let popped = pop_stack_values(stack, pop);
        let use_index = stmts.len();
        record_uses(&popped, uses, use_index);

        if push == 1 {
            self.push_result(instr, &popped, stack, slots, stmts, versions);
        } else if push == 0 {
            self.store_slot(instr, &popped, slots, stmts, versions);
        }
    }

    fn push_result(
        &self,
        instr: &Instruction,
        popped: &[SsaVariable],
        stack: &mut Vec<SsaVariable>,
        slots: &SlotState,
        stmts: &mut Vec<SsaStmt>,
        versions: &mut BTreeMap<String, usize>,
    ) {
        let op = instr.opcode;
        let reaching = slot_name_for(op, &instr.operand).and_then(|name| slots.get(&name).cloned());
        let expr = match reaching {
            Some(var) => SsaExpr::var(var),
            None => self.build_expr(op, instr, popped),
        };
        let base = slot_name_for(op, &instr.operand).unwrap_or_else(|| "t".to_string());
        let target = fresh_var(versions, &base);
        stmts.push(SsaStmt::assign(target.clone(), expr));
        stack.push(target);
    }

    fn store_slot(
        &self,
        instr: &Instruction,
        popped: &[SsaVariable],
        slots: &mut SlotState,
        stmts: &mut Vec<SsaStmt>,
        versions: &mut BTreeMap<String, usize>,
    ) {
        if let Some(name) = slot_name_for(instr.opcode, &instr.operand) {
            if let Some(value) = popped.first().cloned() {
                let target = fresh_var(versions, &name);
                stmts.push(SsaStmt::assign(target.clone(), SsaExpr::var(value)));
                slots.insert(name, target);
            }
        }
    }

    /// Handle fixed-shape stack reorders. New copies get fresh SSA definitions.
    fn apply_reorder(
        &self,
        op: OpCode,
        stack: &mut Vec<SsaVariable>,
        stmts: &mut Vec<SsaStmt>,
        versions: &mut BTreeMap<String, usize>,
    ) {
        let mut fresh_copy =
            |src: SsaVariable, stack: &mut Vec<SsaVariable>, stmts: &mut Vec<SsaStmt>| {
                let target = fresh_var(versions, "t");
                stmts.push(SsaStmt::assign(target.clone(), SsaExpr::var(src)));
                stack.push(target);
            };

        match op {
            OpCode::Dup => {
                if let Some(top) = stack.last().cloned() {
                    fresh_copy(top, stack, stmts);
                }
            }
            OpCode::Over => {
                if stack.len() >= 2 {
                    let second = stack[stack.len() - 2].clone();
                    fresh_copy(second, stack, stmts);
                }
            }
            OpCode::Tuck => {
                if stack.len() >= 2 {
                    let b = stack.pop().unwrap();
                    let a = stack.pop().unwrap();
                    fresh_copy(b.clone(), stack, stmts);
                    stack.push(a);
                    stack.push(b);
                }
            }
            OpCode::Swap => {
                let n = stack.len();
                if n >= 2 {
                    stack.swap(n - 1, n - 2);
                }
            }
            OpCode::Rot => {
                if stack.len() >= 3 {
                    let n = stack.len();
                    let a = stack.remove(n - 3);
                    stack.push(a);
                }
            }
            OpCode::Reverse3 => reverse_top(stack, 3),
            OpCode::Reverse4 => reverse_top(stack, 4),
            OpCode::Depth => {
                let depth = stack.len() as i64;
                let target = fresh_var(versions, "t");
                stmts.push(SsaStmt::assign(
                    target.clone(),
                    SsaExpr::lit(Literal::Int(depth)),
                ));
                stack.push(target);
            }
            OpCode::Drop => {
                stack.pop();
            }
            OpCode::Nip => {
                let n = stack.len();
                if n >= 2 {
                    stack.remove(n - 2);
                }
            }
            _ => {}
        }
    }

    /// Handle operand-dependent specials conservatively.
    fn apply_special(
        &self,
        instr: &Instruction,
        stack: &mut Vec<SsaVariable>,
        uses: &mut Vec<(SsaVariable, usize)>,
    ) {
        match instr.opcode {
            OpCode::Pick | OpCode::Roll | OpCode::Xdrop | OpCode::Reversen => {
                if let Some(v) = stack.pop() {
                    record_use_if_known(v, uses, 0);
                }
            }
            OpCode::Clear => {
                stack.clear();
            }
            OpCode::Syscall => {
                let (pop, push) = syscall_effect(instr);
                let popped = pop_stack_values(stack, pop);
                record_uses(&popped, uses, 0);
                if push {
                    stack.push(unknown_var());
                }
            }
            OpCode::Pack | OpCode::Packmap | OpCode::Packstruct => {
                let count = stack.pop().unwrap_or_else(unknown_var);
                record_use_if_known(count, uses, 0);
                stack.push(unknown_var());
            }
            OpCode::Unpack => {
                let _item = stack.pop();
                stack.push(unknown_var());
            }
            _ => {}
        }
    }

    /// Build the [`SsaExpr`] for a compute opcode from popped operands ordered
    /// deep-to-top.
    fn build_expr(&self, op: OpCode, instr: &Instruction, popped: &[SsaVariable]) -> SsaExpr {
        if let Some(lit) = literal_for_push(op, instr) {
            return SsaExpr::lit(lit);
        }

        if let Some(bin) = binary_op_for(op) {
            let mut it = popped.iter();
            let left = it.next().cloned().unwrap_or_else(unknown_var);
            let right = it.next().cloned().unwrap_or_else(unknown_var);
            return SsaExpr::binary(bin, SsaExpr::var(left), SsaExpr::var(right));
        }

        if matches!(
            op,
            OpCode::Within | OpCode::Substr | OpCode::Modmul | OpCode::Modpow
        ) {
            return call_expr(op, popped);
        }

        if let Some(un) = unary_op_for(op) {
            let operand = popped.first().cloned().unwrap_or_else(unknown_var);
            return SsaExpr::unary(un, SsaExpr::var(operand));
        }

        call_expr(op, popped)
    }
}

fn pop_stack_values(stack: &mut Vec<SsaVariable>, count: usize) -> Vec<SsaVariable> {
    let mut popped = Vec::with_capacity(count);
    for _ in 0..count {
        popped.push(stack.pop().unwrap_or_else(unknown_var));
    }
    popped.reverse();
    popped
}

fn record_uses(popped: &[SsaVariable], uses: &mut Vec<(SsaVariable, usize)>, index: usize) {
    for var in popped {
        if !is_unknown(var) {
            uses.push((var.clone(), index));
        }
    }
}

fn record_use_if_known(var: SsaVariable, uses: &mut Vec<(SsaVariable, usize)>, index: usize) {
    if !is_unknown(&var) {
        uses.push((var, index));
    }
}

fn call_expr(op: OpCode, popped: &[SsaVariable]) -> SsaExpr {
    SsaExpr::call(
        mnemonic(op),
        popped.iter().cloned().map(SsaExpr::var).collect(),
    )
}