ryo_executor/engine/impls/

introduce_variable.rs

//! V2 ASTRegApply implementation for IntroduceVariableMutation
//!
//! Extracts complex expressions into named variables:
//! - `foo(a + b * c, a + b * c)` → `let x = a + b * c; foo(x, x)`
//!
//! # Implementation Strategy
//!
//! 1. Find target function(s) in the registry
//! 2. Recursively traverse the function body to find matching expressions
//! 3. Replace matching expressions with variable references
//! 4. Insert a `let` statement at the beginning of the function body

use ryo_analysis::SymbolKind;
use ryo_mutations::idiom::IntroduceVariableMutation;
use ryo_mutations::{Mutation, MutationResult};
use ryo_source::pure::{PureBlock, PureExpr, PureItem, PurePattern, PureStmt};

use crate::engine::{ASTMutationContext, ASTRegApply};

impl ASTRegApply for IntroduceVariableMutation {
    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
        let mut total_replacements = 0;

        // Collect function IDs to modify
        let fn_ids: Vec<_> = if let Some(target_id) = self.target_fn {
            // Target function specified: direct lookup
            vec![target_id]
        } else {
            // No target: collect all functions
            ctx.symbol_registry
                .iter()
                .filter(|(id, _)| {
                    matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Function))
                })
                .map(|(id, _)| id)
                .collect()
        };

        for fn_id in fn_ids {
            if let Some(PureItem::Fn(func)) = ctx.ast_registry.get(fn_id) {
                let mut new_func = func.clone();
                let replacements =
                    replace_expr_in_block(&mut new_func.body, &self.target_expr, &self.var_name);

                if replacements > 0 {
                    // Insert the variable declaration at the beginning
                    let let_stmt = PureStmt::Local {
                        pattern: PurePattern::Ident {
                            name: self.var_name.clone(),
                            is_mut: self.is_mut,
                        },
                        ty: None, // Type inference
                        init: Some(self.target_expr.clone()),
                    };
                    new_func.body.stmts.insert(0, let_stmt);

                    ctx.set_ast(fn_id, PureItem::Fn(new_func));
                    total_replacements += replacements;
                }
            }
        }

        MutationResult {
            mutation_type: self.mutation_type().to_string(),
            changes: total_replacements,
            description: if total_replacements > 0 {
                format!(
                    "Introduced variable '{}' ({} replacements)",
                    self.var_name, total_replacements
                )
            } else {
                format!("No matching expressions found for '{}'", self.var_name)
            },
        }
    }
}

/// Replace all occurrences of `target` expression with a variable reference in a block.
/// Returns the number of replacements made.
fn replace_expr_in_block(block: &mut PureBlock, target: &PureExpr, var_name: &str) -> usize {
    let mut count = 0;
    for stmt in &mut block.stmts {
        count += replace_expr_in_stmt(stmt, target, var_name);
    }
    count
}

/// Replace expressions in a statement
fn replace_expr_in_stmt(stmt: &mut PureStmt, target: &PureExpr, var_name: &str) -> usize {
    match stmt {
        PureStmt::Local { init, .. } => {
            if let Some(expr) = init {
                replace_expr(expr, target, var_name)
            } else {
                0
            }
        }
        PureStmt::Semi(expr) | PureStmt::Expr(expr) => replace_expr(expr, target, var_name),
        PureStmt::Item(_) => 0, // Don't descend into nested items
    }
}

