ryo-mutations 0.2.0

//! Wrap expressions with `dbg!()` for debugging.
//!
//! This mutation wraps specified expressions with `dbg!()` macro calls,
//! allowing you to see the value and location of expressions.

use ryo_source::pure::{MacroDelimiter, PureBlock, PureExpr, PureFn, PureStmt};

use super::marker::DebugMarker;
use crate::Mutation;

/// Target for wrapping with dbg!().
#[derive(Debug, Clone)]
pub enum WrapTarget {
    /// Wrap all method chains starting with a specific method.
    MethodChain { starts_with: String },
    /// Wrap method calls to a specific method.
    MethodCall { method: String },
    /// Wrap all variables with a specific name.
    Variable { name: String },
    /// Wrap the return expression of a function.
    ReturnExpr { function: String },
}

/// Wrap expressions with `dbg!()` macro.
///
/// # Example
///
/// Before:
/// ```ignore
/// let result = items.iter().map(|x| x * 2).collect();
/// ```
///
/// After (with `MethodChain { starts_with: "iter" }`):
/// ```ignore
/// let result = dbg!(/* ryo-debug:... */ items.iter().map(|x| x * 2).collect());
/// ```
#[derive(Debug, Clone)]
pub struct DbgWrapMutation {
    /// What to wrap.
    pub target: WrapTarget,
    /// Debug marker for tracking.
    pub marker: DebugMarker,
    /// Only apply in this function (if specified).
    pub in_function: Option<String>,
}

impl DbgWrapMutation {
    /// Create a new mutation to wrap method chains.
    pub fn method_chain(starts_with: impl Into<String>, marker: DebugMarker) -> Self {
        Self {
            target: WrapTarget::MethodChain {
                starts_with: starts_with.into(),
            },
            marker,
            in_function: None,
        }
    }

    /// Create a new mutation to wrap specific method calls.
    pub fn method_call(method: impl Into<String>, marker: DebugMarker) -> Self {
        Self {
            target: WrapTarget::MethodCall {
                method: method.into(),
            },
            marker,
            in_function: None,
        }
    }

    /// Create a new mutation to wrap variable references.
    pub fn variable(name: impl Into<String>, marker: DebugMarker) -> Self {
        Self {
            target: WrapTarget::Variable { name: name.into() },
            marker,
            in_function: None,
        }
    }

    /// Create a new mutation to wrap return expressions.
    pub fn return_expr(function: impl Into<String>, marker: DebugMarker) -> Self {
        Self {
            target: WrapTarget::ReturnExpr {
                function: function.into(),
            },
            marker,
            in_function: None,
        }
    }

    /// Only apply in a specific function.
    pub fn in_function(mut self, name: impl Into<String>) -> Self {
        self.in_function = Some(name.into());
        self
    }

    /// Check if an expression matches the wrap target.
    fn matches_target(&self, expr: &PureExpr) -> bool {
        match &self.target {
            WrapTarget::MethodChain { starts_with } => {
                self.is_chain_starting_with(expr, starts_with)
            }
            WrapTarget::MethodCall { method } => {
                matches!(expr, PureExpr::MethodCall { method: m, .. } if m == method)
            }
            WrapTarget::Variable { name } => {
                matches!(expr, PureExpr::Path(p) if p == name)
            }
            WrapTarget::ReturnExpr { .. } => false, // Handled separately
        }
    }

    /// Check if a method chain starts with a specific method.
    fn is_chain_starting_with(&self, expr: &PureExpr, method: &str) -> bool {
        match expr {
            PureExpr::MethodCall {
                receiver,
                method: m,
                ..
            } => {
                if m == method {
                    // Check if receiver is not a method call (i.e., this is the first method)
                    !matches!(receiver.as_ref(), PureExpr::MethodCall { .. })
                } else {
                    // Check the receiver
                    self.is_chain_starting_with(receiver, method)
                }
            }
            _ => false,
        }
    }

    /// Build a dbg! wrapped expression.
    fn wrap_with_dbg(&self, expr: PureExpr) -> PureExpr {
        let marker_str = self.marker.to_marker_string();

        PureExpr::Macro {
            name: "dbg".to_string(),
            delimiter: MacroDelimiter::Paren,
            tokens: format!("\"{}\", {}", marker_str, expr_to_tokens(&expr)),
        }
    }

