neo-decompiler 0.10.1

//! SSA optimization passes.
//!
//! Operates on a built [`SsaForm`] (see [`crate::decompiler::cfg::ssa::SsaBuilder`])
//! and returns a simplified, semantically-equivalent form. The passes are the
//! classic SSA optimizations from the project roadmap:
//!
//! - **Constant folding / propagation**: `v2 = (1 + 2)` folds to `v2 = 3`, and
//!   that constant is then substituted into every use of `v2`.
//! - **Copy propagation**: `v1 = v0` substitutes `v0` for `v1` at every use.
//! - **Trivial-φ elimination**: a φ all of whose operands are the same value
//!   (or a single operand) is replaced by that value everywhere.
//! - **Dead-code elimination**: pure defs with no remaining uses are removed.
//!
//! The passes share a single substitution table and run to a fixed point.

use std::collections::BTreeMap;

mod fold;
mod indexes;
mod rewrite;

use super::form::{SsaExpr, SsaForm, SsaStmt};
use super::variable::SsaVariable;
use fold::{as_literal, fold_binary, fold_unary};
use indexes::{collect_used, rebuild_indexes};
use rewrite::{resolve_once, rewrite_expr};

#[cfg(test)]
mod tests;

/// Optimize `ssa` in place by running the SSA optimization passes to a fixed
/// point. Returns how many rewrite rounds were applied (0 = already optimal).
pub fn optimize(ssa: &mut SsaForm) -> usize {
    let mut rounds = 0usize;
    loop {
        let rewrites = one_round(ssa);
        if rewrites == 0 {
            break;
        }
        rounds += 1;
        // Guard against pathological non-convergence (monotone rewrites always
        // terminate, but cap defensively).
        if rounds > ssa.blocks.len() + 4 {
            break;
        }
    }
    rounds
}

/// A single substitution table mapping a variable to the simpler expression
/// that replaces every reference to it.
type Subst = BTreeMap<SsaVariable, SsaExpr>;

fn one_round(ssa: &mut SsaForm) -> usize {
    // Gather constants/copies from current assignments.
    let mut subst: Subst = BTreeMap::new();

    for (_bid, block) in ssa.blocks.iter() {
        for stmt in &block.stmts {
            if let SsaStmt::Assign { target, value } = stmt {
                match value {
                    // Direct constant: propagate the literal — but NOT out of a
                    // named slot variable (loc/arg/static). A slot is a
                    // user-visible variable; substituting its constant into uses
                    // would erase the variable and fold away branch conditions /
                    // arithmetic the structurer needs (e.g. `loc0 = 0; if loc0<3`
                    // must survive, not become `if true`). Anonymous temps fold.
                    SsaExpr::Literal(_) => {
                        if !is_slot_var(target) {
                            subst.insert(target.clone(), value.clone());
                        }
                    }
                    // Copy of another var: `t = v` -> use v. For a slot target,
                    // only collapse to ANOTHER slot (a redundant load-alias);
                    // never to a temp/literal, so the slot reference stays
                    // visible at use sites.
                    SsaExpr::Variable(src) => {
                        if !is_slot_var(target) || is_slot_var(src) {
                            subst.insert(target.clone(), value.clone());
                        }
                    }
                    // Foldable pure binary on two constants.
                    SsaExpr::Binary { op, left, right } => {
                        if let (Some(a), Some(b)) = (as_literal(left), as_literal(right)) {
                            if let Some(folded) = fold_binary(*op, &a, &b) {
                                subst.insert(target.clone(), SsaExpr::lit(folded));
                            }
                        }
                    }
                    // Foldable pure unary on a constant.
                    SsaExpr::Unary { op, operand } => {
                        if let Some(a) = as_literal(operand) {
                            if let Some(folded) = fold_unary(*op, &a) {
                                subst.insert(target.clone(), SsaExpr::lit(folded));
                            }
                        }
                    }
                    _ => {}
                }
            }
        }
    }

    // Resolve transitive substitutions (v2->3, v1->v2 -> v1->3) and chase copy
    // chains (v1->v0, v0->3 -> v1->3).
    for (_bid, block) in ssa.blocks.iter() {
        for stmt in &block.stmts {
            if let SsaStmt::Assign { target, value } = stmt {
                if !subst.contains_key(target) {
                    continue;
                }
                // Only chase if the current mapping is a var/lit we can improve.
                if let SsaExpr::Variable(src) = value {
                    if let Some(resolved) = subst.get(src) {
                        subst.insert(target.clone(), resolved.clone());
                    }
                }
            }
        }
    }

    // Trivial-φ elimination: φ whose operands are all the same value (after
    // substitution) collapses to that value.
    for (_bid, block) in ssa.blocks.iter() {
        for phi in &block.phi_nodes {
            let operands: Vec<&SsaVariable> = phi.operands.values().collect();
            if operands.is_empty() {
                continue;
            }
            let first = operands[0];
            let mut resolved_first = subst
                .get(first)
                .cloned()
                .unwrap_or_else(|| SsaExpr::var(first.clone()));
            resolved_first = resolve_once(&subst, &resolved_first);
            let all_same = operands.iter().all(|v| {
                let r = subst
                    .get(*v)
                    .cloned()
                    .unwrap_or_else(|| SsaExpr::var((*v).clone()));
                resolve_once(&subst, &r) == resolved_first
            });
            if all_same {
                subst.insert(phi.target.clone(), resolved_first);
            }
        }
    }

    if subst.is_empty() {
        return 0;
    }

    // Apply substitutions to every expression reference and drop defs that have
    // been folded away (an assignment whose target is now a pure literal/copy
    // and is unused is dead).
    let used = collect_used(ssa, &subst);

    let mut rewrites = 0usize;
    for (_bid, block) in ssa.blocks.iter_mut() {
        // Rewrite φ operands.
        for phi in &mut block.phi_nodes {
            for (_pred, var) in phi.operands.iter_mut() {
                if let Some(SsaExpr::Variable(rep_var)) = subst.get(var) {
                    // φ operands are variables; only substitute when the
                    // replacement is itself a variable (constant φ results are
                    // reflected through the target substitution instead).
                    *var = rep_var.clone();
                    rewrites += 1;
                }
            }
        }
        // Rewrite statement RHS expressions and drop dead defs.
        block.stmts.retain(|stmt| match stmt {
            SsaStmt::Assign { target, .. } => {
                // Keep if the target has uses, or if it isn't a pure
                // constant/copy (i.e. we didn't substitute it).
                !subst.contains_key(target) || used.contains(target)
            }
            _ => true,
        });
        for stmt in &mut block.stmts {
            if let SsaStmt::Assign { target, value } = stmt {
                let new_value = rewrite_expr(value, &subst);
                if new_value != *value {
                    *value = new_value;
                    rewrites += 1;
                }
                // Avoid leaving a redundant `vN = vN` after substitution.
                if let SsaExpr::Variable(src) = value {
                    if src == target {
                        // self-assignment: will be DCE'd next round if unused.
                    }
                }
            }
        }
    }

    // Rebuild the definitions/uses indexes so downstream consumers stay correct.
    rebuild_indexes(ssa);
    rewrites.max(1)
}

/// Whether `v` is a named slot variable (a local/argument/static field). Such
/// variables are user-visible; the optimizer keeps them symbolic instead of
/// substituting their constant value into uses (see `one_round`). Slot bases are
/// `locN` / `argN` / `staticN` as produced by the SSA builder's `slot_name_for`.
fn is_slot_var(v: &SsaVariable) -> bool {
    let b = v.base.as_str();
    b.starts_with("loc") || b.starts_with("arg") || b.starts_with("static")
}