ryo_mutations/idiom/

assign_op.rs

//! AssignOpMutation: Convert assignments to compound operators
//!
//! Transforms:
//! - `a = a + b` → `a += b`
//! - `a = a - b` → `a -= b`
//! - `a = a * b` → `a *= b`
//! - `a = a / b` → `a /= b`
//! - `a = a % b` → `a %= b`
//! - `a = a & b` → `a &= b`
//! - `a = a | b` → `a |= b`
//! - `a = a ^ b` → `a ^= b`
//! - `a = a << b` → `a <<= b`
//! - `a = a >> b` → `a >>= b`
//!
//! Corresponds to Clippy lint: `clippy::assign_op_pattern`

use ryo_source::pure::{PureBlock, PureExpr, PureStmt};
use ryo_symbol::SymbolId;

use crate::Mutation;

/// Convert assignments to compound assignment operators
///
/// # Example
///
/// ```rust,ignore
/// use ryo_mutations::idiom::AssignOpMutation;
///
/// let mutation = AssignOpMutation::new();
/// // Transforms: a = a + 1;
/// // Into:       a += 1;
/// ```
#[derive(Debug, Clone, Default)]
pub struct AssignOpMutation {
    /// Target function SymbolId. If None, applies to all functions.
    pub target_fn: Option<SymbolId>,
}

impl AssignOpMutation {
    pub fn new() -> Self {
        Self::default()
    }

    /// Only apply in a specific function
    pub fn in_function(mut self, id: SymbolId) -> Self {
        self.target_fn = Some(id);
        self
    }

    /// Get compound operator for binary operator
    fn compound_op(op: &str) -> Option<&'static str> {
        match op {
            "+" => Some("+="),
            "-" => Some("-="),
            "*" => Some("*="),
            "/" => Some("/="),
            "%" => Some("%="),
            "&" => Some("&="),
            "|" => Some("|="),
            "^" => Some("^="),
            "<<" => Some("<<="),
            ">>" => Some(">>="),
            _ => None,
        }
    }

    /// Check if two expressions are structurally equivalent
    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),
            // Support dereference: *x == *x
            (PureExpr::Unary { op: opa, expr: ea }, PureExpr::Unary { op: opb, expr: eb }) => {
                opa == opb && Self::expr_eq(ea, eb)
            }
            // Support reference: &x == &x, &mut x == &mut x
            (
                PureExpr::Ref {
                    is_mut: ma,
                    expr: ea,
                },
                PureExpr::Ref {
                    is_mut: mb,
                    expr: eb,
                },
            ) => ma == mb && Self::expr_eq(ea, eb),
            _ => false,
        }
    }

    /// Transform statements in a block
    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
    }

    /// Transform a single statement
    fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
        match stmt {
            PureStmt::Semi(expr) | PureStmt::Expr(expr) => {
                // Check for assignment pattern: target = target op rhs
                if let PureExpr::Binary { op, left, right } = expr {
                    if op == "=" {
                        // Check if right side is: left op something
                        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) {
                                    // Transform: a = a + b => a += b
                                    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;
                                }
                            }
                        }
                    }
                }

                // Recursively transform expressions
                self.transform_expr(expr)
            }
            PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
            _ => 0,
        }
    }

    /// Transform expressions (for nested blocks/closures)
    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() {
        // *x == *x
        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));
    }
}