Skip to main content

ryo_executor/engine/impls/
loop_to_iter.rs

1//! V2 ASTRegApply implementation for LoopToIteratorMutation
2//!
3//! Converts for loops to iterator chains:
4//! - `let mut v = Vec::new(); for x in iter { v.push(f(x)) }` → `let v = iter.map(f).collect()`
5//! - `for x in iter { if cond { v.push(x) } }` → `iter.filter(cond).collect()`
6
7use ryo_analysis::SymbolKind;
8use ryo_mutations::idiom::LoopToIteratorMutation;
9use ryo_mutations::{Mutation, MutationResult};
10use ryo_source::pure::{PureImplItem, PureItem};
11
12use crate::engine::{ASTMutationContext, ASTRegApply};
13
14impl ASTRegApply for LoopToIteratorMutation {
15    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
16        let mut total_changes = 0;
17
18        // Process standalone functions
19        let fn_ids: Vec<_> = ctx
20            .symbol_registry
21            .iter()
22            .filter(|(id, _)| matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Function)))
23            .map(|(id, _)| id)
24            .collect();
25
26        for id in fn_ids {
27            if let Some(PureItem::Fn(f)) = ctx.ast_registry.get_mut(id) {
28                total_changes += self.transform_block(&mut f.body);
29            }
30        }
31
32        // Process impl blocks (methods)
33        let impl_ids: Vec<_> = ctx
34            .symbol_registry
35            .iter()
36            .filter(|(id, _)| matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Impl)))
37            .map(|(id, _)| id)
38            .collect();
39
40        for id in impl_ids {
41            if let Some(PureItem::Impl(imp)) = ctx.ast_registry.get_mut(id) {
42                for impl_item in &mut imp.items {
43                    if let PureImplItem::Fn(f) = impl_item {
44                        total_changes += self.transform_block(&mut f.body);
45                    }
46                }
47            }
48        }
49
50        MutationResult {
51            mutation_type: self.mutation_type().to_string(),
52            changes: total_changes,
53            description: if total_changes > 0 {
54                format!("Converted {} for loop(s) to iterator chains", total_changes)
55            } else {
56                "No for loops converted".to_string()
57            },
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use crate::engine::ASTMutationEngine;
66    use ryo_analysis::testing::ContextBuilder;
67
68    #[test]
69    fn test_v2_loop_to_iter_map_collect() {
70        let mut ctx = ContextBuilder::new()
71            .with_file(
72                "src/lib.rs",
73                r#"
74fn process(items: Vec<i32>) -> Vec<i32> {
75    let mut result = Vec::new();
76    for x in items {
77        result.push(x * 2);
78    }
79    result
80}
81"#,
82            )
83            .build();
84
85        let mutation = LoopToIteratorMutation::new();
86        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
87
88        assert_eq!(result.result.changes, 1);
89    }
90
91    #[test]
92    fn test_v2_loop_to_iter_no_changes() {
93        let mut ctx = ContextBuilder::new()
94            .with_file(
95                "src/lib.rs",
96                r#"
97fn process(items: Vec<i32>) -> Vec<i32> {
98    items.into_iter().map(|x| x * 2).collect()
99}
100"#,
101            )
102            .build();
103
104        let mutation = LoopToIteratorMutation::new();
105        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
106
107        assert_eq!(result.result.changes, 0);
108    }
109}