analyssa 0.1.0

//! Symbolic evaluator for building expression trees from SSA operations.
//!
//! This module provides [`SymbolicEvaluator`], which tracks SSA operations
//! symbolically to build [`super::SymbolicExpr`] trees. Unlike
//! concrete evaluation (which computes known values), symbolic evaluation
//! preserves the relationship between operations as expression trees, enabling
//! host-side constraint solving.
//!
//! # Algorithm
//!
//! For each SSA operation, the evaluator:
//! 1. Looks up symbolic expressions for the operand variables
//! 2. If an operand has no expression (unknown), creates a `SymbolicExpr::Variable`
//!    reference to the operand's `SsaVarId`
//! 3. Combines operand expressions with the appropriate `SymbolicOp`
//! 4. Applies `simplify()` for constant folding and algebraic simplification
//! 5. Stores the result expression for the destination variable
//!
//! For `LoadArg` and `LoadLocal`, the evaluator traces back to the version-0
//! variable's expression, preserving the symbolic chain from the original
//! argument/local definition.
//!
//! # Relationship to SsaEvaluator
//!
//! `SymbolicEvaluator` builds expression trees without computing concrete values.
//! In contrast, `SsaEvaluator` computes both concrete and symbolic values using
//! the same `SymbolicExpr` representation but with concrete tracking.

use std::collections::HashMap;

use crate::{
    analysis::symbolic::{expr::SymbolicExpr, ops::SymbolicOp},
    ir::{function::SsaFunction, ops::SsaOp, value::ConstValue, variable::SsaVarId},
    target::Target,
    PointerSize,
};

/// Symbolic evaluator that builds expression trees from SSA operations.
///
/// Unlike `SsaEvaluator` which computes concrete values, `SymbolicEvaluator`
/// tracks operations symbolically, building `SymbolicExpr` trees that can
/// later be solved using Z3.
#[derive(Debug)]
pub struct SymbolicEvaluator<'a, T: Target> {
    ssa: &'a SsaFunction<T>,
    /// Expressions computed for each variable.
    expressions: HashMap<SsaVarId, SymbolicExpr<T>>,
    /// Target pointer size for native int/uint masking.
    pointer_size: PointerSize,
}

impl<'a, T: Target> SymbolicEvaluator<'a, T> {
    /// Creates a new symbolic evaluator for the given SSA function.
    ///
    /// The evaluator starts with no known expressions. Use [`set_symbolic`](Self::set_symbolic)
    /// or [`set_constant`](Self::set_constant) to initialize variable values before
    /// evaluating blocks.
    ///
    /// # Arguments
    ///
    /// * `ssa` - The SSA function to evaluate.
    /// * `ptr_size` - Target pointer size for native int/uint masking.
    ///
    /// # Returns
    ///
    /// A new evaluator with no expressions.
    #[must_use]
    pub fn new(ssa: &'a SsaFunction<T>, ptr_size: PointerSize) -> Self {
        Self {
            ssa,
            expressions: HashMap::new(),
            pointer_size: ptr_size,
        }
    }

    /// Sets a variable to a named symbolic value.
    ///
    /// Named symbolic values represent external inputs whose concrete values
    /// are unknown. Use this to mark the "state" variable in control flow
    /// unflattening.
    ///
    /// # Arguments
    ///
    /// * `var` - The SSA variable to set.
    /// * `name` - The symbolic name (e.g., "state").
    pub fn set_symbolic(&mut self, var: SsaVarId, name: impl Into<String>) {
        self.expressions.insert(var, SymbolicExpr::named(name));
    }

    /// Sets a variable to a typed constant value.
    ///
    /// Use this to provide known initial values for variables with type preservation.
    /// The caller is responsible for providing the correct `ConstValue` type.
    ///
    /// # Arguments
    ///
    /// * `var` - The SSA variable to set.
    /// * `value` - The typed constant value.
    pub fn set_constant(&mut self, var: SsaVarId, value: ConstValue<T>) {
        self.expressions.insert(var, SymbolicExpr::constant(value));
    }

    /// Gets the expression for a variable, if known.
    ///
    /// # Arguments
    ///
    /// * `var` - The SSA variable to look up.
    ///
    /// # Returns
    ///
    /// The expression for the variable, or `None` if not yet evaluated.
    #[must_use]
    pub fn get_expression(&self, var: SsaVarId) -> Option<&SymbolicExpr<T>> {
        self.expressions.get(&var)
    }

