rsleigh-decompile 0.5.0

use crate::fold::CallingConv;
use crate::ir::*;
use crate::provenance::{OperationOrigin, SourceScope};
use pcode_ir::{get_output, AddressSpaceId, PcodeOp, Varnode};
use std::collections::{HashMap, HashSet};

/// Convert a CFG into SSA form (SysV calling convention).
pub fn build_ssa(cfg: &Cfg) -> SsaCfg {
    build_ssa_with_cc(cfg, CallingConv::SysV)
}

/// Convert a CFG into SSA form with a specific calling convention.
/// The `cc` parameter controls which registers are invalidated after Call sites.
pub fn build_ssa_with_cc(cfg: &Cfg, cc: CallingConv) -> SsaCfg {
    build_ssa_with_return_register(cfg, cc, crate::fold::abi(cc).return_reg_int)
}

/// Use the native architecture's register layout for return-value selection.
/// MIPS/RISC-V argument binding remains unsupported; their return registers
/// must not be confused with an x86 register at a similar offset.
pub fn build_ssa_for_arch(cfg: &Cfg, cc: CallingConv, arch: rsleigh_api::Architecture) -> SsaCfg {
    let offset = match arch {
        rsleigh_api::Architecture::MIPS32 => 8,
        rsleigh_api::Architecture::RiscV64 => 0x2050,
        _ => crate::fold::abi(cc).return_reg_int.unwrap_or(0),
    };
    build_ssa_with_return_register(cfg, cc, Some(offset))
}

