use ryo_source::pure::{PureBlock, PureExpr, PureStmt};
use ryo_symbol::SymbolId;
use crate::Mutation;
#[derive(Debug, Clone, Default)]
pub struct AssignOpMutation {
pub target_fn: Option<SymbolId>,
}
impl AssignOpMutation {
pub fn new() -> Self {
Self::default()
}
pub fn in_function(mut self, id: SymbolId) -> Self {
self.target_fn = Some(id);
self
}
fn compound_op(op: &str) -> Option<&'static str> {
match op {
"+" => Some("+="),
"-" => Some("-="),
"*" => Some("*="),
"/" => Some("/="),
"%" => Some("%="),
"&" => Some("&="),
"|" => Some("|="),
"^" => Some("^="),
"<<" => Some("<<="),
">>" => Some(">>="),
_ => None,
}
}
fn expr_eq(a: &PureExpr, b: &PureExpr) -> bool {
match (a, b) {
(PureExpr::Path(pa), PureExpr::Path(pb)) => pa == pb,
(
PureExpr::Field {
expr: ea,
field: fa,
},
PureExpr::Field {
expr: eb,
field: fb,
},
) => fa == fb && Self::expr_eq(ea, eb),
(
PureExpr::Index {
expr: ea,
index: ia,
},
PureExpr::Index {
expr: eb,
index: ib,
},
) => Self::expr_eq(ea, eb) && Self::expr_eq(ia, ib),
(PureExpr::Unary { op: opa, expr: ea }, PureExpr::Unary { op: opb, expr: eb }) => {
opa == opb && Self::expr_eq(ea, eb)
}
(
PureExpr::Ref {
is_mut: ma,
expr: ea,
},
PureExpr::Ref {
is_mut: mb,
expr: eb,
},
) => ma == mb && Self::expr_eq(ea, eb),
_ => false,
}
}
pub fn transform_block(&self, block: &mut PureBlock) -> usize {
let mut changes = 0;
for stmt in &mut block.stmts {
changes += self.transform_stmt(stmt);
}
changes
}
fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
match stmt {
PureStmt::Semi(expr) | PureStmt::Expr(expr) => {
if let PureExpr::Binary { op, left, right } = expr {
if op == "=" {
if let PureExpr::Binary {
op: bin_op,
left: bin_left,
right: bin_right,
} = right.as_ref()
{
if Self::expr_eq(left, bin_left) {
if let Some(compound) = Self::compound_op(bin_op) {
let target = std::mem::replace(
left.as_mut(),
PureExpr::Path("__placeholder".to_string()),
);
let rhs = bin_right.as_ref().clone();
*expr = PureExpr::Binary {
op: compound.to_string(),
left: Box::new(target),
right: Box::new(rhs),
};
return 1;
}
}
}
}
}
self.transform_expr(expr)
}
PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
_ => 0,
}
}
fn transform_expr(&self, expr: &mut PureExpr) -> usize {
let mut changes = 0;
match expr {
PureExpr::Block { block, .. } => {
changes += self.transform_block(block);
}
PureExpr::If {
cond,
then_branch,
else_branch,
} => {
changes += self.transform_expr(cond);
changes += self.transform_block(then_branch);
if let Some(else_expr) = else_branch {
changes += self.transform_expr(else_expr);
}
}
PureExpr::Match { expr: e, arms } => {
changes += self.transform_expr(e);
for arm in arms {
changes += self.transform_expr(&mut arm.body);
}
}
PureExpr::Loop { body: block, .. } | PureExpr::While { body: block, .. } => {
changes += self.transform_block(block);
}
PureExpr::For { body, .. } => {
changes += self.transform_block(body);
}
PureExpr::Closure { body, .. } => {
changes += self.transform_expr(body);
}
_ => {}
}
changes
}
}
impl Mutation for AssignOpMutation {
fn describe(&self) -> String {
"Convert assignments to compound operators (a = a + b → a += b)".to_string()
}
fn mutation_type(&self) -> &'static str {
"AssignOp"
}
fn box_clone(&self) -> Box<dyn Mutation> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compound_op() {
assert_eq!(AssignOpMutation::compound_op("+"), Some("+="));
assert_eq!(AssignOpMutation::compound_op("-"), Some("-="));
assert_eq!(AssignOpMutation::compound_op("*"), Some("*="));
assert_eq!(AssignOpMutation::compound_op("/"), Some("/="));
assert_eq!(AssignOpMutation::compound_op("&&"), None);
}
#[test]
fn test_expr_eq() {
let a = PureExpr::Path("x".to_string());
let b = PureExpr::Path("x".to_string());
let c = PureExpr::Path("y".to_string());
assert!(AssignOpMutation::expr_eq(&a, &b));
assert!(!AssignOpMutation::expr_eq(&a, &c));
}
#[test]
fn test_expr_eq_unary() {
let a = PureExpr::Unary {
op: "*".to_string(),
expr: Box::new(PureExpr::Path("x".to_string())),
};
let b = PureExpr::Unary {
op: "*".to_string(),
expr: Box::new(PureExpr::Path("x".to_string())),
};
let c = PureExpr::Unary {
op: "*".to_string(),
expr: Box::new(PureExpr::Path("y".to_string())),
};
assert!(AssignOpMutation::expr_eq(&a, &b));
assert!(!AssignOpMutation::expr_eq(&a, &c));
}
#[test]
fn test_expr_eq_field() {
let a = PureExpr::Field {
expr: Box::new(PureExpr::Path("self".to_string())),
field: "count".to_string(),
};
let b = PureExpr::Field {
expr: Box::new(PureExpr::Path("self".to_string())),
field: "count".to_string(),
};
let c = PureExpr::Field {
expr: Box::new(PureExpr::Path("self".to_string())),
field: "other".to_string(),
};
assert!(AssignOpMutation::expr_eq(&a, &b));
assert!(!AssignOpMutation::expr_eq(&a, &c));
}
}