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, ModificationType};
13
14impl ASTRegApply for LoopToIteratorMutation {
15    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
16        let mut total_changes = 0;
17
18        if let Some(target_id) = self.target_fn {
19            // Target function specified: standalone Function only.
20            // Methods inside Impl blocks are not handled in scoped mode
21            // (pre-existing limitation across all scoped idioms).
22            if let Some(PureItem::Fn(f)) = ctx.ast_registry.get_mut(target_id) {
23                let changes = self.transform_block(&mut f.body);
24                if changes > 0 {
25                    ctx.emit_modified(target_id, ModificationType::BodyModified);
26                    total_changes += changes;
27                }
28            }
29        } else {
30            // Process standalone functions
31            let fn_ids: Vec<_> = ctx
32                .symbol_registry
33                .iter()
34                .filter(|(id, _)| {
35                    matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Function))
36                })
37                .map(|(id, _)| id)
38                .collect();
39
40            for id in fn_ids {
41                if let Some(PureItem::Fn(f)) = ctx.ast_registry.get_mut(id) {
42                    let changes = self.transform_block(&mut f.body);
43                    if changes > 0 {
44                        ctx.emit_modified(id, ModificationType::BodyModified);
45                        total_changes += changes;
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                let mut impl_changes = 0;
60                if let Some(PureItem::Impl(imp)) = ctx.ast_registry.get_mut(id) {
61                    for impl_item in &mut imp.items {
62                        if let PureImplItem::Fn(f) = impl_item {
63                            impl_changes += self.transform_block(&mut f.body);
64                        }
65                    }
66                }
67                if impl_changes > 0 {
68                    ctx.emit_modified(id, ModificationType::BodyModified);
69                    total_changes += impl_changes;
70                }
71            }
72        }
73
74        MutationResult {
75            mutation_type: self.mutation_type().to_string(),
76            changes: total_changes,
77            description: if total_changes > 0 {
78                format!("Converted {} for loop(s) to iterator chains", total_changes)
79            } else {
80                "No for loops converted".to_string()
81            },
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::engine::ASTMutationEngine;
90    use ryo_analysis::testing::ContextBuilder;
91
92    #[test]
93    fn test_v2_loop_to_iter_target_fn_scopes_to_target_only() {
94        // Boundary regression: target_fn must restrict transform to ONE fn.
95        // Pre-fix the converter discarded module_id via `..`, leaking across
96        // every fn that contained a convertible for-loop.
97        let mut ctx = ContextBuilder::new()
98            .with_file(
99                "src/lib.rs",
100                r#"
101fn foo(items: Vec<i32>) -> Vec<i32> {
102    let mut result = Vec::new();
103    for x in items {
104        result.push(x * 2);
105    }
106    result
107}
108
109fn bar(items: Vec<i32>) -> Vec<i32> {
110    let mut result = Vec::new();
111    for x in items {
112        result.push(x + 1);
113    }
114    result
115}
116"#,
117            )
118            .build();
119
120        let foo_id = ctx
121            .registry
122            .iter()
123            .find(|(id, path)| {
124                matches!(ctx.registry.kind(*id), Some(SymbolKind::Function)) && path.name() == "foo"
125            })
126            .map(|(id, _)| id)
127            .expect("foo function not found");
128
129        let mutation = LoopToIteratorMutation::new().in_function(foo_id);
130        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
131
132        assert_eq!(result.result.changes, 1, "target_fn must scope to foo only");
133    }
134
135    #[test]
136    fn test_v2_loop_to_iter_map_collect() {
137        let mut ctx = ContextBuilder::new()
138            .with_file(
139                "src/lib.rs",
140                r#"
141fn process(items: Vec<i32>) -> Vec<i32> {
142    let mut result = Vec::new();
143    for x in items {
144        result.push(x * 2);
145    }
146    result
147}
148"#,
149            )
150            .build();
151
152        let mutation = LoopToIteratorMutation::new();
153        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
154
155        assert_eq!(result.result.changes, 1);
156    }
157
158    #[test]
159    fn test_v2_loop_to_iter_no_changes() {
160        let mut ctx = ContextBuilder::new()
161            .with_file(
162                "src/lib.rs",
163                r#"
164fn process(items: Vec<i32>) -> Vec<i32> {
165    items.into_iter().map(|x| x * 2).collect()
166}
167"#,
168            )
169            .build();
170
171        let mutation = LoopToIteratorMutation::new();
172        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
173
174        assert_eq!(result.result.changes, 0);
175    }
176}