fn build_ssa_with_return_register(
    cfg: &Cfg,
    cc: CallingConv,
    return_register: Option<u64>,
) -> SsaCfg {
    let mut ssa = SsaCfg {
        blocks: Vec::new(),
        vars: Vec::new(),
        entry: cfg.entry,
        diagnostics: cfg.diagnostics.clone(),
    };

    let preds = cfg.predecessors();
    let classified_edges = cfg.classified_edges();
    let back_edges: HashSet<(BlockId, BlockId)> = classified_edges
        .iter()
        .filter(|edge| edge.kind == CfgEdgeKind::Back)
        .map(|edge| (edge.from, edge.to))
        .collect();
    let acyclic_edges: HashSet<(BlockId, BlockId)> = classified_edges
        .iter()
        .filter(|edge| edge.kind != CfgEdgeKind::Back)
        .map(|edge| (edge.from, edge.to))
        .collect();

    // Per-block: map from varnode -> VarId at block exit
    let mut block_exit_vars: Vec<HashMap<Varnode, VarId>> = vec![HashMap::new(); cfg.blocks.len()];

    // Iterative dataflow: re-process blocks until exit vars stabilize (max 4 passes).
    // Track the blocks changed by the previous complete pass. Comparing a
    // snapshot with the live map at the start of the next pass misses changes
    // made by predecessors that appear later in storage/layout order.
    let mut changed_blocks: HashSet<BlockId> = HashSet::new();
    for iteration in 0..4u32 {
        let changed_last_iteration = std::mem::take(&mut changed_blocks);

        for (block_idx, block) in cfg.blocks.iter().enumerate() {
            crate::budget::work("ssa", 1);
            let block_preds = &preds[block.id.0];

            // On iteration > 0, skip blocks whose predecessors haven't changed.
            // Also always skip the entry block — it has no predecessors, so its
            // register state (function parameters) should never be modified by
            // loop convergence iterations.
            if iteration > 0 {
                if block_preds.is_empty() {
                    continue; // Entry block — never re-process
                }
                let any_pred_changed = block_preds
                    .iter()
                    .any(|pred| changed_last_iteration.contains(pred));
                // Self-loop blocks (block is its own predecessor) must be re-processed
                // on iteration 1 so that early Phi nodes can be created for loop accumulators.
                // Without this, the skip condition prevents the block from ever seeing its
                // own back-edge exit vars.
                let has_back_edge = block_preds
                    .iter()
                    .any(|pred| back_edges.contains(&(*pred, block.id)));
                if !any_pred_changed && !(has_back_edge && iteration == 1) {
                    continue;
                }
            }

            let mut current: HashMap<Varnode, VarId> = HashMap::new();
            // Inherit from the first already-processed acyclic predecessor.
            // DFS edge classification, rather than block layout order, identifies
            // loop-carried inputs. Back-edge values are merged via Phi nodes below.
            if !block_preds.is_empty() {
                // Tree, forward, and cross edges all establish acyclic input
                // state. Only a DFS back edge is loop-carried.
                for pred in block_preds {
                    if acyclic_edges.contains(&(*pred, block.id))
                        && !block_exit_vars[pred.0].is_empty()
                    {
                        current = block_exit_vars[pred.0].clone();
                        break;
                    }
                }
                // Fallback: if no acyclic predecessor has data (entry or unreachable),
                // use any predecessor
                if current.is_empty() {
                    for pred in block_preds {
                        if !block_exit_vars[pred.0].is_empty() {
                            current = block_exit_vars[pred.0].clone();
                            break;
                        }
                    }
                }
            }

            let mut stmts = Vec::new();

            // Note: Phi nodes for loop-carried variables are created in the late Phi
            // pass (after all iterations) and then re-linked into loop body expressions.

            // Group P-code ops by instruction address for correct intra-instruction
            // register handling. x86-64 generates IntZext(EAX→RAX) before address
            // calculations that read RAX — the Zext must be deferred until after all
            // reads from the same instruction are resolved.
            let mut ops_iter = block.ops.iter().peekable();
            while ops_iter.peek().is_some() {
                // Collect all ops from the same instruction (same address)
                let inst_addr = ops_iter.peek().unwrap().0;
                let mut inst_ops: Vec<&PcodeOp> = Vec::new();
                while ops_iter.peek().map_or(false, |(a, _)| *a == inst_addr) {
                    inst_ops.push(&ops_iter.next().unwrap().1);
                }

                // Check for the sub-register Zext clobber pattern:
                // IntZext{out=(R,off,big), input=(R,off,small)} appears before
                // other ops that read (R,off,big).
                // If found, snapshot the pre-Zext value and defer the Zext write.
                let mut deferred_zext: Vec<(Varnode, VarId)> = Vec::new();

                // Find Zext ops that write to a register that is also read by later ops
                for (i, op) in inst_ops.iter().enumerate() {
                    let _source = SourceScope::enter(Some(OperationOrigin {
                        instruction_address: inst_addr,
                        operation_index: i,
                    }));
                    if let PcodeOp::IntZext { out, input } = op {
                        if out.space == AddressSpaceId::Register
                            && input.space == AddressSpaceId::Register
                            && out.offset == input.offset
                            && out.size > input.size
                        {
                            // Check if any later op in this instruction reads the output register
                            let reads_later = inst_ops[i + 1..]
                                .iter()
                                .any(|later_op| pcode_ir::reads_varnode(later_op, out));
                            if reads_later {
                                // Snapshot the current value of the super-register
                                // Process the Zext to get its VarId, but don't update current yet
                                let input_var = resolve_input(&mut ssa, &mut current, input);
                                let expr = Expr::UnaryOp(UnaryOpKind::Zext, input_var);
                                let var_id = ssa.new_var(*out, expr, out.size);
                                stmts.push(Stmt::Assign(var_id));
                                deferred_zext.push((*out, var_id));
                                continue;
                            }
                        }
                    }
                }

                // Detect MOVSD zero-clobber pattern:
                // Load { out: XMM(off>=4608, sz:16) } followed by
                // Copy { out: same_XMM, input: Const(0) }
                // The Copy zeros upper bytes — drop it to preserve the Load result.
                let mut skip_zero_copy: HashSet<usize> = HashSet::new();
                for (i, op) in inst_ops.iter().enumerate() {
                    if let PcodeOp::Load { out, .. } = op {
                        if out.space == AddressSpaceId::Register
                            && out.offset >= 4608
                            && out.size == 16
                        {
                            if i + 1 < inst_ops.len() {
                                if let PcodeOp::Copy {
                                    out: copy_out,
                                    input,
                                } = inst_ops[i + 1]
                                {
                                    if copy_out.space == out.space
                                        && copy_out.offset == out.offset
                                        && input.space == AddressSpaceId::Const
                                        && input.offset == 0
                                    {
                                        skip_zero_copy.insert(i + 1);
                                    }
                                }
                            }
                        }
                    }
                }

                // Detect intra-instruction CBranch (AArch64 CSEL/CSINC/CNEG pattern)
                // Pattern: [pre-ops..., CBranch{Const,cond}, else-ops..., post-op]
                // CBranch condition TRUE → skip else → use "then" value (from pre-ops)
                // CBranch condition FALSE → execute else → use "else" value
                let cbranch_idx = inst_ops.iter().position(|op| {
                    matches!(op, PcodeOp::CBranch { dest, .. } if dest.space == AddressSpaceId::Const)
                });

                // CSEL-style intra-instruction CBranch must have at least
                // one op AFTER the branch (else no else-path exists and
                // slicing `[cb_idx+1..last_idx]` panics). If the CBranch
                // is the last op of the instruction, fall through to the
                // regular per-op path below.
                let cbranch_idx = cbranch_idx.filter(|&i| i + 1 < inst_ops.len());

                if let Some(cb_idx) = cbranch_idx {
                    // Get the CBranch condition varnode
                    let cond_vn = if let PcodeOp::CBranch { cond, .. } = inst_ops[cb_idx] {
                        *cond
                    } else {
                        unreachable!()
                    };

                    // Process pre-CBranch ops normally (condition setup + then-value copies)
                    for (op_idx, op) in inst_ops[..cb_idx].iter().enumerate() {
                        if skip_zero_copy.contains(&op_idx) {
                            continue;
                        }
                        if let PcodeOp::IntZext { out, .. } = op {
                            if deferred_zext.iter().any(|(vn, _)| vn == out) {
                                continue;
                            }
                        }
                        process_op(
                            &mut ssa,
                            &mut current,
                            &mut stmts,
                            op,
                            cc,
                            OperationOrigin {
                                instruction_address: inst_addr,
                                operation_index: op_idx,
                            },
                        );
                    }

                    let cond_var = {
                        let _source = SourceScope::enter(Some(OperationOrigin {
                            instruction_address: inst_addr,
                            operation_index: cb_idx,
                        }));
                        resolve_input(&mut ssa, &mut current, &cond_vn)
                    };

                    // Snapshot current state — Unique varnodes hold "then" values
                    let then_state: HashMap<Varnode, VarId> = current
                        .iter()
                        .filter(|(vn, _)| vn.space == AddressSpaceId::Unique)
                        .map(|(vn, vid)| (*vn, *vid))
                        .collect();

                    // Process else-path ops (between CBranch and last op)
                    let last_idx = inst_ops.len() - 1;
                    for (relative_idx, op) in inst_ops[cb_idx + 1..last_idx].iter().enumerate() {
                        process_op(
                            &mut ssa,
                            &mut current,
                            &mut stmts,
                            op,
                            cc,
                            OperationOrigin {
                                instruction_address: inst_addr,
                                operation_index: cb_idx + 1 + relative_idx,
                            },
                        );
                    }

                    // For each Unique varnode written in both then and else paths,
                    // create a Ternary expression
                    for (vn, then_var) in &then_state {
                        if let Some(&else_var) = current.get(vn) {
                            if else_var != *then_var {
                                let ternary_expr = Expr::Ternary(cond_var, *then_var, else_var);
                                let ternary_id = ssa.new_var(*vn, ternary_expr, vn.size);
                                ssa.vars[ternary_id.0 as usize]
                                    .origins
                                    .insert(OperationOrigin {
                                        instruction_address: inst_addr,
                                        operation_index: cb_idx,
                                    });
                                current.insert(*vn, ternary_id);
                                stmts.push(Stmt::Assign(ternary_id));
                            }
                        }
                    }

                    // Process post-label ops (final assignment like IntZext)
                    if last_idx < inst_ops.len() {
                        process_op(
                            &mut ssa,
                            &mut current,
                            &mut stmts,
                            inst_ops[last_idx],
                            cc,
                            OperationOrigin {
                                instruction_address: inst_addr,
                                operation_index: last_idx,
                            },
                        );
                    }
                } else {
                    // Process remaining ops normally
                    for (op_idx, op) in inst_ops.iter().enumerate() {
                        // Skip MOVSD zero-clobber copies
                        if skip_zero_copy.contains(&op_idx) {
                            continue;
                        }
                        // Skip ops we already handled as deferred Zext
                        if let PcodeOp::IntZext { out, input: _ } = op {
                            if deferred_zext.iter().any(|(vn, _)| vn == out) {
                                continue;
                            }
                        }

                        process_op(
                            &mut ssa,
                            &mut current,
                            &mut stmts,
                            op,
                            cc,
                            OperationOrigin {
                                instruction_address: inst_addr,
                                operation_index: op_idx,
                            },
                        );
                    }
                }

                // Now apply deferred Zext writes
                for (vn, var_id) in deferred_zext {
                    current.insert(vn, var_id);
                }
            }

            let terminator = {
                let _source = SourceScope::enter(block.terminator_origin);
                convert_terminator(
                    &mut ssa,
                    &mut current,
                    &block.terminator,
                    cc,
                    &mut stmts,
                    return_register,
                )
            };

            if block_exit_vars[block.id.0] != current {
                changed_blocks.insert(block.id);
            }
            block_exit_vars[block.id.0] = current;

            // On first iteration, push new blocks; on subsequent iterations, replace
            if iteration == 0 {
                ssa.blocks.push(SsaBlock {
                    id: block.id,
                    addr: block.addr,
                    stmts,
                    terminator,
                });
            } else {
                ssa.blocks[block_idx].stmts = stmts;
                ssa.blocks[block_idx].terminator = terminator;
            }
        }

        if iteration > 0 && changed_blocks.is_empty() {
            break;
        }
    }

    // Second pass: insert Phi nodes at join points
    for bid in 0..cfg.blocks.len() {
        let block_preds = &preds[bid];
        if block_preds.len() <= 1 {
            continue;
        }

        // Find varnodes that differ across predecessors
        let mut all_varnodes: HashMap<Varnode, Vec<(BlockId, VarId)>> = HashMap::new();
        for &pred_id in block_preds {
            for (vn, &var_id) in &block_exit_vars[pred_id.0] {
                // Skip flag registers and tiny temporaries for cleaner output
                if vn.space == AddressSpaceId::Unique {
                    continue;
                }
                all_varnodes.entry(*vn).or_default().push((pred_id, var_id));
            }
        }

        // Sort varnodes deterministically so Phi creation order is stable.
        // HashMap iteration is non-deterministic, which cascades into VarId
        // assignment and downstream passes that depend on statement ordering.
        let mut sorted_vns: Vec<Varnode> = all_varnodes.keys().copied().collect();
        sorted_vns.sort_by_key(|vn| (vn.space, vn.offset, vn.size));

        let mut phi_stmts = Vec::new();
        for vn in &sorted_vns {
            let entries = &all_varnodes[vn];
            if entries.len() < 2 {
                continue;
            }
            // Check if all predecessors agree
            let first_var = entries[0].1;
            if entries.iter().all(|(_, v)| *v == first_var) {
                continue;
            }
            // Insert Phi
            let phi_inputs: Vec<VarId> = entries.iter().map(|(_, v)| *v).collect();
            let phi_var = ssa.new_var(*vn, Expr::Phi(phi_inputs.clone()), vn.size);
            phi_stmts.push(Stmt::Assign(phi_var));
        }

        // Prepend phis to block and re-link loop body expressions
        if !phi_stmts.is_empty() {
            // Build a replacement map: for each Phi, map the forward-predecessor's
            // VarId to the Phi VarId. This allows re-linking loop body expressions
            // so they read the Phi output instead of the stale pre-loop value.
            let mut relink: HashMap<VarId, VarId> = HashMap::new();
            for stmt in &phi_stmts {
                if let Stmt::Assign(phi_vid) = stmt {
                    if let Expr::Phi(_inputs) = &ssa.vars[phi_vid.0 as usize].expr {
                        // Relink values from every acyclic predecessor. Cross
                        // edges are ordinary join inputs, not loop-carried values.
                        let phi_vn = ssa.vars[phi_vid.0 as usize].varnode;
                        for &pred_id in block_preds {
                            if acyclic_edges.contains(&(pred_id, BlockId(bid))) {
                                if let Some(&fwd_var) = block_exit_vars[pred_id.0].get(&phi_vn) {
                                    relink.insert(fwd_var, *phi_vid);
                                }
                            }
                        }
                    }
                }
            }

            // Also build back-edge relink: map back-edge VarIds to Phi VarIds.
            // This ensures post-loop blocks reference the Phi (loop variable)
            // instead of the raw loop body result.
            let mut back_relink: HashMap<VarId, VarId> = HashMap::new();
            for stmt in &phi_stmts {
                if let Stmt::Assign(phi_vid) = stmt {
                    if let Expr::Phi(_inputs) = &ssa.vars[phi_vid.0 as usize].expr {
                        let phi_vn = ssa.vars[phi_vid.0 as usize].varnode;
                        for &pred_id in block_preds {
                            if back_edges.contains(&(pred_id, BlockId(bid))) {
                                if let Some(&back_var) = block_exit_vars[pred_id.0].get(&phi_vn) {
                                    back_relink.insert(back_var, *phi_vid);
                                }
                            }
                        }
                    }
                }
            }

            // Re-link: replace stale forward-pred references with Phi VarIds
            // in all expressions within this block.
            if !relink.is_empty() {
                let block = &mut ssa.blocks[bid];
                for stmt in &block.stmts {
                    if let Stmt::Assign(vid) = stmt {
                        let vi = vid.0 as usize;
                        let rewritten = relink_expr(&ssa.vars[vi].expr, &relink);
                        crate::provenance::rewrite(&mut ssa.vars, vi, rewritten);
                    }
                }
                // Also re-link the terminator condition
                if let SsaTerminator::CBranch {
                    cond,
                    taken,
                    fallthrough,
                } = &block.terminator
                {
                    if let Some(&new_cond) = relink.get(cond) {
                        let t = *taken;
                        let f = *fallthrough;
                        ssa.blocks[bid].terminator = SsaTerminator::CBranch {
                            cond: new_cond,
                            taken: t,
                            fallthrough: f,
                        };
                    }
                }
            }

            // Re-link successor blocks: replace back-edge VarIds with Phi VarIds.
            // This ensures post-loop returns reference the Phi (the loop variable)
            // instead of the raw ADD result from the last iteration.
            if !back_relink.is_empty() {
                // Find successor blocks (exit targets from this loop header)
                let successors: Vec<usize> = match &ssa.blocks[bid].terminator {
                    SsaTerminator::CBranch {
                        taken, fallthrough, ..
                    } => {
                        let mut s = Vec::new();
                        if taken.0 != bid {
                            s.push(taken.0);
                        }
                        if fallthrough.0 != bid {
                            s.push(fallthrough.0);
                        }
                        s
                    }
                    SsaTerminator::Fallthrough(b) | SsaTerminator::Branch(b) => {
                        if b.0 != bid {
                            vec![b.0]
                        } else {
                            vec![]
                        }
                    }
                    _ => vec![],
                };
                for succ_bid in successors {
                    if succ_bid >= ssa.blocks.len() {
                        continue;
                    }
                    for stmt in &ssa.blocks[succ_bid].stmts {
                        if let Stmt::Assign(vid) = stmt {
                            let vi = vid.0 as usize;
                            let rewritten = relink_expr(&ssa.vars[vi].expr, &back_relink);
                            crate::provenance::rewrite(&mut ssa.vars, vi, rewritten);
                        }
                    }
                    // Re-link return value
                    if let SsaTerminator::Return(Some(ret_var)) = &ssa.blocks[succ_bid].terminator {
                        if let Some(&phi_var) = back_relink.get(ret_var) {
                            ssa.blocks[succ_bid].terminator = SsaTerminator::Return(Some(phi_var));
                        } else {
                            // Also check: the return might reference a Var/Zext chain
                            // that wraps a back-edge VarId. Follow one level.
                            let rv = &ssa.vars[ret_var.0 as usize];
                            let inner = match &rv.expr {
                                Expr::Var(v) => Some(*v),
                                Expr::UnaryOp(UnaryOpKind::Zext, v) => Some(*v),
                                _ => None,
                            };
                            if let Some(inner_id) = inner {
                                if let Some(&phi_var) = back_relink.get(&inner_id) {
                                    ssa.blocks[succ_bid].terminator =
                                        SsaTerminator::Return(Some(phi_var));
                                }
                            }
                        }
                    }
                }
            }

            let block = &mut ssa.blocks[bid];
            let mut new_stmts = phi_stmts;
            new_stmts.append(&mut block.stmts);
            block.stmts = new_stmts;
        }
    }

    for block in &ssa.blocks {
        if matches!(block.terminator, SsaTerminator::Return(Some(id)) if ssa.var(id).call_return) {
            ssa.diagnostics.push(Diagnostic {
                severity: Severity::Info,
                kind: DiagKind::StaleReturnInherited,
                addr: Some(block.addr),
                detail:
                    "return-register value comes from a call; source-level function may be void"
                        .into(),
            });
        }
    }
    if return_register == crate::fold::abi(cc).return_reg_int {
        crate::memory::forward(&mut ssa, cfg, cc);
    }

    crate::provenance::propagate(&mut ssa.vars);

    // Count uses (after Phase 2 + Phase 3 may have changed expressions)
    count_uses(&mut ssa);

    ssa
}

