ryo-mutations 0.1.0

//! LoopToIteratorMutation: Convert for/while loops to iterator chains
//!
//! Patterns:
//! - `let mut v = Vec::new(); for x in iter { v.push(f(x)) }` → `let v: Vec<_> = iter.into_iter().map(|x| f(x)).collect()`
//! - `let mut v = Vec::new(); for x in iter { if cond { v.push(x) } }` → `let v: Vec<_> = iter.into_iter().filter(|x| cond).collect()`
//! - `for x in 0..n { ... }` → `(0..n).for_each(|x| ...)`
//!
//! # Design
//!
//! This mutation analyzes blocks to detect the pattern:
//! 1. `let mut v = Vec::new();` (or `vec![]`)
//! 2. Followed by `for x in iter { v.push(...) }`
//!
//! Both statements are replaced with a single `let v = iter.into_iter()...collect()`.

use ryo_source::pure::{PureBlock, PureClosureParam, PureExpr, PurePattern, PureStmt, PureType};

use crate::Mutation;

/// Pattern detected in a for loop
#[derive(Debug, Clone, PartialEq)]
pub enum LoopPattern {
    /// `for x in iter { result.push(f(x)) }` → `iter.map(f).collect()`
    MapCollect {
        iter_expr: PureExpr,
        var_name: String,
        target_var: String,
        transform: PureExpr,
    },
    /// `for x in iter { if cond { result.push(x) } }` → `iter.filter(cond).collect()`
    FilterCollect {
        iter_expr: PureExpr,
        var_name: String,
        target_var: String,
        condition: PureExpr,
    },
    /// `for x in iter { if cond { result.push(f(x)) } }` → `iter.filter(cond).map(f).collect()`
    FilterMapCollect {
        iter_expr: PureExpr,
        var_name: String,
        target_var: String,
        condition: PureExpr,
        transform: PureExpr,
    },
    /// `for x in iter { body }` → `iter.for_each(|x| body)`
    ForEach {
        iter_expr: PureExpr,
        var_name: String,
        body: PureBlock,
    },
}

/// Pattern detected in a block: `let mut v = Vec::new(); for x in iter { v.push(...) }`
#[derive(Debug, Clone)]
pub struct BlockLoopPattern {
    /// Index of the `let mut v = Vec::new()` statement
    pub vec_decl_idx: usize,
    /// Index of the `for` loop statement
    pub for_loop_idx: usize,
    /// Name of the target variable (`v`)
    pub target_var: String,
    /// The loop pattern detected
    pub loop_pattern: LoopPattern,
}

/// Convert for loops to iterator chains
#[derive(Debug, Clone, Default)]
pub struct LoopToIteratorMutation {
    /// Only convert loops with this iterator variable name
    pub target_var: Option<String>,
    /// Convert to for_each even if no clear pattern
    pub aggressive: bool,
}

impl LoopToIteratorMutation {
    pub fn new() -> Self {
        Self::default()
    }

    /// Only convert loops using a specific variable
    pub fn with_target(mut self, var: impl Into<String>) -> Self {
        self.target_var = Some(var.into());
        self
    }

    /// Aggressively convert to for_each when no clear pattern
    pub fn aggressive(mut self) -> Self {
        self.aggressive = true;
        self
    }

    /// Analyze a for loop and detect the pattern
    fn detect_pattern(
        var_pattern: &PurePattern,
        iter_expr: &PureExpr,
        body: &PureBlock,
    ) -> Option<LoopPattern> {
        let var_name = match var_pattern {
            PurePattern::Ident { name, .. } => name.clone(),
            _ => return None, // Complex patterns not supported yet
        };

        // Check for single-statement body
        if body.stmts.len() == 1 {
            if let Some(pattern) =
                Self::detect_single_stmt_pattern(&var_name, iter_expr, &body.stmts[0])
            {
                return Some(pattern);
            }
        }

        // Check for filter pattern: if cond { push }
        if body.stmts.len() == 1 {
            if let PureStmt::Semi(PureExpr::If {
                cond,
                then_branch,
                else_branch: None,
            })
            | PureStmt::Expr(PureExpr::If {
                cond,
                then_branch,
                else_branch: None,
            }) = &body.stmts[0]
            {
                if then_branch.stmts.len() == 1 {
                    if let Some((target_var, pushed_expr)) =
                        Self::extract_push(&then_branch.stmts[0])
                    {
                        // Check if it's filter (push var) or filter-map (push f(var))
                        if Self::is_simple_var(&pushed_expr, &var_name) {
                            return Some(LoopPattern::FilterCollect {
                                iter_expr: iter_expr.clone(),
                                var_name,
                                target_var,
                                condition: *cond.clone(),
                            });
                        } else {
                            return Some(LoopPattern::FilterMapCollect {
                                iter_expr: iter_expr.clone(),
                                var_name,
                                target_var,
                                condition: *cond.clone(),
                                transform: pushed_expr,
                            });
                        }
                    }
                }
            }
        }

        // Default: for_each pattern
        Some(LoopPattern::ForEach {
            iter_expr: iter_expr.clone(),
            var_name,
            body: body.clone(),
        })
    }