    /// Gets the expression for a variable, simplified.
    ///
    /// Returns a copy of the expression with constant folding and algebraic
    /// simplifications applied.
    ///
    /// # Arguments
    ///
    /// * `var` - The SSA variable to look up.
    ///
    /// # Returns
    ///
    /// A simplified copy of the expression, or `None` if not yet evaluated.
    #[must_use]
    pub fn get_simplified(&self, var: SsaVarId) -> Option<SymbolicExpr<T>> {
        self.expressions
            .get(&var)
            .map(|e| e.simplify(self.pointer_size))
    }

    /// Returns all computed expressions.
    ///
    /// # Returns
    ///
    /// A reference to the map from variable IDs to their symbolic expressions.
    #[must_use]
    pub fn expressions(&self) -> &HashMap<SsaVarId, SymbolicExpr<T>> {
        &self.expressions
    }

    /// Evaluates all instructions in a block symbolically.
    ///
    /// Processes each instruction in the block, building symbolic expressions
    /// for any variables they define.
    ///
    /// # Arguments
    ///
    /// * `block_idx` - The index of the block to evaluate.
    pub fn evaluate_block(&mut self, block_idx: usize) {
        let Some(block) = self.ssa.block(block_idx) else {
            return;
        };

        for instr in block.instructions() {
            self.evaluate_op(instr.op());
        }
    }

    /// Evaluates a sequence of blocks in order.
    ///
    /// # Arguments
    ///
    /// * `block_indices` - The indices of blocks to evaluate, in order.
    pub fn evaluate_blocks(&mut self, block_indices: &[usize]) {
        for &block_idx in block_indices {
            self.evaluate_block(block_idx);
        }
    }