/// Process a single P-code op: resolve inputs, build SSA expression, update current map.
/// Extracted to avoid duplication between normal path and CSEL path.
fn process_op(
    ssa: &mut SsaCfg,
    current: &mut HashMap<Varnode, VarId>,
    stmts: &mut Vec<Stmt>,
    op: &PcodeOp,
    _cc: CallingConv,
    origin: OperationOrigin,
) {
    let _source = SourceScope::enter(Some(origin));
    match op.clone() {
        PcodeOp::Store { space, ptr, val } => {
            let addr_var = resolve_input(ssa, current, &ptr);
            let val_var = resolve_input(ssa, current, &val);
            // Give the stored value its own SSA definition so copies and folds
            // retain the Store operation without attributing later writes to
            // the original value's earlier uses.
            let stored = ssa.new_var(
                Varnode::unique(0xF300_0000 + ssa.vars.len() as u64, val.size),
                Expr::Var(val_var),
                val.size,
            );
            let address_origins = ssa.var(addr_var).origins.clone();
            ssa.var_mut(stored).origins.merge(&address_origins);
            ssa.var_mut(stored).origins.synthetic = true;
            ssa.var_mut(stored).memory = Some(crate::memory::Access::Store { space });
            stmts.push(Stmt::Store {
                addr: addr_var,
                val: stored,
            });
        }
        PcodeOp::CallOther {
            func_id,
            inputs,
            out: None,
        } => {
            // Void user-pcodeop (e.g. `software_interrupt(0x71)` on ARM swi).
            // Emit as a statement even though there's no output varnode — the
            // side effect itself is meaningful (it changes machine state the
            // decompiler cannot model, so surfacing the call keeps the
            // analyst informed).
            let resolved: Vec<VarId> = inputs
                .iter()
                .map(|vn| resolve_input(ssa, current, vn))
                .collect();
            // Allocate a synthetic var to hold the UserOp expr so the printer
            // can process it through the usual Stmt::Assign path.
            let placeholder_vn = Varnode {
                space: AddressSpaceId::Unique,
                offset: u64::MAX - func_id,
                size: 0,
            };
            let var_id = ssa.new_var(
                placeholder_vn,
                Expr::UserOp {
                    func_id,
                    inputs: resolved,
                },
                0,
            );
            stmts.push(Stmt::Assign(var_id));
        }
        ref op => {
            if let Some(out_vn) = get_output(op) {
                let expr = build_expr(ssa, current, op);
                let effective_size = float_semantic_size(&expr, &ssa.vars).unwrap_or(out_vn.size);
                let var_id = ssa.new_var(out_vn, expr, effective_size);
                if let PcodeOp::Load { space, .. } = op {
                    ssa.var_mut(var_id).memory = Some(crate::memory::Access::Load {
                        space: *space,
                        stores: vec![],
                        boundary: Some(crate::memory::Boundary::UnmodeledMemory),
                    });
                }
                // An unsupported output still has a defining raw operation.
                if matches!(ssa.vars[var_id.0 as usize].expr, Expr::Unknown) {
                    ssa.vars[var_id.0 as usize].origins.attach_source();
                }
                current.insert(out_vn, var_id);
                // Sub-register propagation: when writing to a larger register (e.g., RAX 8-byte),
                // also update the smaller sub-register at the same offset (e.g., EAX 4-byte).
                // This ensures that return value detection finds the correct value when the
                // function uses 64-bit ops (LEA/INC on RAX) but the return checks EAX first.
                if out_vn.space == AddressSpaceId::Register && out_vn.size == 8 {
                    let sub_vn = Varnode {
                        space: out_vn.space,
                        offset: out_vn.offset,
                        size: 4,
                    };
                    current.insert(sub_vn, var_id);
                }
                // Reverse sub-register propagation: writing to a SMALLER register
                // (e.g. AL = 1 byte from `setne al`) must also update the
                // parent aliases (EAX 4, RAX 8) so subsequent reads of the parent
                // see the merged value instead of a stale pre-write value.
                //
                // The merge model is `parent = (old & ~mask) | Zext(new)` where
                // mask covers the low `out_vn.size` bytes. This preserves the
                // high bytes of the prior parent value, which a pure Zext would
                // discard. Canonical case where the high bytes matter:
                //   mov eax, 0x12345678   (RAX low 4 bytes = 0x12345678)
                //   mov al,  0x01         (RAX low byte = 0x01 → 0x12345601)
                // Constant folding collapses the masked-or back to a single
                // Const when both inputs are constants, so the bool-return
                // idiom (xor eax,eax; setne al) still folds to a clean Zext.
                if out_vn.space == AddressSpaceId::Register
                    && (out_vn.size == 1 || out_vn.size == 2)
                {
                    for parent_size in [4u32, 8u32] {
                        if parent_size <= out_vn.size {
                            continue;
                        }
                        let parent_vn = Varnode {
                            space: out_vn.space,
                            offset: out_vn.offset,
                            size: parent_size,
                        };
                        if let Some(&parent_old) = current.get(&parent_vn) {
                            // mask: high bytes of parent (above the sub-write).
                            let high_mask: u64 = if parent_size >= 8 {
                                !((1u128 << (out_vn.size as u128 * 8)) - 1) as u64
                            } else {
                                let parent_bits = parent_size * 8;
                                let sub_bits = out_vn.size * 8;
                                let parent_mask = if parent_bits >= 64 {
                                    u64::MAX
                                } else {
                                    (1u64 << parent_bits) - 1
                                };
                                let sub_mask = (1u64 << sub_bits) - 1;
                                parent_mask & !sub_mask
                            };
                            let mask_const_id = ssa.new_var(
                                Varnode {
                                    space: AddressSpaceId::Const,
                                    offset: high_mask,
                                    size: parent_size,
                                },
                                Expr::Const(high_mask, parent_size),
                                parent_size,
                            );
                            let high_id = ssa.new_var(
                                parent_vn,
                                Expr::BinOp(BinOpKind::And, parent_old, mask_const_id),
                                parent_size,
                            );
                            let zext_id = ssa.new_var(
                                parent_vn,
                                Expr::UnaryOp(UnaryOpKind::Zext, var_id),
                                parent_size,
                            );
                            let merged_id = ssa.new_var(
                                parent_vn,
                                Expr::BinOp(BinOpKind::Or, high_id, zext_id),
                                parent_size,
                            );
                            for id in [mask_const_id, high_id, zext_id, merged_id] {
                                ssa.vars[id.0 as usize].origins.synthetic = true;
                            }
                            current.insert(parent_vn, merged_id);
                        }
                    }
                }
                stmts.push(Stmt::Assign(var_id));
            }
        }
    }
}