/// Replace target expression with variable reference, recursively.
/// Returns the number of replacements made.
fn replace_expr(expr: &mut PureExpr, target: &PureExpr, var_name: &str) -> usize {
    // Check if this expression matches the target
    if expr == target {
        *expr = PureExpr::Path(var_name.to_string());
        return 1;
    }

    // Recursively check children
    match expr {
        PureExpr::Binary { left, right, .. } => {
            replace_expr(left, target, var_name) + replace_expr(right, target, var_name)
        }
        PureExpr::Unary { expr: inner, .. } => replace_expr(inner, target, var_name),
        PureExpr::Call { func, args } => {
            let mut count = replace_expr(func, target, var_name);
            for arg in args {
                count += replace_expr(arg, target, var_name);
            }
            count
        }
        PureExpr::MethodCall { receiver, args, .. } => {
            let mut count = replace_expr(receiver, target, var_name);
            for arg in args {
                count += replace_expr(arg, target, var_name);
            }
            count
        }
        PureExpr::Field { expr: inner, .. } => replace_expr(inner, target, var_name),
        PureExpr::Index { expr: e, index } => {
            replace_expr(e, target, var_name) + replace_expr(index, target, var_name)
        }
        PureExpr::Block { block, .. } => replace_expr_in_block(block, target, var_name),
        PureExpr::If {
            cond,
            then_branch,
            else_branch,
        } => {
            let mut count = replace_expr(cond, target, var_name);
            count += replace_expr_in_block(then_branch, target, var_name);
            if let Some(else_expr) = else_branch {
                count += replace_expr(else_expr, target, var_name);
            }
            count
        }
        PureExpr::Match { expr: e, arms } => {
            let mut count = replace_expr(e, target, var_name);
            for arm in arms {
                count += replace_expr(&mut arm.body, target, var_name);
                if let Some(guard) = &mut arm.guard {
                    count += replace_expr(guard, target, var_name);
                }
            }
            count
        }
        PureExpr::Loop { body: block, .. } | PureExpr::While { body: block, .. } => {
            replace_expr_in_block(block, target, var_name)
        }
        PureExpr::For { expr: e, body, .. } => {
            replace_expr(e, target, var_name) + replace_expr_in_block(body, target, var_name)
        }
        PureExpr::Return(Some(inner))
        | PureExpr::Break {
            expr: Some(inner), ..
        } => replace_expr(inner, target, var_name),
        PureExpr::Closure { body, .. } => replace_expr(body, target, var_name),
        PureExpr::Struct { fields, .. } => {
            let mut count = 0;
            for (_, field_expr) in fields {
                count += replace_expr(field_expr, target, var_name);
            }
            count
        }
        PureExpr::Tuple(exprs) | PureExpr::Array(exprs) => {
            let mut count = 0;
            for e in exprs {
                count += replace_expr(e, target, var_name);
            }
            count
        }
        PureExpr::Ref { expr: inner, .. } => replace_expr(inner, target, var_name),
        PureExpr::Await(inner) | PureExpr::Try(inner) => replace_expr(inner, target, var_name),
        PureExpr::Range { start, end, .. } => {
            let mut count = 0;
            if let Some(s) = start {
                count += replace_expr(s, target, var_name);
            }
            if let Some(e) = end {
                count += replace_expr(e, target, var_name);
            }
            count
        }
        PureExpr::Cast { expr: inner, .. } => replace_expr(inner, target, var_name),
        PureExpr::Let { expr: inner, .. } => replace_expr(inner, target, var_name),
        PureExpr::Async { body, .. } | PureExpr::Unsafe(body) => {
            replace_expr_in_block(body, target, var_name)
        }
        PureExpr::Repeat { expr: e, len } => {
            replace_expr(e, target, var_name) + replace_expr(len, target, var_name)
        }
        // Leaf nodes - no children to recurse into
        PureExpr::Lit(_)
        | PureExpr::Path(_)
        | PureExpr::Macro { .. }
        | PureExpr::Return(None)
        | PureExpr::Break { expr: None, .. }
        | PureExpr::Continue { .. }
        | PureExpr::Other(_) => 0,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::ASTMutationEngine;
    use ryo_analysis::testing::ContextBuilder;

    #[test]
    fn test_v2_introduce_variable() {
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"
fn compute() -> i32 {
    let x = 1 + 2;
    let y = 1 + 2;
    x + y
}
"#,
            )
            .build();

        // Create target expression: 1 + 2
        let target = PureExpr::Binary {
            op: "+".to_string(),
            left: Box::new(PureExpr::Lit("1".to_string())),
            right: Box::new(PureExpr::Lit("2".to_string())),
        };

        let mutation = IntroduceVariableMutation {
            target_expr: target,
            var_name: "sum".to_string(),
            target_fn: None, // Note: target_fn requires SymbolId, not name
            is_mut: false,
        };

        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);

        // Should have found and replaced the expression
        // Note: The actual number depends on how the parser represents the code
        println!("Result: {:?}", result.result);
    }

    #[test]
    fn test_v2_introduce_variable_no_match() {
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"
fn simple() -> i32 {
    42
}
"#,
            )
            .build();

        let target = PureExpr::Binary {
            op: "+".to_string(),
            left: Box::new(PureExpr::Lit("1".to_string())),
            right: Box::new(PureExpr::Lit("2".to_string())),
        };

        let mutation = IntroduceVariableMutation {
            target_expr: target,
            var_name: "sum".to_string(),
            target_fn: None,
            is_mut: false,
        };

        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);

        // No matching expression should be found
        assert_eq!(result.result.changes, 0);
    }
}