    /// Evaluates a single SSA operation symbolically.
    ///
    /// Builds a symbolic expression for the operation's result based on its
    /// operands. If operands have known expressions, those are used; otherwise,
    /// the operands are treated as symbolic variables.
    ///
    /// # Arguments
    ///
    /// * `op` - The SSA operation to evaluate.
    pub fn evaluate_op(&mut self, op: &SsaOp<T>) {
        match op {
            SsaOp::Const { dest, value } => {
                self.expressions
                    .insert(*dest, SymbolicExpr::constant(value.clone()));
            }

            SsaOp::Copy { dest, src } => {
                if let Some(expr) = self.expressions.get(src) {
                    self.expressions.insert(*dest, expr.clone());
                } else {
                    self.expressions.insert(*dest, SymbolicExpr::variable(*src));
                }
            }

            SsaOp::Add {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Add);
            }
            SsaOp::Sub {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Sub);
            }
            SsaOp::Mul {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Mul);
            }
            SsaOp::Div {
                dest,
                left,
                right,
                unsigned,
                ..
            } => {
                let op = if *unsigned {
                    SymbolicOp::DivU
                } else {
                    SymbolicOp::DivS
                };
                self.eval_binary(*dest, *left, *right, op);
            }
            SsaOp::Rem {
                dest,
                left,
                right,
                unsigned,
                ..
            } => {
                let op = if *unsigned {
                    SymbolicOp::RemU
                } else {
                    SymbolicOp::RemS
                };
                self.eval_binary(*dest, *left, *right, op);
            }
            SsaOp::Xor {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Xor);
            }
            SsaOp::And {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::And);
            }
            SsaOp::Or {
                dest, left, right, ..
            } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Or);
            }
            SsaOp::Shl {
                dest,
                value,
                amount,
                ..
            } => {
                self.eval_binary(*dest, *value, *amount, SymbolicOp::Shl);
            }
            SsaOp::Shr {
                dest,
                value,
                amount,
                unsigned,
                ..
            } => {
                let op = if *unsigned {
                    SymbolicOp::ShrU
                } else {
                    SymbolicOp::ShrS
                };
                self.eval_binary(*dest, *value, *amount, op);
            }
            SsaOp::Neg { dest, operand, .. } => {
                self.eval_unary(*dest, *operand, SymbolicOp::Neg);
            }
            SsaOp::Not { dest, operand, .. } => {
                self.eval_unary(*dest, *operand, SymbolicOp::Not);
            }
            SsaOp::Ceq { dest, left, right } => {
                self.eval_binary(*dest, *left, *right, SymbolicOp::Eq);
            }
            SsaOp::Cgt {
                dest,
                left,
                right,
                unsigned,
            } => {
                let op = if *unsigned {
                    SymbolicOp::GtU
                } else {
                    SymbolicOp::GtS
                };
                self.eval_binary(*dest, *left, *right, op);
            }
            SsaOp::Clt {
                dest,
                left,
                right,
                unsigned,
            } => {
                let op = if *unsigned {
                    SymbolicOp::LtU
                } else {
                    SymbolicOp::LtS
                };
                self.eval_binary(*dest, *left, *right, op);
            }
            SsaOp::Conv { dest, operand, .. } => {
                if let Some(expr) = self.expressions.get(operand) {
                    self.expressions.insert(*dest, expr.clone());
                } else {
                    self.expressions
                        .insert(*dest, SymbolicExpr::variable(*operand));
                }
            }

            // LoadArg/LoadLocal: propagate expression from the source argument/local variable
            SsaOp::LoadArg { dest, arg_index } => {
                // Find the argument variable (version 0) and propagate its expression
                if let Some(arg_var) = self
                    .ssa
                    .variables_from_argument(*arg_index)
                    .find(|v| v.version() == 0)
                {
                    let src = arg_var.id();
                    if let Some(expr) = self.expressions.get(&src).cloned() {
                        self.expressions.insert(*dest, expr);
                    } else {
                        self.expressions.insert(*dest, SymbolicExpr::variable(src));
                    }
                }
            }
            SsaOp::LoadLocal { dest, local_index } => {
                if let Some(local_var) = self
                    .ssa
                    .variables_from_local(*local_index)
                    .find(|v| v.version() == 0)
                {
                    let src = local_var.id();
                    if let Some(expr) = self.expressions.get(&src).cloned() {
                        self.expressions.insert(*dest, expr);
                    } else {
                        self.expressions.insert(*dest, SymbolicExpr::variable(src));
                    }
                }
            }

            // Rotate operations
            SsaOp::Rol {
                dest,
                value,
                amount,
            } => {
                self.eval_binary(*dest, *value, *amount, SymbolicOp::Rol);
            }
            SsaOp::Ror {
                dest,
                value,
                amount,
            } => {
                self.eval_binary(*dest, *value, *amount, SymbolicOp::Ror);
            }
            SsaOp::Rcl {
                dest,
                value,
                amount,
            } => {
                self.eval_binary(*dest, *value, *amount, SymbolicOp::Rcl);
            }
            SsaOp::Rcr {
                dest,
                value,
                amount,
            } => {
                self.eval_binary(*dest, *value, *amount, SymbolicOp::Rcr);
            }

            // Bit manipulation operations
            SsaOp::BSwap { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::BSwap);
            }
            SsaOp::BRev { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::BRev);
            }
            SsaOp::BitScanForward { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::BitScanForward);
            }
            SsaOp::BitScanReverse { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::BitScanReverse);
            }
            SsaOp::Popcount { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::Popcount);
            }
            SsaOp::Parity { dest, src } => {
                self.eval_unary(*dest, *src, SymbolicOp::Parity);
            }

            // Side-effecting operations have no meaningful symbolic result,
            // so they fall through to the catch-all.
            _ => {}
        }
    }

    /// Evaluates a binary operation and stores the result expression.
    ///
    /// Looks up expressions for the operands (or creates variable references
    /// if unknown), combines them with the operation, simplifies, and stores.
    ///
    /// # Arguments
    ///
    /// * `dest` - The destination variable for the result.
    /// * `left` - The left operand variable.
    /// * `right` - The right operand variable.
    /// * `op` - The binary operation to apply.
    fn eval_binary(&mut self, dest: SsaVarId, left: SsaVarId, right: SsaVarId, op: SymbolicOp) {
        let left_expr = self
            .expressions
            .get(&left)
            .cloned()
            .unwrap_or_else(|| SymbolicExpr::variable(left));
        let right_expr = self
            .expressions
            .get(&right)
            .cloned()
            .unwrap_or_else(|| SymbolicExpr::variable(right));

        let result = SymbolicExpr::binary(op, left_expr, right_expr).simplify(self.pointer_size);
        self.expressions.insert(dest, result);
    }

    /// Evaluates a unary operation and stores the result expression.
    ///
    /// Looks up the expression for the operand (or creates a variable reference
    /// if unknown), applies the operation, simplifies, and stores.
    ///
    /// # Arguments
    ///
    /// * `dest` - The destination variable for the result.
    /// * `operand` - The operand variable.
    /// * `op` - The unary operation to apply.
    fn eval_unary(&mut self, dest: SsaVarId, operand: SsaVarId, op: SymbolicOp) {
        let operand_expr = self
            .expressions
            .get(&operand)
            .cloned()
            .unwrap_or_else(|| SymbolicExpr::variable(operand));

        let result = SymbolicExpr::unary(op, operand_expr).simplify(self.pointer_size);
        self.expressions.insert(dest, result);
    }
}