fn resolve_input(ssa: &mut SsaCfg, current: &mut HashMap<Varnode, VarId>, vn: &Varnode) -> VarId {
    if vn.space == AddressSpaceId::Const {
        return ssa.new_var(*vn, Expr::Const(vn.offset, vn.size), vn.size);
    }
    if let Some(&var_id) = current.get(vn) {
        return var_id;
    }
    // Sub-register aliasing at the same offset:
    // Case 1: Reading smaller (w8) when larger (x8) was written → reuse directly
    //   Common on AArch64 where CSETM writes x8 and CSINC reads w8.
    // Case 2: Reading larger (RDX) when smaller (EDX) was written → zero-extend
    //   Common on x86-64 where 32-bit ops implicitly zero-extend to 64-bit.
    if vn.space == AddressSpaceId::Register {
        for (&existing_vn, &existing_var) in current.iter() {
            if existing_vn.space == AddressSpaceId::Register
                && existing_vn.offset == vn.offset
                && existing_vn.size != vn.size
            {
                if existing_vn.size > vn.size {
                    // Case 1: read smaller from larger — reuse directly
                    return existing_var;
                } else {
                    // Case 2: read larger from smaller — zero-extend
                    let expr = Expr::UnaryOp(UnaryOpKind::Zext, existing_var);
                    let var_id = ssa.new_var(*vn, expr, vn.size);
                    ssa.vars[var_id.0 as usize].origins.synthetic = true;
                    current.insert(*vn, var_id);
                    return var_id;
                }
            }
        }
    }
    // Unknown — function parameter or uninitialized
    let var_id = ssa.new_var(*vn, Expr::Unknown, vn.size);
    current.insert(*vn, var_id);
    var_id
}