    /// Detect pattern from a single statement
    fn detect_single_stmt_pattern(
        var_name: &str,
        iter_expr: &PureExpr,
        stmt: &PureStmt,
    ) -> Option<LoopPattern> {
        // Check for vec.push(expr) pattern
        if let Some((target_var, pushed_expr)) = Self::extract_push(stmt) {
            if Self::is_simple_var(&pushed_expr, var_name) {
                // Simple collect: for x in iter { v.push(x) }
                return Some(LoopPattern::MapCollect {
                    iter_expr: iter_expr.clone(),
                    var_name: var_name.to_string(),
                    target_var,
                    transform: pushed_expr,
                });
            } else {
                // Map collect: for x in iter { v.push(f(x)) }
                return Some(LoopPattern::MapCollect {
                    iter_expr: iter_expr.clone(),
                    var_name: var_name.to_string(),
                    target_var,
                    transform: pushed_expr,
                });
            }
        }
        None
    }

    /// Extract target and expression from a push statement
    fn extract_push(stmt: &PureStmt) -> Option<(String, PureExpr)> {
        let expr = match stmt {
            PureStmt::Semi(e) | PureStmt::Expr(e) => e,
            _ => return None,
        };

        // Look for method call: target.push(arg)
        if let PureExpr::MethodCall {
            receiver,
            method,
            args,
            ..
        } = expr
        {
            if method == "push" && args.len() == 1 {
                if let PureExpr::Path(target_var) = receiver.as_ref() {
                    return Some((target_var.clone(), args[0].clone()));
                }
            }
        }
        None
    }

    /// Check if expression is just a simple variable reference
    fn is_simple_var(expr: &PureExpr, var_name: &str) -> bool {
        matches!(expr, PureExpr::Path(name) if name == var_name)
    }

