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;
let fn_ids: Vec<_> = if let Some(target_id) = self.target_fn {
vec![target_id]
} else {
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 {
let let_stmt = PureStmt::Local {
pattern: PurePattern::Ident {
name: self.var_name.clone(),
is_mut: self.is_mut,
},
ty: None, 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)
},
}
}
}
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
}
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, }
}
fn replace_expr(expr: &mut PureExpr, target: &PureExpr, var_name: &str) -> usize {
if expr == target {
*expr = PureExpr::Path(var_name.to_string());
return 1;
}
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)
}
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();
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);
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);
assert_eq!(result.result.changes, 0);
}
}