/// Replace VarId references in an expression according to a replacement map.
/// Used to re-link loop body expressions to read from Phi nodes instead of
/// stale pre-loop values.
fn relink_expr(expr: &Expr, relink: &HashMap<VarId, VarId>) -> Expr {
    match expr {
        Expr::Var(id) => Expr::Var(*relink.get(id).unwrap_or(id)),
        Expr::BinOp(k, l, r) => {
            Expr::BinOp(*k, *relink.get(l).unwrap_or(l), *relink.get(r).unwrap_or(r))
        }
        Expr::UnaryOp(k, i) => Expr::UnaryOp(*k, *relink.get(i).unwrap_or(i)),
        Expr::Load(p) => Expr::Load(*relink.get(p).unwrap_or(p)),
        Expr::Ternary(c, t, e) => Expr::Ternary(
            *relink.get(c).unwrap_or(c),
            *relink.get(t).unwrap_or(t),
            *relink.get(e).unwrap_or(e),
        ),
        Expr::Phi(inputs) => {
            Expr::Phi(inputs.iter().map(|i| *relink.get(i).unwrap_or(i)).collect())
        }
        Expr::UserOp { func_id, inputs } => Expr::UserOp {
            func_id: *func_id,
            inputs: inputs.iter().map(|i| *relink.get(i).unwrap_or(i)).collect(),
        },
        _ => expr.clone(),
    }
}

