ryo_mutations/idiom/

unwrap_to_question.rs

//! UnwrapToQuestionMutation: Convert .unwrap()/.expect() to ? operator
//!
//! Converts:
//! - `x.unwrap()` → `x?`
//! - `x.expect("msg")` → `x?`
//! - `x.unwrap_or_else(|| panic!(...))` → `x?`
//!
//! Note: Only applies in functions that return Result/Option

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

use crate::Mutation;

/// Convert .unwrap() and .expect() to ? operator
#[derive(Debug, Clone, Default)]
pub struct UnwrapToQuestionMutation {
    /// Also convert .expect() calls
    pub include_expect: bool,
    /// Target function SymbolId. If None, applies to all functions.
    pub target_fn: Option<SymbolId>,
}

impl UnwrapToQuestionMutation {
    pub fn new() -> Self {
        Self {
            include_expect: true,
            target_fn: None,
        }
    }

    /// Don't convert .expect() calls (keep the error message)
    pub fn unwrap_only(mut self) -> Self {
        self.include_expect = false;
        self
    }

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

    /// Check if method call is .unwrap() or .expect()
    fn is_unwrap_call(&self, method: &str, args: &[PureExpr]) -> bool {
        match method {
            "unwrap" if args.is_empty() => true,
            "expect" if args.len() == 1 && self.include_expect => true,
            "unwrap_or_else" if args.len() == 1 => {
                // Check if the closure panics
                if let PureExpr::Closure { body, .. } = &args[0] {
                    Self::is_panic_expr(body)
                } else {
                    false
                }
            }
            _ => false,
        }
    }

    /// Check if expression is a panic
    fn is_panic_expr(expr: &PureExpr) -> bool {
        match expr {
            PureExpr::Macro { name, .. } => {
                name == "panic"
                    || name == "unreachable"
                    || name == "todo"
                    || name == "unimplemented"
            }
            PureExpr::Block { block, .. } if block.stmts.len() == 1 => match &block.stmts[0] {
                PureStmt::Expr(e) | PureStmt::Semi(e) => Self::is_panic_expr(e),
                _ => false,
            },
            _ => false,
        }
    }

    /// Transform an expression (recursively), returns (new_expr, changes_count)
    fn transform_expr(&self, expr: &mut PureExpr) -> usize {
        let mut changes = 0;

        // First, check if this is an unwrap call we should transform
        if let PureExpr::MethodCall {
            receiver,
            method,
            args,
            ..
        } = expr
        {
            if self.is_unwrap_call(method, args) {
                // Transform: receiver.unwrap() → receiver?
                // First transform the receiver recursively
                changes += self.transform_expr(receiver);

                // Then wrap in Try
                let inner = std::mem::replace(
                    receiver.as_mut(),
                    PureExpr::Path("__placeholder".to_string()),
                );
                *expr = PureExpr::Try(Box::new(inner));
                return changes + 1;
            }
        }

        // Recursively transform sub-expressions
        match expr {
            PureExpr::Binary { left, right, .. } => {
                changes += self.transform_expr(left);
                changes += self.transform_expr(right);
            }
            PureExpr::Unary { expr: inner, .. } => {
                changes += self.transform_expr(inner);
            }
            PureExpr::Call { func, args } => {
                changes += self.transform_expr(func);
                for arg in args {
                    changes += self.transform_expr(arg);
                }
            }
            PureExpr::MethodCall { receiver, args, .. } => {
                changes += self.transform_expr(receiver);
                for arg in args {
                    changes += self.transform_expr(arg);
                }
            }
            PureExpr::Field { expr: inner, .. } => {
                changes += self.transform_expr(inner);
            }
            PureExpr::Index { expr: inner, index } => {
                changes += self.transform_expr(inner);
                changes += self.transform_expr(index);
            }
            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 {
                expr: iter_expr,
                body,
                ..
            } => {
                changes += self.transform_expr(iter_expr);
                changes += self.transform_block(body);
            }
            PureExpr::Closure { body, .. } => {
                changes += self.transform_expr(body);
            }
            PureExpr::Tuple(exprs) | PureExpr::Array(exprs) => {
                for e in exprs {
                    changes += self.transform_expr(e);
                }
            }
            PureExpr::Struct { fields, .. } => {
                for (_, e) in fields {
                    changes += self.transform_expr(e);
                }
            }
            PureExpr::Ref { expr: inner, .. } => {
                changes += self.transform_expr(inner);
            }
            PureExpr::Return(Some(inner)) => {
                changes += self.transform_expr(inner);
            }
            PureExpr::Try(inner) => {
                changes += self.transform_expr(inner);
            }
            PureExpr::Await(inner) => {
                changes += self.transform_expr(inner);
            }
            _ => {}
        }

        changes
    }

    /// Transform 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 statement
    fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
        match stmt {
            PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
            PureStmt::Semi(e) | PureStmt::Expr(e) => self.transform_expr(e),
            _ => 0,
        }
    }
}

impl Mutation for UnwrapToQuestionMutation {
    fn describe(&self) -> String {
        "Convert .unwrap()/.expect() to ? operator".to_string()
    }

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

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

#[cfg(test)]
mod tests {
    use super::*;

    fn make_unwrap_expr() -> PureExpr {
        // x.unwrap()
        PureExpr::MethodCall {
            receiver: Box::new(PureExpr::Path("x".to_string())),
            method: "unwrap".to_string(),
            turbofish: None,
            args: vec![],
        }
    }

    fn make_expect_expr() -> PureExpr {
        // x.expect("error")
        PureExpr::MethodCall {
            receiver: Box::new(PureExpr::Path("x".to_string())),
            method: "expect".to_string(),
            turbofish: None,
            args: vec![PureExpr::Lit("\"error\"".to_string())],
        }
    }

    #[test]
    fn test_is_unwrap_call() {
        let mutation = UnwrapToQuestionMutation::new();
        assert!(mutation.is_unwrap_call("unwrap", &[]));
        assert!(mutation.is_unwrap_call("expect", &[PureExpr::Lit("\"msg\"".to_string())]));
        assert!(!mutation.is_unwrap_call("map", &[]));
    }

    #[test]
    fn test_transform_unwrap() {
        let mutation = UnwrapToQuestionMutation::new();
        let mut expr = make_unwrap_expr();
        let changes = mutation.transform_expr(&mut expr);

        assert_eq!(changes, 1);
        assert!(matches!(expr, PureExpr::Try(_)));
    }

    #[test]
    fn test_transform_expect() {
        let mutation = UnwrapToQuestionMutation::new();
        let mut expr = make_expect_expr();
        let changes = mutation.transform_expr(&mut expr);

        assert_eq!(changes, 1);
        assert!(matches!(expr, PureExpr::Try(_)));
    }

    #[test]
    fn test_skip_expect_when_unwrap_only() {
        let mutation = UnwrapToQuestionMutation::new().unwrap_only();
        let mut expr = make_expect_expr();
        let changes = mutation.transform_expr(&mut expr);

        assert_eq!(changes, 0);
        assert!(matches!(expr, PureExpr::MethodCall { .. }));
    }
}