ryo_mutations/idiom/

map_unwrap_or.rs

//! MapUnwrapOrMutation: Convert `.map().unwrap_or()` to `.map_or()`
//!
//! Transforms:
//! - `opt.map(f).unwrap_or(default)` -> `opt.map_or(default, f)`
//! - `opt.map(f).unwrap_or_else(g)` -> `opt.map_or_else(g, f)`
//!
//! Corresponds to Clippy lints: `clippy::map_unwrap_or`, `clippy::option_map_or_none`

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

use crate::Mutation;

/// Convert `.map().unwrap_or()` chains to `.map_or()`
///
/// # Example
///
/// ```rust,ignore
/// use ryo_mutations::idiom::MapUnwrapOrMutation;
///
/// let mutation = MapUnwrapOrMutation::new();
/// // Transforms: opt.map(|x| x + 1).unwrap_or(0)
/// // Into:       opt.map_or(0, |x| x + 1)
/// ```
#[derive(Debug, Clone, Default)]
pub struct MapUnwrapOrMutation {
    /// Target function SymbolId. If None, applies to all functions.
    pub target_fn: Option<SymbolId>,
}

impl MapUnwrapOrMutation {
    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
    }

    /// Check if expression is a `.map(f).unwrap_or(default)` pattern
    /// Returns (receiver, map_fn, default_value, is_lazy) if matched
    fn is_map_unwrap_or_pattern(expr: &PureExpr) -> Option<(PureExpr, PureExpr, PureExpr, bool)> {
        if let PureExpr::MethodCall {
            receiver,
            method,
            args,
            ..
        } = expr
        {
            // Check for .unwrap_or(default) or .unwrap_or_else(f)
            let (is_lazy, default_arg) = if method == "unwrap_or" && args.len() == 1 {
                (false, &args[0])
            } else if method == "unwrap_or_else" && args.len() == 1 {
                (true, &args[0])
            } else {
                return None;
            };

            // Check receiver is .map(f)
            if let PureExpr::MethodCall {
                receiver: map_receiver,
                method: map_method,
                args: map_args,
                ..
            } = receiver.as_ref()
            {
                if map_method == "map" && map_args.len() == 1 {
                    return Some((
                        map_receiver.as_ref().clone(),
                        map_args[0].clone(),
                        default_arg.clone(),
                        is_lazy,
                    ));
                }
            }
        }
        None
    }

    /// Check if expression is a `.map(f).unwrap_or(None)` pattern (-> and_then)
    /// Returns (receiver, map_fn) if matched
    fn is_map_or_none_pattern(expr: &PureExpr) -> Option<(PureExpr, PureExpr)> {
        if let PureExpr::MethodCall {
            receiver,
            method,
            args,
            ..
        } = expr
        {
            // Check for .unwrap_or(None)
            if method == "unwrap_or"
                && args.len() == 1
                && matches!(&args[0], PureExpr::Path(p) if p == "None")
            {
                // Check receiver is .map(f)
                if let PureExpr::MethodCall {
                    receiver: map_receiver,
                    method: map_method,
                    args: map_args,
                    ..
                } = receiver.as_ref()
                {
                    if map_method == "map" && map_args.len() == 1 {
                        return Some((map_receiver.as_ref().clone(), map_args[0].clone()));
                    }
                }
            }
        }
        None
    }

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

        // Check for map().unwrap_or(None) -> and_then pattern first
        if let Some((receiver, map_fn)) = Self::is_map_or_none_pattern(expr) {
            *expr = PureExpr::MethodCall {
                receiver: Box::new(receiver),
                method: "and_then".to_string(),
                turbofish: None,
                args: vec![map_fn],
            };
            return 1;
        }

        // Check for map().unwrap_or() pattern
        if let Some((receiver, map_fn, default_value, is_lazy)) =
            Self::is_map_unwrap_or_pattern(expr)
        {
            let method = if is_lazy { "map_or_else" } else { "map_or" };
            *expr = PureExpr::MethodCall {
                receiver: Box::new(receiver),
                method: method.to_string(),
                turbofish: None,
                // Note: map_or takes (default, f), map_or_else takes (default_fn, f)
                args: vec![default_value, map_fn],
            };
            return 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) | PureExpr::Await(inner) => {
                changes += self.transform_expr(inner);
            }
            _ => {}
        }

        changes
    }

    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::Local { init: Some(e), .. } => self.transform_expr(e),
            PureStmt::Semi(e) | PureStmt::Expr(e) => self.transform_expr(e),
            _ => 0,
        }
    }
}

impl Mutation for MapUnwrapOrMutation {
    fn describe(&self) -> String {
        "Convert .map().unwrap_or() to .map_or()".to_string()
    }

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

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

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

    #[test]
    fn test_is_map_unwrap_or_pattern() {
        // opt.map(f).unwrap_or(default)
        let expr = PureExpr::MethodCall {
            receiver: Box::new(PureExpr::MethodCall {
                receiver: Box::new(PureExpr::Path("opt".to_string())),
                method: "map".to_string(),
                turbofish: None,
                args: vec![PureExpr::Path("f".to_string())],
            }),
            method: "unwrap_or".to_string(),
            turbofish: None,
            args: vec![PureExpr::Lit("0".to_string())],
        };

        let result = MapUnwrapOrMutation::is_map_unwrap_or_pattern(&expr);
        assert!(result.is_some());
        let (_, _, _, is_lazy) = result.unwrap();
        assert!(!is_lazy);
    }

    #[test]
    fn test_is_map_unwrap_or_else_pattern() {
        // opt.map(f).unwrap_or_else(g)
        let expr = PureExpr::MethodCall {
            receiver: Box::new(PureExpr::MethodCall {
                receiver: Box::new(PureExpr::Path("opt".to_string())),
                method: "map".to_string(),
                turbofish: None,
                args: vec![PureExpr::Path("f".to_string())],
            }),
            method: "unwrap_or_else".to_string(),
            turbofish: None,
            args: vec![PureExpr::Path("g".to_string())],
        };

        let result = MapUnwrapOrMutation::is_map_unwrap_or_pattern(&expr);
        assert!(result.is_some());
        let (_, _, _, is_lazy) = result.unwrap();
        assert!(is_lazy);
    }

    #[test]
    fn test_is_map_or_none_pattern() {
        // opt.map(f).unwrap_or(None)
        let expr = PureExpr::MethodCall {
            receiver: Box::new(PureExpr::MethodCall {
                receiver: Box::new(PureExpr::Path("opt".to_string())),
                method: "map".to_string(),
                turbofish: None,
                args: vec![PureExpr::Path("f".to_string())],
            }),
            method: "unwrap_or".to_string(),
            turbofish: None,
            args: vec![PureExpr::Path("None".to_string())],
        };

        assert!(MapUnwrapOrMutation::is_map_or_none_pattern(&expr).is_some());
    }
}