/// Caller-saved (volatile) integer register offsets per ABI.
/// These registers must be invalidated in the SSA `current` map after any Call.
///
/// x86-64 offsets: RAX=0, RCX=8, RDX=16, RSI=48, RDI=56, R8=128, R9=136, R10=144, R11=152
/// AArch64: x0=16384 stride 8, x0..x18 are caller-saved
/// ARM32/x86-32: r0/EAX=0, r1/ECX=8, r2/EDX=16, r3=44(ARM) or nothing extra
/// MIPS/RISC-V: covered by SysV default as fallback
const WIN64_CALLER_SAVED: &[u64] = &[
    0,   // RAX
    8,   // RCX
    16,  // RDX
    128, // R8
    136, // R9
    144, // R10
    152, // R11
];

const SYSV64_CALLER_SAVED: &[u64] = &[
    0,   // RAX
    8,   // RCX
    16,  // RDX
    48,  // RSI
    56,  // RDI
    128, // R8
    136, // R9
    144, // R10
    152, // R11
];

/// AArch64 AAPCS64 caller-saved: x0..x18 at stride 8 starting at 16384.
const AARCH64_CALLER_SAVED: &[u64] = &[
    16384, 16392, 16400, 16408, 16416, 16424, 16432, 16440, // x0..x7
    16448, 16456, 16464, 16472, 16480, 16488, 16496, 16504, // x8..x15
    16512, 16520, 16528, // x16..x18
];

/// x86-32 cdecl caller-saved: EAX, ECX, EDX in the native 32-bit register space.
const X86_32_CALLER_SAVED: &[u64] = &[0, 4, 8];

/// ARM32 AAPCS caller-saved: r0-r3 (args), r12 (IP scratch), r14 (LR).
const ARM32_CALLER_SAVED: &[u64] = &[
    32, 36, 40, 44, // r0..r3
    80, // r12 (offset 0x20 + 12*4 = 0x50 = 80)
    88, // r14 / lr (0x20 + 14*4 = 0x58 = 88)
];

/// Register offset of the return register per calling convention.
fn return_reg_offset(cc: CallingConv) -> u64 {
    crate::fold::abi(cc).return_reg_int.unwrap_or(0)
}

/// Size in bytes of the return register per calling convention.
fn return_reg_size(cc: CallingConv) -> u32 {
    match cc {
        CallingConv::SysV | CallingConv::Win64 | CallingConv::GoAmd64 | CallingConv::AArch64 => 8,
        CallingConv::Cdecl32
        | CallingConv::Stdcall32
        | CallingConv::Thiscall32
        | CallingConv::Fastcall32
        | CallingConv::Arm32 => 4,
    }
}

fn caller_saved_offsets(cc: CallingConv) -> &'static [u64] {
    match cc {
        CallingConv::Win64 => WIN64_CALLER_SAVED,
        CallingConv::SysV | CallingConv::GoAmd64 => SYSV64_CALLER_SAVED,
        CallingConv::AArch64 => AARCH64_CALLER_SAVED,
        CallingConv::Cdecl32
        | CallingConv::Stdcall32
        | CallingConv::Thiscall32
        | CallingConv::Fastcall32 => X86_32_CALLER_SAVED,
        CallingConv::Arm32 => ARM32_CALLER_SAVED,
    }
}

/// Invalidate caller-saved registers in `current` after a Call.
/// Emits one `Stmt::Assign(ret_var)` for the return register with `call_return=true`.
/// Other caller-saved registers receive unknown definitions attributed to the
/// call, so later reads cannot be mislabeled as incoming function parameters.
fn clobber_caller_saved(
    ssa: &mut SsaCfg,
    current: &mut HashMap<Varnode, VarId>,
    cc: CallingConv,
    stmts: &mut Vec<Stmt>,
) {
    let offsets = caller_saved_offsets(cc);
    let ret_off = return_reg_offset(cc);
    let ret_size = return_reg_size(cc);

    // Drop every current entry at any caller-saved offset, regardless of size.
    current.retain(|vn, _| !(vn.space == AddressSpaceId::Register && offsets.contains(&vn.offset)));

    for &offset in offsets {
        if offset == ret_off {
            continue;
        }
        let vn = Varnode::register(offset, ret_size);
        let unknown = ssa.new_var(vn, Expr::Unknown, ret_size);
        ssa.var_mut(unknown).origins.attach_source();
        current.insert(vn, unknown);
        if ret_size == 8 {
            current.insert(Varnode::register(offset, 4), unknown);
        }
    }

    // Create a fresh return-register clobber with call_return=true.
    let ret_vn = Varnode {
        space: AddressSpaceId::Register,
        offset: ret_off,
        size: ret_size,
    };
    let ret_var = ssa.new_var(ret_vn, Expr::Unknown, ret_size);
    ssa.vars[ret_var.0 as usize].call_return = true;
    ssa.vars[ret_var.0 as usize].origins.attach_source();
    current.insert(ret_vn, ret_var);

    // Seed size-4 sub-register too (so `mov eax, ...` reads see the same VarId).
    if ret_size == 8 {
        let sub_vn = Varnode {
            space: AddressSpaceId::Register,
            offset: ret_off,
            size: 4,
        };
        current.insert(sub_vn, ret_var);
    }

    stmts.push(Stmt::Assign(ret_var));
}

