Skip to main content

ryo_executor/engine/impls/
unwrap_to_question.rs

1//! V2 ASTRegApply implementation for UnwrapToQuestionMutation
2//!
3//! Converts .unwrap()/.expect() to ? operator:
4//! - `x.unwrap()` → `x?`
5//! - `x.expect("msg")` → `x?`
6//!
7//! Only applies in functions that return Result/Option.
8
9use ryo_analysis::SymbolKind;
10use ryo_mutations::idiom::UnwrapToQuestionMutation;
11use ryo_mutations::{Mutation, MutationResult};
12use ryo_source::pure::{PureImplItem, PureItem, PureType};
13
14use crate::engine::{ASTMutationContext, ASTRegApply};
15
16/// Check if a return type is Option or Result (compatible with ? operator)
17fn returns_option_or_result(ret: &Option<PureType>) -> bool {
18    match ret {
19        Some(PureType::Path(path)) => {
20            path.starts_with("Option<")
21                || path.starts_with("Result<")
22                || path == "Option"
23                || path == "Result"
24        }
25        _ => false,
26    }
27}
28
29impl ASTRegApply for UnwrapToQuestionMutation {
30    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
31        let mut total_changes = 0;
32
33        // Process standalone functions
34        let fn_ids: Vec<_> = ctx
35            .symbol_registry
36            .iter()
37            .filter(|(id, _)| matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Function)))
38            .map(|(id, _)| id)
39            .collect();
40
41        for id in fn_ids {
42            if let Some(PureItem::Fn(f)) = ctx.ast_registry.get_mut(id) {
43                // Only transform if function returns Option or Result
44                if returns_option_or_result(&f.ret) {
45                    total_changes += self.transform_block(&mut f.body);
46                }
47            }
48        }
49
50        // Process impl blocks (methods)
51        let impl_ids: Vec<_> = ctx
52            .symbol_registry
53            .iter()
54            .filter(|(id, _)| matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Impl)))
55            .map(|(id, _)| id)
56            .collect();
57
58        for id in impl_ids {
59            if let Some(PureItem::Impl(imp)) = ctx.ast_registry.get_mut(id) {
60                for impl_item in &mut imp.items {
61                    if let PureImplItem::Fn(f) = impl_item {
62                        // Only transform if method returns Option or Result
63                        if returns_option_or_result(&f.ret) {
64                            total_changes += self.transform_block(&mut f.body);
65                        }
66                    }
67                }
68            }
69        }
70
71        MutationResult {
72            mutation_type: self.mutation_type().to_string(),
73            changes: total_changes,
74            description: if total_changes > 0 {
75                format!("Converted {} .unwrap()/.expect() to ?", total_changes)
76            } else {
77                "No unwrap calls converted".to_string()
78            },
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::engine::ASTMutationEngine;
87    use ryo_analysis::testing::ContextBuilder;
88
89    #[test]
90    fn test_v2_unwrap_to_question_basic() {
91        let mut ctx = ContextBuilder::new()
92            .with_file(
93                "src/lib.rs",
94                r#"
95fn process(opt: Option<i32>) -> Option<i32> {
96    let x = opt.unwrap();
97    Some(x + 1)
98}
99"#,
100            )
101            .build();
102
103        let mutation = UnwrapToQuestionMutation::new();
104        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
105
106        assert_eq!(result.result.changes, 1);
107    }
108
109    #[test]
110    fn test_v2_unwrap_to_question_non_option_return() {
111        let mut ctx = ContextBuilder::new()
112            .with_file(
113                "src/lib.rs",
114                r#"
115fn process(opt: Option<i32>) -> i32 {
116    opt.unwrap()
117}
118"#,
119            )
120            .build();
121
122        let mutation = UnwrapToQuestionMutation::new();
123        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
124
125        // Should not convert because function doesn't return Option/Result
126        assert_eq!(result.result.changes, 0);
127    }
128
129    #[test]
130    fn test_v2_unwrap_to_question_expect() {
131        let mut ctx = ContextBuilder::new()
132            .with_file(
133                "src/lib.rs",
134                r#"
135fn process(opt: Option<i32>) -> Option<i32> {
136    let x = opt.expect("should not be none");
137    Some(x + 1)
138}
139"#,
140            )
141            .build();
142
143        let mutation = UnwrapToQuestionMutation::new();
144        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
145
146        assert_eq!(result.result.changes, 1);
147    }
148}