    /// Convert a detected pattern to iterator expression
    fn pattern_to_iter_expr(pattern: &LoopPattern) -> PureExpr {
        match pattern {
            LoopPattern::MapCollect {
                iter_expr,
                var_name,
                transform,
                ..
            } => {
                let is_identity = Self::is_simple_var(transform, var_name);

                if is_identity {
                    // iter.collect()
                    PureExpr::MethodCall {
                        receiver: Box::new(PureExpr::MethodCall {
                            receiver: Box::new(iter_expr.clone()),
                            method: "into_iter".to_string(),
                            turbofish: None,
                            args: vec![],
                        }),
                        method: "collect".to_string(),
                        turbofish: None,
                        args: vec![],
                    }
                } else {
                    // iter.map(|var| transform).collect()
                    PureExpr::MethodCall {
                        receiver: Box::new(PureExpr::MethodCall {
                            receiver: Box::new(PureExpr::MethodCall {
                                receiver: Box::new(iter_expr.clone()),
                                method: "into_iter".to_string(),
                                turbofish: None,
                                args: vec![],
                            }),
                            method: "map".to_string(),
                            turbofish: None,
                            args: vec![PureExpr::Closure {
                                is_async: false,
                                is_move: false,
                                params: vec![PureClosureParam::untyped(PurePattern::Ident {
                                    name: var_name.clone(),
                                    is_mut: false,
                                })],
                                ret: None,
                                body: Box::new(transform.clone()),
                            }],
                        }),
                        method: "collect".to_string(),
                        turbofish: None,
                        args: vec![],
                    }
                }
            }
            LoopPattern::FilterCollect {
                iter_expr,
                var_name,
                condition,
                ..
            } => {
                // iter.filter(|&var| cond).collect()
                // Note: Use reference pattern |&var| for filter closures since filter
                // passes items by reference, but the original condition uses the value.
                PureExpr::MethodCall {
                    receiver: Box::new(PureExpr::MethodCall {
                        receiver: Box::new(PureExpr::MethodCall {
                            receiver: Box::new(iter_expr.clone()),
                            method: "into_iter".to_string(),
                            turbofish: None,
                            args: vec![],
                        }),
                        method: "filter".to_string(),
                        turbofish: None,
                        args: vec![PureExpr::Closure {
                            is_async: false,
                            is_move: false,
                            params: vec![PureClosureParam::untyped(PurePattern::Ref {
                                is_mut: false,
                                pattern: Box::new(PurePattern::Ident {
                                    name: var_name.clone(),
                                    is_mut: false,
                                }),
                            })],
                            ret: None,
                            body: Box::new(condition.clone()),
                        }],
                    }),
                    method: "collect".to_string(),
                    turbofish: None,
                    args: vec![],
                }
            }
            LoopPattern::FilterMapCollect {
                iter_expr,
                var_name,
                condition,
                transform,
                ..
            } => {
                // iter.filter(|&var| cond).map(|var| transform).collect()
                // Note: Use reference pattern |&var| for filter closures since filter
                // passes items by reference.
                PureExpr::MethodCall {
                    receiver: Box::new(PureExpr::MethodCall {
                        receiver: Box::new(PureExpr::MethodCall {
                            receiver: Box::new(PureExpr::MethodCall {
                                receiver: Box::new(iter_expr.clone()),
                                method: "into_iter".to_string(),
                                turbofish: None,
                                args: vec![],
                            }),
                            method: "filter".to_string(),
                            turbofish: None,
                            args: vec![PureExpr::Closure {
                                is_async: false,
                                is_move: false,
                                params: vec![PureClosureParam::untyped(PurePattern::Ref {
                                    is_mut: false,
                                    pattern: Box::new(PurePattern::Ident {
                                        name: var_name.clone(),
                                        is_mut: false,
                                    }),
                                })],
                                ret: None,
                                body: Box::new(condition.clone()),
                            }],
                        }),
                        method: "map".to_string(),
                        turbofish: None,
                        args: vec![PureExpr::Closure {
                            is_async: false,
                            is_move: false,
                            params: vec![PureClosureParam::untyped(PurePattern::Ident {
                                name: var_name.clone(),
                                is_mut: false,
                            })],
                            ret: None,
                            body: Box::new(transform.clone()),
                        }],
                    }),
                    method: "collect".to_string(),
                    turbofish: None,
                    args: vec![],
                }
            }
            LoopPattern::ForEach {
                iter_expr,
                var_name,
                body,
            } => {
                // iter.for_each(|var| { body })
                PureExpr::MethodCall {
                    receiver: Box::new(PureExpr::MethodCall {
                        receiver: Box::new(iter_expr.clone()),
                        method: "into_iter".to_string(),
                        turbofish: None,
                        args: vec![],
                    }),
                    method: "for_each".to_string(),
                    turbofish: None,
                    args: vec![PureExpr::Closure {
                        is_async: false,
                        is_move: false,
                        params: vec![PureClosureParam::untyped(PurePattern::Ident {
                            name: var_name.clone(),
                            is_mut: false,
                        })],
                        ret: None,
                        body: Box::new(PureExpr::Block {
                            label: None,
                            block: body.clone(),
                        }),
                    }],
                }
            }
        }
    }

    /// Detect block-level patterns: `let mut v = Vec::new(); for x in iter { v.push(...) }`
    fn detect_block_patterns(stmts: &[PureStmt]) -> Vec<BlockLoopPattern> {
        let mut patterns = Vec::new();

        for (i, stmt) in stmts.iter().enumerate() {
            // Look for `let mut v = Vec::new()` or `let mut v = vec![]`
            if let Some((var_name, is_vec_init)) = Self::extract_vec_init(stmt) {
                if !is_vec_init {
                    continue;
                }

                // Look for a following for loop that pushes to this variable
                for (j, jstmt) in stmts.iter().enumerate().skip(i + 1) {
                    if let Some(loop_pattern) = Self::extract_for_loop_push(jstmt, &var_name) {
                        patterns.push(BlockLoopPattern {
                            vec_decl_idx: i,
                            for_loop_idx: j,
                            target_var: var_name.clone(),
                            loop_pattern,
                        });
                        break; // Only match the first for loop
                    }

                    // If we hit a statement that uses `var_name` in a different way, stop
                    if Self::stmt_uses_var(jstmt, &var_name) {
                        break;
                    }
                }
            }
        }

        patterns
    }