    /// Transform an expression, wrapping matches.
    fn transform_expr(&self, expr: &PureExpr) -> (PureExpr, usize) {
        // Check if this expression should be wrapped
        if self.matches_target(expr) {
            return (self.wrap_with_dbg(expr.clone()), 1);
        }

        // Recursively transform nested expressions
        match expr {
            PureExpr::MethodCall {
                receiver,
                method,
                args,
                ..
            } => {
                let (new_receiver, c1) = self.transform_expr(receiver);
                let mut total = c1;
                let new_args: Vec<_> = args
                    .iter()
                    .map(|a| {
                        let (new_a, c) = self.transform_expr(a);
                        total += c;
                        new_a
                    })
                    .collect();
                (
                    PureExpr::MethodCall {
                        receiver: Box::new(new_receiver),
                        method: method.clone(),
                        turbofish: None,
                        args: new_args,
                    },
                    total,
                )
            }

            PureExpr::Call { func, args } => {
                let (new_func, mut total) = self.transform_expr(func);
                let new_args: Vec<_> = args
                    .iter()
                    .map(|a| {
                        let (new_a, c) = self.transform_expr(a);
                        total += c;
                        new_a
                    })
                    .collect();
                (
                    PureExpr::Call {
                        func: Box::new(new_func),
                        args: new_args,
                    },
                    total,
                )
            }

            PureExpr::Binary { op, left, right } => {
                let (new_left, c1) = self.transform_expr(left);
                let (new_right, c2) = self.transform_expr(right);
                (
                    PureExpr::Binary {
                        op: op.clone(),
                        left: Box::new(new_left),
                        right: Box::new(new_right),
                    },
                    c1 + c2,
                )
            }

            PureExpr::Block { label, block } => {
                let (new_block, count) = self.transform_block(block);
                (
                    PureExpr::Block {
                        label: label.clone(),
                        block: new_block,
                    },
                    count,
                )
            }

            PureExpr::If {
                cond,
                then_branch,
                else_branch,
            } => {
                let (new_cond, c1) = self.transform_expr(cond);
                let (new_then, c2) = self.transform_block(then_branch);
                let (new_else, c3) = if let Some(e) = else_branch {
                    let (ne, c) = self.transform_expr(e);
                    (Some(Box::new(ne)), c)
                } else {
                    (None, 0)
                };
                (
                    PureExpr::If {
                        cond: Box::new(new_cond),
                        then_branch: new_then,
                        else_branch: new_else,
                    },
                    c1 + c2 + c3,
                )
            }

            PureExpr::Closure {
                params, ret, body, ..
            } => {
                let (new_body, count) = self.transform_expr(body);
                (
                    PureExpr::Closure {
                        is_async: false,
                        is_move: false,
                        params: params.clone(),
                        ret: ret.clone(),
                        body: Box::new(new_body),
                    },
                    count,
                )
            }

            _ => (expr.clone(), 0),
        }
    }

    /// Transform a block.
    fn transform_block(&self, block: &PureBlock) -> (PureBlock, usize) {
        let mut count = 0;
        let new_stmts: Vec<_> = block
            .stmts
            .iter()
            .map(|stmt| {
                let (new_stmt, c) = self.transform_stmt(stmt);
                count += c;
                new_stmt
            })
            .collect();
        (PureBlock { stmts: new_stmts }, count)
    }

    /// Transform a statement.
    fn transform_stmt(&self, stmt: &PureStmt) -> (PureStmt, usize) {
        match stmt {
            PureStmt::Local {
                pattern,
                ty,
                init,
                else_branch,
            } => {
                if let Some(init_expr) = init {
                    let (new_init, count) = self.transform_expr(init_expr);
                    (
                        PureStmt::Local {
                            pattern: pattern.clone(),
                            ty: ty.clone(),
                            init: Some(new_init),
                            else_branch: else_branch.clone(),
                        },
                        count,
                    )
                } else {
                    (stmt.clone(), 0)
                }
            }
            PureStmt::Expr(expr) => {
                let (new_expr, count) = self.transform_expr(expr);
                (PureStmt::Expr(new_expr), count)
            }
            PureStmt::Semi(expr) => {
                let (new_expr, count) = self.transform_expr(expr);
                (PureStmt::Semi(new_expr), count)
            }
            _ => (stmt.clone(), 0),
        }
    }

    /// Transform a function, handling return expression wrapping.
    pub fn transform_fn(&self, func: &PureFn) -> (PureFn, usize) {
        // Check if we should only apply to a specific function
        if let Some(ref target_fn) = self.in_function {
            if &func.name != target_fn {
                return (func.clone(), 0);
            }
        }

        let mut new_func = func.clone();
        let mut count = 0;

        // Handle return expression wrapping
        if let WrapTarget::ReturnExpr { function } = &self.target {
            if &func.name == function {
                // Wrap the last expression if it's the return value
                if let Some(last_stmt) = new_func.body.stmts.last_mut() {
                    if let PureStmt::Expr(expr) = last_stmt {
                        *last_stmt = PureStmt::Expr(self.wrap_with_dbg(expr.clone()));
                        count += 1;
                    }
                }
                return (new_func, count);
            }
        }

        // Regular transformation
        let (new_body, block_count) = self.transform_block(&func.body);
        new_func.body = new_body;
        (new_func, count + block_count)
    }
}

impl Mutation for DbgWrapMutation {
    fn describe(&self) -> String {
        let target_desc = match &self.target {
            WrapTarget::MethodChain { starts_with } => {
                format!("method chains starting with .{}()", starts_with)
            }
            WrapTarget::MethodCall { method } => {
                format!(".{}() calls", method)
            }
            WrapTarget::Variable { name } => {
                format!("variable '{}'", name)
            }
            WrapTarget::ReturnExpr { function } => {
                format!("return expression in {}", function)
            }
        };
        format!(
            "Wrap {} with dbg!() [session: {}]",
            target_desc, self.marker.session_id
        )
    }

    fn mutation_type(&self) -> &'static str {
        "DbgWrap"
    }

    fn box_clone(&self) -> Box<dyn Mutation> {
        Box::new(self.clone())
    }
}

/// Convert an expression to a token string (simplified).
fn expr_to_tokens(expr: &PureExpr) -> String {
    // This is a simplified version - the actual rendering would use to_source()
    match expr {
        PureExpr::Path(p) => p.clone(),
        PureExpr::Lit(l) => l.clone(),
        PureExpr::MethodCall {
            receiver,
            method,
            args,
            ..
        } => {
            let receiver_str = expr_to_tokens(receiver);
            let args_str: Vec<_> = args.iter().map(expr_to_tokens).collect();
            format!("{}.{}({})", receiver_str, method, args_str.join(", "))
        }
        PureExpr::Call { func, args } => {
            let func_str = expr_to_tokens(func);
            let args_str: Vec<_> = args.iter().map(expr_to_tokens).collect();
            format!("{}({})", func_str, args_str.join(", "))
        }
        PureExpr::Binary { op, left, right } => {
            format!("{} {} {}", expr_to_tokens(left), op, expr_to_tokens(right))
        }
        _ => "...".to_string(),
    }
}