fn build_expr(ssa: &mut SsaCfg, current: &mut HashMap<Varnode, VarId>, op: &PcodeOp) -> Expr {
    macro_rules! bin {
        ($kind:ident, $left:expr, $right:expr) => {{
            let l = resolve_input(ssa, current, $left);
            let r = resolve_input(ssa, current, $right);
            Expr::BinOp(BinOpKind::$kind, l, r)
        }};
    }
    macro_rules! unary {
        ($kind:ident, $input:expr) => {{
            let i = resolve_input(ssa, current, $input);
            Expr::UnaryOp(UnaryOpKind::$kind, i)
        }};
    }

    match op {
        PcodeOp::Copy { input, .. } => {
            let v = resolve_input(ssa, current, input);
            Expr::Var(v)
        }
        PcodeOp::Load { ptr, .. } => {
            let p = resolve_input(ssa, current, ptr);
            Expr::Load(p)
        }
        PcodeOp::IntAdd { left, right, .. } => bin!(Add, left, right),
        PcodeOp::IntSub { left, right, .. } => bin!(Sub, left, right),
        PcodeOp::IntMult { left, right, .. } => bin!(Mult, left, right),
        PcodeOp::IntDiv { left, right, .. } => bin!(Div, left, right),
        PcodeOp::IntSDiv { left, right, .. } => bin!(SDiv, left, right),
        PcodeOp::IntRem { left, right, .. } => bin!(Rem, left, right),
        PcodeOp::IntSRem { left, right, .. } => bin!(SRem, left, right),
        PcodeOp::IntAnd { left, right, .. } => bin!(And, left, right),
        PcodeOp::IntOr { left, right, .. } => bin!(Or, left, right),
        PcodeOp::IntXor {
            left, right, out, ..
        } => {
            // XOR reg, reg → 0 (common zero-init: XORPS/XORPD/XOR EAX,EAX)
            if left.space == right.space
                && left.offset == right.offset
                && left.size == right.size
                && left.space == AddressSpaceId::Register
            {
                Expr::Const(0, out.size)
            } else {
                bin!(Xor, left, right)
            }
        }
        PcodeOp::IntLsl { left, right, .. } => bin!(Lsl, left, right),
        PcodeOp::IntLsr { left, right, .. } => bin!(Lsr, left, right),
        PcodeOp::IntAsr { left, right, .. } => bin!(Asr, left, right),
        PcodeOp::IntEq { left, right, .. } => bin!(Eq, left, right),
        PcodeOp::IntNotEq { left, right, .. } => bin!(NotEq, left, right),
        PcodeOp::IntLess { left, right, .. } => bin!(Less, left, right),
        PcodeOp::IntLessEq { left, right, .. } => bin!(LessEq, left, right),
        PcodeOp::IntSLess { left, right, .. } => bin!(SLess, left, right),
        PcodeOp::IntSLessEq { left, right, .. } => bin!(SLessEq, left, right),
        PcodeOp::IntCarry { left, right, .. } => bin!(Carry, left, right),
        PcodeOp::IntSCarry { left, right, .. } => bin!(SCarry, left, right),
        PcodeOp::IntSBorrow { left, right, .. } => bin!(SBorrow, left, right),
        PcodeOp::IntNeg { input, .. } => unary!(Neg, input),
        PcodeOp::IntNot { input, .. } => unary!(Not, input),
        PcodeOp::IntZext { input, .. } => unary!(Zext, input),
        PcodeOp::IntSext { input, .. } => unary!(Sext, input),
        PcodeOp::BoolAnd { left, right, .. } => bin!(BoolAnd, left, right),
        PcodeOp::BoolOr { left, right, .. } => bin!(BoolOr, left, right),
        PcodeOp::BoolXor { left, right, .. } => bin!(BoolXor, left, right),
        PcodeOp::BoolNot { input, .. } => unary!(BoolNot, input),
        PcodeOp::FloatAdd { left, right, .. } => bin!(FloatAdd, left, right),
        PcodeOp::FloatSub { left, right, .. } => bin!(FloatSub, left, right),
        PcodeOp::FloatMult { left, right, .. } => bin!(FloatMult, left, right),
        PcodeOp::FloatDiv { left, right, .. } => bin!(FloatDiv, left, right),
        PcodeOp::FloatEq { left, right, .. } => bin!(FloatEq, left, right),
        PcodeOp::FloatNotEq { left, right, .. } => bin!(FloatNotEq, left, right),
        PcodeOp::FloatLess { left, right, .. } => bin!(FloatLess, left, right),
        PcodeOp::FloatLessEq { left, right, .. } => bin!(FloatLessEq, left, right),
        PcodeOp::FloatNeg { input, .. } => unary!(FloatNeg, input),
        PcodeOp::FloatAbs { input, .. } => unary!(FloatAbs, input),
        PcodeOp::FloatSqrt { input, .. } => unary!(FloatSqrt, input),
        PcodeOp::FloatNan { input, .. } => unary!(FloatNan, input),
        PcodeOp::Int2Float { input, .. } => unary!(Int2Float, input),
        PcodeOp::Float2Float { input, .. } => unary!(Float2Float, input),
        PcodeOp::Trunc { input, .. } => unary!(Trunc, input),
        PcodeOp::FloatCeil { input, .. } => unary!(FloatCeil, input),
        PcodeOp::FloatFloor { input, .. } => unary!(FloatFloor, input),
        PcodeOp::FloatRound { input, .. } => unary!(FloatRound, input),
        PcodeOp::Popcount { input, .. } => unary!(Popcount, input),
        PcodeOp::Lzcount { input, .. } => unary!(Lzcount, input),
        PcodeOp::CallOther {
            func_id, inputs, ..
        } => {
            let resolved: Vec<VarId> = inputs
                .iter()
                .map(|vn| resolve_input(ssa, current, vn))
                .collect();
            Expr::UserOp {
                func_id: *func_id,
                inputs: resolved,
            }
        }
        PcodeOp::Subpiece { input, lsb, out: _ } => {
            let i = resolve_input(ssa, current, input);
            if *lsb == 0 {
                // Truncation — just treat as a variable reference
                Expr::Var(i)
            } else {
                let shift_amt = ssa.new_var(
                    Varnode::constant((*lsb as u64) * 8, 4),
                    Expr::Const((*lsb as u64) * 8, 4),
                    4,
                );
                ssa.vars[shift_amt.0 as usize].origins.synthetic = true;
                Expr::BinOp(BinOpKind::Lsr, i, shift_amt)
            }
        }
        // Branching ops (Branch/CBranch/BranchInd/Call/CallInd/Return) are
        // consumed by the CFG builder before SSA, and Store has no Expr value,
        // so this fallthrough only fires when a new PcodeOp variant is added
        // without lowering. Surface it.
        other => {
            ssa.diagnostics.push(Diagnostic {
                severity: Severity::Warn,
                kind: DiagKind::UnknownPcodeOp,
                addr: None,
                detail: format!("build_expr: no lowering for {:?}", other),
            });
            Expr::Unknown
        }
    }
}