    /// Extract variable name and check if it's `Vec::new()` or `vec![]`
    fn extract_vec_init(stmt: &PureStmt) -> Option<(String, bool)> {
        if let PureStmt::Local {
            pattern: PurePattern::Ident { name, is_mut: true },
            init: Some(init_expr),
            ..
        } = stmt
        {
            let is_vec_init = Self::is_vec_new_call(init_expr) || Self::is_vec_macro(init_expr);
            return Some((name.clone(), is_vec_init));
        }
        None
    }

    /// Check if expression is `Vec::new()`
    fn is_vec_new_call(expr: &PureExpr) -> bool {
        if let PureExpr::Call { func, args } = expr {
            if args.is_empty() {
                // Check for Vec::new() pattern
                if let PureExpr::Path(path) = func.as_ref() {
                    return path == "Vec::new" || path.ends_with("::Vec::new");
                }
            }
        }
        false
    }

    /// Check if expression is `vec![]`
    fn is_vec_macro(expr: &PureExpr) -> bool {
        if let PureExpr::Macro { name, .. } = expr {
            return name == "vec" || name.ends_with("::vec");
        }
        false
    }

    /// Extract for loop pattern if it pushes to the target variable
    fn extract_for_loop_push(stmt: &PureStmt, target_var: &str) -> Option<LoopPattern> {
        let for_expr = match stmt {
            PureStmt::Semi(e) | PureStmt::Expr(e) => e,
            _ => return None,
        };

        if let PureExpr::For {
            pat,
            expr: iter_expr,
            body,
            ..
        } = for_expr
        {
            let pattern = Self::detect_pattern(pat, iter_expr, body)?;

            // Check if the pattern targets our variable
            let pattern_target = match &pattern {
                LoopPattern::MapCollect { target_var, .. } => target_var,
                LoopPattern::FilterCollect { target_var, .. } => target_var,
                LoopPattern::FilterMapCollect { target_var, .. } => target_var,
                LoopPattern::ForEach { .. } => return None, // ForEach doesn't push
            };

            if pattern_target == target_var {
                return Some(pattern);
            }
        }
        None
    }

    /// Check if a statement uses the given variable (other than in a for loop push)
    fn stmt_uses_var(stmt: &PureStmt, var_name: &str) -> bool {
        match stmt {
            PureStmt::Semi(expr) | PureStmt::Expr(expr) => Self::expr_uses_var(expr, var_name),
            PureStmt::Local { init: Some(e), .. } => Self::expr_uses_var(e, var_name),
            _ => false,
        }
    }

    /// Check if an expression uses the given variable
    fn expr_uses_var(expr: &PureExpr, var_name: &str) -> bool {
        match expr {
            PureExpr::Path(name) => name == var_name,
            PureExpr::MethodCall { receiver, args, .. } => {
                Self::expr_uses_var(receiver, var_name)
                    || args.iter().any(|a| Self::expr_uses_var(a, var_name))
            }
            PureExpr::Call { func, args } => {
                Self::expr_uses_var(func, var_name)
                    || args.iter().any(|a| Self::expr_uses_var(a, var_name))
            }
            PureExpr::Binary { left, right, .. } => {
                Self::expr_uses_var(left, var_name) || Self::expr_uses_var(right, var_name)
            }
            PureExpr::If {
                cond,
                then_branch,
                else_branch,
            } => {
                Self::expr_uses_var(cond, var_name)
                    || then_branch
                        .stmts
                        .iter()
                        .any(|s| Self::stmt_uses_var(s, var_name))
                    || else_branch
                        .as_ref()
                        .map(|e| Self::expr_uses_var(e, var_name))
                        .unwrap_or(false)
            }
            PureExpr::For { .. } => false, // Don't recurse into for loops (they're handled separately)
            _ => false,
        }
    }

    /// Convert a block pattern to a let statement
    fn block_pattern_to_let_stmt(pattern: &BlockLoopPattern) -> PureStmt {
        let iter_expr = Self::pattern_to_iter_expr(&pattern.loop_pattern);

        PureStmt::Local {
            pattern: PurePattern::Ident {
                name: pattern.target_var.clone(),
                is_mut: false, // No longer mutable!
            },
            ty: Some(PureType::Other("Vec<_>".to_string())), // Type annotation for collect()
            init: Some(iter_expr),
        }
    }

