ryo_mutations/idiom/
loop_to_iter.rs

1//! LoopToIteratorMutation: Convert for/while loops to iterator chains
2//!
3//! Patterns:
4//! - `let mut v = Vec::new(); for x in iter { v.push(f(x)) }` → `let v: Vec<_> = iter.into_iter().map(|x| f(x)).collect()`
5//! - `let mut v = Vec::new(); for x in iter { if cond { v.push(x) } }` → `let v: Vec<_> = iter.into_iter().filter(|x| cond).collect()`
6//! - `for x in 0..n { ... }` → `(0..n).for_each(|x| ...)`
7//!
8//! # Design
9//!
10//! This mutation analyzes blocks to detect the pattern:
11//! 1. `let mut v = Vec::new();` (or `vec![]`)
12//! 2. Followed by `for x in iter { v.push(...) }`
13//!
14//! Both statements are replaced with a single `let v = iter.into_iter()...collect()`.
15
16use ryo_source::pure::{PureBlock, PureClosureParam, PureExpr, PurePattern, PureStmt, PureType};
17
18use crate::Mutation;
19
20/// Pattern detected in a for loop
21#[derive(Debug, Clone, PartialEq)]
22pub enum LoopPattern {
23    /// `for x in iter { result.push(f(x)) }` → `iter.map(f).collect()`
24    MapCollect {
25        iter_expr: PureExpr,
26        var_name: String,
27        target_var: String,
28        transform: PureExpr,
29    },
30    /// `for x in iter { if cond { result.push(x) } }` → `iter.filter(cond).collect()`
31    FilterCollect {
32        iter_expr: PureExpr,
33        var_name: String,
34        target_var: String,
35        condition: PureExpr,
36    },
37    /// `for x in iter { if cond { result.push(f(x)) } }` → `iter.filter(cond).map(f).collect()`
38    FilterMapCollect {
39        iter_expr: PureExpr,
40        var_name: String,
41        target_var: String,
42        condition: PureExpr,
43        transform: PureExpr,
44    },
45    /// `for x in iter { body }` → `iter.for_each(|x| body)`
46    ForEach {
47        iter_expr: PureExpr,
48        var_name: String,
49        body: PureBlock,
50    },
51}
52
53/// Pattern detected in a block: `let mut v = Vec::new(); for x in iter { v.push(...) }`
54#[derive(Debug, Clone)]
55pub struct BlockLoopPattern {
56    /// Index of the `let mut v = Vec::new()` statement
57    pub vec_decl_idx: usize,
58    /// Index of the `for` loop statement
59    pub for_loop_idx: usize,
60    /// Name of the target variable (`v`)
61    pub target_var: String,
62    /// The loop pattern detected
63    pub loop_pattern: LoopPattern,
64}
65
66/// Convert for loops to iterator chains
67#[derive(Debug, Clone, Default)]
68pub struct LoopToIteratorMutation {
69    /// Only convert loops with this iterator variable name
70    pub target_var: Option<String>,
71    /// Convert to for_each even if no clear pattern
72    pub aggressive: bool,
73}
74
75impl LoopToIteratorMutation {
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    /// Only convert loops using a specific variable
81    pub fn with_target(mut self, var: impl Into<String>) -> Self {
82        self.target_var = Some(var.into());
83        self
84    }
85
86    /// Aggressively convert to for_each when no clear pattern
87    pub fn aggressive(mut self) -> Self {
88        self.aggressive = true;
89        self
90    }
91
92    /// Analyze a for loop and detect the pattern
93    fn detect_pattern(
94        var_pattern: &PurePattern,
95        iter_expr: &PureExpr,
96        body: &PureBlock,
97    ) -> Option<LoopPattern> {
98        let var_name = match var_pattern {
99            PurePattern::Ident { name, .. } => name.clone(),
100            _ => return None, // Complex patterns not supported yet
101        };
102
103        // Check for single-statement body
104        if body.stmts.len() == 1 {
105            if let Some(pattern) =
106                Self::detect_single_stmt_pattern(&var_name, iter_expr, &body.stmts[0])
107            {
108                return Some(pattern);
109            }
110        }
111
112        // Check for filter pattern: if cond { push }
113        if body.stmts.len() == 1 {
114            if let PureStmt::Semi(PureExpr::If {
115                cond,
116                then_branch,
117                else_branch: None,
118            })
119            | PureStmt::Expr(PureExpr::If {
120                cond,
121                then_branch,
122                else_branch: None,
123            }) = &body.stmts[0]
124            {
125                if then_branch.stmts.len() == 1 {
126                    if let Some((target_var, pushed_expr)) =
127                        Self::extract_push(&then_branch.stmts[0])
128                    {
129                        // Check if it's filter (push var) or filter-map (push f(var))
130                        if Self::is_simple_var(&pushed_expr, &var_name) {
131                            return Some(LoopPattern::FilterCollect {
132                                iter_expr: iter_expr.clone(),
133                                var_name,
134                                target_var,
135                                condition: *cond.clone(),
136                            });
137                        } else {
138                            return Some(LoopPattern::FilterMapCollect {
139                                iter_expr: iter_expr.clone(),
140                                var_name,
141                                target_var,
142                                condition: *cond.clone(),
143                                transform: pushed_expr,
144                            });
145                        }
146                    }
147                }
148            }
149        }
150
151        // Default: for_each pattern
152        Some(LoopPattern::ForEach {
153            iter_expr: iter_expr.clone(),
154            var_name,
155            body: body.clone(),
156        })
157    }
158
159    /// Detect pattern from a single statement
160    fn detect_single_stmt_pattern(
161        var_name: &str,
162        iter_expr: &PureExpr,
163        stmt: &PureStmt,
164    ) -> Option<LoopPattern> {
165        // Check for vec.push(expr) pattern
166        if let Some((target_var, pushed_expr)) = Self::extract_push(stmt) {
167            if Self::is_simple_var(&pushed_expr, var_name) {
168                // Simple collect: for x in iter { v.push(x) }
169                return Some(LoopPattern::MapCollect {
170                    iter_expr: iter_expr.clone(),
171                    var_name: var_name.to_string(),
172                    target_var,
173                    transform: pushed_expr,
174                });
175            } else {
176                // Map collect: for x in iter { v.push(f(x)) }
177                return Some(LoopPattern::MapCollect {
178                    iter_expr: iter_expr.clone(),
179                    var_name: var_name.to_string(),
180                    target_var,
181                    transform: pushed_expr,
182                });
183            }
184        }
185        None
186    }
187
188    /// Extract target and expression from a push statement
189    fn extract_push(stmt: &PureStmt) -> Option<(String, PureExpr)> {
190        let expr = match stmt {
191            PureStmt::Semi(e) | PureStmt::Expr(e) => e,
192            _ => return None,
193        };
194
195        // Look for method call: target.push(arg)
196        if let PureExpr::MethodCall {
197            receiver,
198            method,
199            args,
200            ..
201        } = expr
202        {
203            if method == "push" && args.len() == 1 {
204                if let PureExpr::Path(target_var) = receiver.as_ref() {
205                    return Some((target_var.clone(), args[0].clone()));
206                }
207            }
208        }
209        None
210    }
211
212    /// Check if expression is just a simple variable reference
213    fn is_simple_var(expr: &PureExpr, var_name: &str) -> bool {
214        matches!(expr, PureExpr::Path(name) if name == var_name)
215    }
216
217    /// Convert a detected pattern to iterator expression
218    fn pattern_to_iter_expr(pattern: &LoopPattern) -> PureExpr {
219        match pattern {
220            LoopPattern::MapCollect {
221                iter_expr,
222                var_name,
223                transform,
224                ..
225            } => {
226                let is_identity = Self::is_simple_var(transform, var_name);
227
228                if is_identity {
229                    // iter.collect()
230                    PureExpr::MethodCall {
231                        receiver: Box::new(PureExpr::MethodCall {
232                            receiver: Box::new(iter_expr.clone()),
233                            method: "into_iter".to_string(),
234                            turbofish: None,
235                            args: vec![],
236                        }),
237                        method: "collect".to_string(),
238                        turbofish: None,
239                        args: vec![],
240                    }
241                } else {
242                    // iter.map(|var| transform).collect()
243                    PureExpr::MethodCall {
244                        receiver: Box::new(PureExpr::MethodCall {
245                            receiver: Box::new(PureExpr::MethodCall {
246                                receiver: Box::new(iter_expr.clone()),
247                                method: "into_iter".to_string(),
248                                turbofish: None,
249                                args: vec![],
250                            }),
251                            method: "map".to_string(),
252                            turbofish: None,
253                            args: vec![PureExpr::Closure {
254                                is_async: false,
255                                is_move: false,
256                                params: vec![PureClosureParam::untyped(PurePattern::Ident {
257                                    name: var_name.clone(),
258                                    is_mut: false,
259                                })],
260                                ret: None,
261                                body: Box::new(transform.clone()),
262                            }],
263                        }),
264                        method: "collect".to_string(),
265                        turbofish: None,
266                        args: vec![],
267                    }
268                }
269            }
270            LoopPattern::FilterCollect {
271                iter_expr,
272                var_name,
273                condition,
274                ..
275            } => {
276                // iter.filter(|&var| cond).collect()
277                // Note: Use reference pattern |&var| for filter closures since filter
278                // passes items by reference, but the original condition uses the value.
279                PureExpr::MethodCall {
280                    receiver: Box::new(PureExpr::MethodCall {
281                        receiver: Box::new(PureExpr::MethodCall {
282                            receiver: Box::new(iter_expr.clone()),
283                            method: "into_iter".to_string(),
284                            turbofish: None,
285                            args: vec![],
286                        }),
287                        method: "filter".to_string(),
288                        turbofish: None,
289                        args: vec![PureExpr::Closure {
290                            is_async: false,
291                            is_move: false,
292                            params: vec![PureClosureParam::untyped(PurePattern::Ref {
293                                is_mut: false,
294                                pattern: Box::new(PurePattern::Ident {
295                                    name: var_name.clone(),
296                                    is_mut: false,
297                                }),
298                            })],
299                            ret: None,
300                            body: Box::new(condition.clone()),
301                        }],
302                    }),
303                    method: "collect".to_string(),
304                    turbofish: None,
305                    args: vec![],
306                }
307            }
308            LoopPattern::FilterMapCollect {
309                iter_expr,
310                var_name,
311                condition,
312                transform,
313                ..
314            } => {
315                // iter.filter(|&var| cond).map(|var| transform).collect()
316                // Note: Use reference pattern |&var| for filter closures since filter
317                // passes items by reference.
318                PureExpr::MethodCall {
319                    receiver: Box::new(PureExpr::MethodCall {
320                        receiver: Box::new(PureExpr::MethodCall {
321                            receiver: Box::new(PureExpr::MethodCall {
322                                receiver: Box::new(iter_expr.clone()),
323                                method: "into_iter".to_string(),
324                                turbofish: None,
325                                args: vec![],
326                            }),
327                            method: "filter".to_string(),
328                            turbofish: None,
329                            args: vec![PureExpr::Closure {
330                                is_async: false,
331                                is_move: false,
332                                params: vec![PureClosureParam::untyped(PurePattern::Ref {
333                                    is_mut: false,
334                                    pattern: Box::new(PurePattern::Ident {
335                                        name: var_name.clone(),
336                                        is_mut: false,
337                                    }),
338                                })],
339                                ret: None,
340                                body: Box::new(condition.clone()),
341                            }],
342                        }),
343                        method: "map".to_string(),
344                        turbofish: None,
345                        args: vec![PureExpr::Closure {
346                            is_async: false,
347                            is_move: false,
348                            params: vec![PureClosureParam::untyped(PurePattern::Ident {
349                                name: var_name.clone(),
350                                is_mut: false,
351                            })],
352                            ret: None,
353                            body: Box::new(transform.clone()),
354                        }],
355                    }),
356                    method: "collect".to_string(),
357                    turbofish: None,
358                    args: vec![],
359                }
360            }
361            LoopPattern::ForEach {
362                iter_expr,
363                var_name,
364                body,
365            } => {
366                // iter.for_each(|var| { body })
367                PureExpr::MethodCall {
368                    receiver: Box::new(PureExpr::MethodCall {
369                        receiver: Box::new(iter_expr.clone()),
370                        method: "into_iter".to_string(),
371                        turbofish: None,
372                        args: vec![],
373                    }),
374                    method: "for_each".to_string(),
375                    turbofish: None,
376                    args: vec![PureExpr::Closure {
377                        is_async: false,
378                        is_move: false,
379                        params: vec![PureClosureParam::untyped(PurePattern::Ident {
380                            name: var_name.clone(),
381                            is_mut: false,
382                        })],
383                        ret: None,
384                        body: Box::new(PureExpr::Block {
385                            label: None,
386                            block: body.clone(),
387                        }),
388                    }],
389                }
390            }
391        }
392    }
393
394    /// Detect block-level patterns: `let mut v = Vec::new(); for x in iter { v.push(...) }`
395    fn detect_block_patterns(stmts: &[PureStmt]) -> Vec<BlockLoopPattern> {
396        let mut patterns = Vec::new();
397
398        for (i, stmt) in stmts.iter().enumerate() {
399            // Look for `let mut v = Vec::new()` or `let mut v = vec![]`
400            if let Some((var_name, is_vec_init)) = Self::extract_vec_init(stmt) {
401                if !is_vec_init {
402                    continue;
403                }
404
405                // Look for a following for loop that pushes to this variable
406                for (j, jstmt) in stmts.iter().enumerate().skip(i + 1) {
407                    if let Some(loop_pattern) = Self::extract_for_loop_push(jstmt, &var_name) {
408                        patterns.push(BlockLoopPattern {
409                            vec_decl_idx: i,
410                            for_loop_idx: j,
411                            target_var: var_name.clone(),
412                            loop_pattern,
413                        });
414                        break; // Only match the first for loop
415                    }
416
417                    // If we hit a statement that uses `var_name` in a different way, stop
418                    if Self::stmt_uses_var(jstmt, &var_name) {
419                        break;
420                    }
421                }
422            }
423        }
424
425        patterns
426    }
427
428    /// Extract variable name and check if it's `Vec::new()` or `vec![]`
429    fn extract_vec_init(stmt: &PureStmt) -> Option<(String, bool)> {
430        if let PureStmt::Local {
431            pattern: PurePattern::Ident { name, is_mut: true },
432            init: Some(init_expr),
433            ..
434        } = stmt
435        {
436            let is_vec_init = Self::is_vec_new_call(init_expr) || Self::is_vec_macro(init_expr);
437            return Some((name.clone(), is_vec_init));
438        }
439        None
440    }
441
442    /// Check if expression is `Vec::new()`
443    fn is_vec_new_call(expr: &PureExpr) -> bool {
444        if let PureExpr::Call { func, args } = expr {
445            if args.is_empty() {
446                // Check for Vec::new() pattern
447                if let PureExpr::Path(path) = func.as_ref() {
448                    return path == "Vec::new" || path.ends_with("::Vec::new");
449                }
450            }
451        }
452        false
453    }
454
455    /// Check if expression is `vec![]`
456    fn is_vec_macro(expr: &PureExpr) -> bool {
457        if let PureExpr::Macro { name, .. } = expr {
458            return name == "vec" || name.ends_with("::vec");
459        }
460        false
461    }
462
463    /// Extract for loop pattern if it pushes to the target variable
464    fn extract_for_loop_push(stmt: &PureStmt, target_var: &str) -> Option<LoopPattern> {
465        let for_expr = match stmt {
466            PureStmt::Semi(e) | PureStmt::Expr(e) => e,
467            _ => return None,
468        };
469
470        if let PureExpr::For {
471            pat,
472            expr: iter_expr,
473            body,
474            ..
475        } = for_expr
476        {
477            let pattern = Self::detect_pattern(pat, iter_expr, body)?;
478
479            // Check if the pattern targets our variable
480            let pattern_target = match &pattern {
481                LoopPattern::MapCollect { target_var, .. } => target_var,
482                LoopPattern::FilterCollect { target_var, .. } => target_var,
483                LoopPattern::FilterMapCollect { target_var, .. } => target_var,
484                LoopPattern::ForEach { .. } => return None, // ForEach doesn't push
485            };
486
487            if pattern_target == target_var {
488                return Some(pattern);
489            }
490        }
491        None
492    }
493
494    /// Check if a statement uses the given variable (other than in a for loop push)
495    fn stmt_uses_var(stmt: &PureStmt, var_name: &str) -> bool {
496        match stmt {
497            PureStmt::Semi(expr) | PureStmt::Expr(expr) => Self::expr_uses_var(expr, var_name),
498            PureStmt::Local { init: Some(e), .. } => Self::expr_uses_var(e, var_name),
499            _ => false,
500        }
501    }
502
503    /// Check if an expression uses the given variable
504    fn expr_uses_var(expr: &PureExpr, var_name: &str) -> bool {
505        match expr {
506            PureExpr::Path(name) => name == var_name,
507            PureExpr::MethodCall { receiver, args, .. } => {
508                Self::expr_uses_var(receiver, var_name)
509                    || args.iter().any(|a| Self::expr_uses_var(a, var_name))
510            }
511            PureExpr::Call { func, args } => {
512                Self::expr_uses_var(func, var_name)
513                    || args.iter().any(|a| Self::expr_uses_var(a, var_name))
514            }
515            PureExpr::Binary { left, right, .. } => {
516                Self::expr_uses_var(left, var_name) || Self::expr_uses_var(right, var_name)
517            }
518            PureExpr::If {
519                cond,
520                then_branch,
521                else_branch,
522            } => {
523                Self::expr_uses_var(cond, var_name)
524                    || then_branch
525                        .stmts
526                        .iter()
527                        .any(|s| Self::stmt_uses_var(s, var_name))
528                    || else_branch
529                        .as_ref()
530                        .map(|e| Self::expr_uses_var(e, var_name))
531                        .unwrap_or(false)
532            }
533            PureExpr::For { .. } => false, // Don't recurse into for loops (they're handled separately)
534            _ => false,
535        }
536    }
537
538    /// Convert a block pattern to a let statement
539    fn block_pattern_to_let_stmt(pattern: &BlockLoopPattern) -> PureStmt {
540        let iter_expr = Self::pattern_to_iter_expr(&pattern.loop_pattern);
541
542        PureStmt::Local {
543            pattern: PurePattern::Ident {
544                name: pattern.target_var.clone(),
545                is_mut: false, // No longer mutable!
546            },
547            ty: Some(PureType::Other("Vec<_>".to_string())), // Type annotation for collect()
548            init: Some(iter_expr),
549        }
550    }
551
552    /// Transform for loops in a block (with block-level pattern detection)
553    pub fn transform_block(&self, block: &mut PureBlock) -> usize {
554        let mut changes = 0;
555
556        // First, detect and apply block-level patterns
557        let block_patterns = Self::detect_block_patterns(&block.stmts);
558
559        if !block_patterns.is_empty() {
560            // Apply patterns in reverse order to preserve indices
561            let mut indices_to_remove = Vec::new();
562
563            for pattern in block_patterns.iter().rev() {
564                // Replace the vec declaration with the iterator expression
565                let new_stmt = Self::block_pattern_to_let_stmt(pattern);
566                block.stmts[pattern.vec_decl_idx] = new_stmt;
567
568                // Mark the for loop for removal
569                indices_to_remove.push(pattern.for_loop_idx);
570                changes += 1;
571            }
572
573            // Remove the for loops (in reverse order to preserve indices)
574            indices_to_remove.sort();
575            indices_to_remove.reverse();
576            for idx in indices_to_remove {
577                block.stmts.remove(idx);
578            }
579        }
580
581        // Then, recursively transform nested blocks
582        for stmt in &mut block.stmts {
583            changes += self.transform_stmt(stmt);
584        }
585
586        changes
587    }
588
589    /// Transform a statement
590    fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
591        match stmt {
592            PureStmt::Semi(expr) | PureStmt::Expr(expr) => self.transform_expr(expr),
593            PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
594            _ => 0,
595        }
596    }
597
598    /// Transform an expression (recursively)
599    fn transform_expr(&self, expr: &mut PureExpr) -> usize {
600        let mut changes = 0;
601
602        match expr {
603            PureExpr::For {
604                pat,
605                expr: iter_expr,
606                body,
607                ..
608            } => {
609                // Check if we should transform this loop
610                if let Some(ref target) = self.target_var {
611                    if let PurePattern::Ident { name, .. } = pat {
612                        if name != target {
613                            // Not our target, still recurse into body
614                            changes += self.transform_block(body);
615                            return changes;
616                        }
617                    }
618                }
619
620                // Standalone for loops: only ForEach pattern with aggressive mode
621                if self.aggressive {
622                    if let Some(pattern) = Self::detect_pattern(pat, iter_expr, body) {
623                        if matches!(pattern, LoopPattern::ForEach { .. }) {
624                            *expr = Self::pattern_to_iter_expr(&pattern);
625                            changes += 1;
626                        }
627                    }
628                } else {
629                    // Recurse into body
630                    changes += self.transform_block(body);
631                }
632            }
633            PureExpr::Block { block, .. } => {
634                changes += self.transform_block(block);
635            }
636            PureExpr::If {
637                cond,
638                then_branch,
639                else_branch,
640            } => {
641                changes += self.transform_expr(cond);
642                changes += self.transform_block(then_branch);
643                if let Some(else_expr) = else_branch {
644                    changes += self.transform_expr(else_expr);
645                }
646            }
647            PureExpr::Match { expr: e, arms } => {
648                changes += self.transform_expr(e);
649                for arm in arms {
650                    changes += self.transform_expr(&mut arm.body);
651                }
652            }
653            PureExpr::Loop { body: block, .. } | PureExpr::While { body: block, .. } => {
654                changes += self.transform_block(block);
655            }
656            PureExpr::Closure { body, .. } => {
657                changes += self.transform_expr(body);
658            }
659            _ => {}
660        }
661
662        changes
663    }
664}
665
666impl Mutation for LoopToIteratorMutation {
667    fn describe(&self) -> String {
668        "Convert for loops to iterator chains".to_string()
669    }
670
671    fn mutation_type(&self) -> &'static str {
672        "LoopToIterator"
673    }
674
675    fn box_clone(&self) -> Box<dyn Mutation> {
676        Box::new(self.clone())
677    }
678}
679
680#[cfg(test)]
681mod tests {
682    use super::*;
683
684    #[test]
685    fn test_detect_map_pattern() {
686        let pat = PurePattern::Ident {
687            name: "x".to_string(),
688            is_mut: false,
689        };
690        let iter = PureExpr::Path("items".to_string());
691        let body = PureBlock {
692            stmts: vec![PureStmt::Semi(PureExpr::MethodCall {
693                receiver: Box::new(PureExpr::Path("result".to_string())),
694                method: "push".to_string(),
695                turbofish: None,
696                args: vec![PureExpr::Binary {
697                    op: "*".to_string(),
698                    left: Box::new(PureExpr::Path("x".to_string())),
699                    right: Box::new(PureExpr::Lit("2".to_string())),
700                }],
701            })],
702        };
703
704        let pattern = LoopToIteratorMutation::detect_pattern(&pat, &iter, &body);
705        assert!(matches!(pattern, Some(LoopPattern::MapCollect { .. })));
706    }
707
708    #[test]
709    fn test_detect_filter_pattern() {
710        let pat = PurePattern::Ident {
711            name: "x".to_string(),
712            is_mut: false,
713        };
714        let iter = PureExpr::Path("items".to_string());
715        let body = PureBlock {
716            stmts: vec![PureStmt::Semi(PureExpr::If {
717                cond: Box::new(PureExpr::Binary {
718                    op: ">".to_string(),
719                    left: Box::new(PureExpr::Path("x".to_string())),
720                    right: Box::new(PureExpr::Lit("0".to_string())),
721                }),
722                then_branch: PureBlock {
723                    stmts: vec![PureStmt::Semi(PureExpr::MethodCall {
724                        receiver: Box::new(PureExpr::Path("result".to_string())),
725                        method: "push".to_string(),
726                        turbofish: None,
727                        args: vec![PureExpr::Path("x".to_string())],
728                    })],
729                },
730                else_branch: None,
731            })],
732        };
733
734        let pattern = LoopToIteratorMutation::detect_pattern(&pat, &iter, &body);
735        assert!(matches!(pattern, Some(LoopPattern::FilterCollect { .. })));
736    }
737}
ryo_mutations/idiom/loop_to_iter.rs

ryo_mutations/idiom/
loop_to_iter.rs