/// For float ops, return the semantic operand size (4=float, 8=double).
/// SSE scalar instructions write to full 16-byte XMM registers but the
/// meaningful result is only the low 4 or 8 bytes.
fn float_semantic_size(expr: &Expr, vars: &[VarDef]) -> Option<u32> {
    match expr {
        Expr::BinOp(kind, left, right) => {
            use BinOpKind::*;
            match kind {
                FloatAdd | FloatSub | FloatMult | FloatDiv => {
                    let ls = vars[left.0 as usize].size;
                    let rs = vars[right.0 as usize].size;
                    Some(ls.min(rs))
                }
                _ => None,
            }
        }
        Expr::UnaryOp(kind, input) => {
            use UnaryOpKind::*;
            match kind {
                FloatNeg | FloatAbs | FloatSqrt | FloatCeil | FloatFloor | FloatRound => {
                    Some(vars[input.0 as usize].size)
                }
                Int2Float => {
                    let is = vars[input.0 as usize].size;
                    Some(if is >= 8 { 8 } else { 4 })
                }
                Float2Float => None,
                _ => None,
            }
        }
        _ => None,
    }
}

fn convert_terminator(
    ssa: &mut SsaCfg,
    current: &mut HashMap<Varnode, VarId>,
    term: &Terminator,
    cc: CallingConv,
    stmts: &mut Vec<Stmt>,
    return_register: Option<u64>,
) -> SsaTerminator {
    match term {
        Terminator::Fallthrough(b) => SsaTerminator::Fallthrough(*b),
        Terminator::Branch(b) => SsaTerminator::Branch(*b),
        Terminator::CBranch {
            cond,
            taken,
            fallthrough,
        } => {
            let cond_var = resolve_input(ssa, current, cond);
            SsaTerminator::CBranch {
                cond: cond_var,
                taken: *taken,
                fallthrough: *fallthrough,
            }
        }
        Terminator::Call {
            target,
            fallthrough,
        } => {
            if return_register == crate::fold::abi(cc).return_reg_int {
                clobber_caller_saved(ssa, current, cc, stmts);
            } else {
                // The fallback convention does not describe this architecture.
                // Invalidate all live registers conservatively; only the native
                // return register gets a call-result dependency. Do not retain
                // a pre-call value using an unrelated architecture's ABI.
                let mut registers: Vec<_> = current
                    .keys()
                    .copied()
                    .filter(|vn| vn.space == AddressSpaceId::Register)
                    .collect();
                registers.sort_by_key(|vn| (vn.offset, vn.size));
                for vn in registers {
                    let value = ssa.new_var(vn, Expr::Unknown, vn.size);
                    ssa.var_mut(value).origins.attach_source();
                    current.insert(vn, value);
                }
                if let Some(offset) = return_register {
                    let size = if offset == 8 { 4 } else { 8 };
                    let vn = Varnode::register(offset, size);
                    let result = ssa.new_var(vn, Expr::Unknown, size);
                    ssa.var_mut(result).call_return = true;
                    ssa.var_mut(result).origins.attach_source();
                    current.insert(vn, result);
                    stmts.push(Stmt::Assign(result));
                }
            }
            SsaTerminator::Call {
                target: target.clone(),
                args: vec![],
                out: None,
                fallthrough: *fallthrough,
            }
        }
        Terminator::Return => {
            // Select only the native return register, preferring its narrow
            // definition when present. Entry unknowns do not imply a return.
            let ret_val = return_register
                .into_iter()
                .flat_map(|offset| [Varnode::register(offset, 4), Varnode::register(offset, 8)])
                .find_map(|vn| {
                    let var_id = current.get(&vn).copied()?;
                    let vdef = &ssa.vars[var_id.0 as usize];
                    // Skip if this is just the entry parameter value (Unknown)
                    // — the function didn't explicitly set a return value.
                    // Also skip bare Unknown without param_name (uninitialized reads).
                    if matches!(&vdef.expr, Expr::Unknown) && !vdef.call_return {
                        return None;
                    }
                    Some(var_id)
                });
            SsaTerminator::Return(ret_val)
        }
        Terminator::Indirect(vn) => {
            let v = resolve_input(ssa, current, vn);
            SsaTerminator::Indirect(v)
        }
    }
}

fn count_uses(ssa: &mut SsaCfg) {
    // Collect all referenced VarIds first, then update counts
    let mut use_counts = vec![0u32; ssa.vars.len()];

    for v in 0..ssa.vars.len() {
        crate::budget::work("ssa", 1);
        let refs = collect_expr_refs(&ssa.vars[v].expr);
        for id in refs {
            use_counts[id.0 as usize] += 1;
        }
    }

    for block in &ssa.blocks {
        crate::budget::work("ssa", 1);
        for stmt in &block.stmts {
            match stmt {
                Stmt::Store { addr, val } => {
                    use_counts[addr.0 as usize] += 1;
                    use_counts[val.0 as usize] += 1;
                }
                Stmt::Call { args, out: _, .. } => {
                    for a in args {
                        use_counts[a.0 as usize] += 1;
                    }
                }
                _ => {}
            }
        }
        match &block.terminator {
            SsaTerminator::CBranch { cond, .. } => {
                use_counts[cond.0 as usize] += 1;
            }
            SsaTerminator::Return(Some(v)) | SsaTerminator::Indirect(v) => {
                use_counts[v.0 as usize] += 1;
            }
            _ => {}
        }
    }

    for (i, count) in use_counts.into_iter().enumerate() {
        ssa.vars[i].use_count = count;
    }
}

pub(crate) fn collect_expr_refs(expr: &Expr) -> Vec<VarId> {
    match expr {
        Expr::Var(id) => vec![*id],
        Expr::BinOp(_, l, r) => vec![*l, *r],
        Expr::UnaryOp(_, i) | Expr::Load(i) | Expr::FieldAccess(i, _) => vec![*i],
        Expr::Phi(inputs) => inputs.clone(),
        Expr::Ternary(c, t, e) => vec![*c, *t, *e],
        Expr::UserOp { inputs, .. } => inputs.clone(),
        Expr::Const(_, _) | Expr::Unknown => vec![],
    }
}