    /// Transform for loops in a block (with block-level pattern detection)
    pub fn transform_block(&self, block: &mut PureBlock) -> usize {
        let mut changes = 0;

        // First, detect and apply block-level patterns
        let block_patterns = Self::detect_block_patterns(&block.stmts);

        if !block_patterns.is_empty() {
            // Apply patterns in reverse order to preserve indices
            let mut indices_to_remove = Vec::new();

            for pattern in block_patterns.iter().rev() {
                // Replace the vec declaration with the iterator expression
                let new_stmt = Self::block_pattern_to_let_stmt(pattern);
                block.stmts[pattern.vec_decl_idx] = new_stmt;

                // Mark the for loop for removal
                indices_to_remove.push(pattern.for_loop_idx);
                changes += 1;
            }

            // Remove the for loops (in reverse order to preserve indices)
            indices_to_remove.sort();
            indices_to_remove.reverse();
            for idx in indices_to_remove {
                block.stmts.remove(idx);
            }
        }

        // Then, recursively transform nested blocks
        for stmt in &mut block.stmts {
            changes += self.transform_stmt(stmt);
        }

        changes
    }

    /// Transform a statement
    fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
        match stmt {
            PureStmt::Semi(expr) | PureStmt::Expr(expr) => self.transform_expr(expr),
            PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
            _ => 0,
        }
    }

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

        match expr {
            PureExpr::For {
                pat,
                expr: iter_expr,
                body,
                ..
            } => {
                // Check if we should transform this loop
                if let Some(ref target) = self.target_var {
                    if let PurePattern::Ident { name, .. } = pat {
                        if name != target {
                            // Not our target, still recurse into body
                            changes += self.transform_block(body);
                            return changes;
                        }
                    }
                }

                // Standalone for loops: only ForEach pattern with aggressive mode
                if self.aggressive {
                    if let Some(pattern) = Self::detect_pattern(pat, iter_expr, body) {
                        if matches!(pattern, LoopPattern::ForEach { .. }) {
                            *expr = Self::pattern_to_iter_expr(&pattern);
                            changes += 1;
                        }
                    }
                } else {
                    // Recurse into body
                    changes += self.transform_block(body);
                }
            }
            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::Closure { body, .. } => {
                changes += self.transform_expr(body);
            }
            _ => {}
        }

        changes
    }
}

impl Mutation for LoopToIteratorMutation {
    fn describe(&self) -> String {
        "Convert for loops to iterator chains".to_string()
    }

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

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

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

    #[test]
    fn test_detect_map_pattern() {
        let pat = PurePattern::Ident {
            name: "x".to_string(),
            is_mut: false,
        };
        let iter = PureExpr::Path("items".to_string());
        let body = PureBlock {
            stmts: vec![PureStmt::Semi(PureExpr::MethodCall {
                receiver: Box::new(PureExpr::Path("result".to_string())),
                method: "push".to_string(),
                turbofish: None,
                args: vec![PureExpr::Binary {
                    op: "*".to_string(),
                    left: Box::new(PureExpr::Path("x".to_string())),
                    right: Box::new(PureExpr::Lit("2".to_string())),
                }],
            })],
        };

        let pattern = LoopToIteratorMutation::detect_pattern(&pat, &iter, &body);
        assert!(matches!(pattern, Some(LoopPattern::MapCollect { .. })));
    }

    #[test]
    fn test_detect_filter_pattern() {
        let pat = PurePattern::Ident {
            name: "x".to_string(),
            is_mut: false,
        };
        let iter = PureExpr::Path("items".to_string());
        let body = PureBlock {
            stmts: vec![PureStmt::Semi(PureExpr::If {
                cond: Box::new(PureExpr::Binary {
                    op: ">".to_string(),
                    left: Box::new(PureExpr::Path("x".to_string())),
                    right: Box::new(PureExpr::Lit("0".to_string())),
                }),
                then_branch: PureBlock {
                    stmts: vec![PureStmt::Semi(PureExpr::MethodCall {
                        receiver: Box::new(PureExpr::Path("result".to_string())),
                        method: "push".to_string(),
                        turbofish: None,
                        args: vec![PureExpr::Path("x".to_string())],
                    })],
                },
                else_branch: None,
            })],
        };

        let pattern = LoopToIteratorMutation::detect_pattern(&pat, &iter, &body);
        assert!(matches!(pattern, Some(LoopPattern::FilterCollect { .. })));
    }
}