windjammer 0.48.0

//! Phase 15: SIMD Vectorization (legacy stats / future AST rewrite hook)
//!
//! **Active path:** Analyzer [`crate::analyzer::simd_loops`] plus Rust codegen
//! [`crate::codegen::rust::simd_transform`] emit SIMD for select `for` loops when
//! `CompilationTarget::Rust`. This module retains loop walk + stats until the optimizer
//! pipeline is wired to owned AST again.
//! This optimization identifies loops that operate on numeric arrays
//! and transforms them to use SIMD (Single Instruction, Multiple Data)
//! operations for massive performance improvements.
//!
//! ## What is SIMD Vectorization?
//!
//! SIMD allows processing multiple data elements in parallel with a single CPU instruction.
//! Modern CPUs can process 4-8 floats or 8-16 integers simultaneously.
//!
//! ## Vectorization Patterns
//!
//! 1. **Map Operations** - Apply function to each element
//!    ```text
//!    for i in 0..n { result[i] = array[i] * 2.0 }
//!    → SIMD: Process 4-8 elements at once
//!    ```
//!
//! 2. **Reduction Operations** - Sum, product, min, max
//!    ```text
//!    let mut sum = 0.0
//!    for x in array { sum += x }
//!    → SIMD: Parallel accumulation
//!    ```
//!
//! 3. **Element-wise Operations** - Add, multiply, etc.
//!    ```text
//!    for i in 0..n { c[i] = a[i] + b[i] }
//!    → SIMD: Vectorized addition
//!    ```
//!
//! ## Performance Impact
//!
//! - **4-8x faster** for float operations (f32/f64)
//! - **8-16x faster** for integer operations (i32/i64)
//! - Near-zero overhead when not applicable
//!
//! ## Example
//!
//! ```windjammer
//! // You write:
//! fn dot_product(a: &[f64], b: &[f64]) -> f64 {
//!     let mut sum = 0.0
//!     for i in 0..a.len() {
//!         sum += a[i] * b[i]
//!     }
//!     sum
//! }
//!
//! // Compiler generates (with SIMD):
//! fn dot_product(a: &[f64], b: &[f64]) -> f64 {
//!     let mut sum = 0.0;
//!     let chunks = a.len() / 4;
//!     
//!     // Vectorized part (processes 4 f64s at once)
//!     for i in 0..chunks {
//!         let offset = i * 4;
//!         let va = f64x4::from_slice_unaligned(&a[offset..]);
//!         let vb = f64x4::from_slice_unaligned(&b[offset..]);
//!         sum += (va * vb).sum();
//!     }
//!     
//!     // Scalar remainder
//!     for i in (chunks * 4)..a.len() {
//!         sum += a[i] * b[i];
//!     }
//!     sum
//! }
//! ```

use crate::parser::*;

/// Statistics for SIMD vectorization optimization
#[derive(Debug, Clone, Default)]
pub struct SimdStats {
    /// Number of loops vectorized
    pub loops_vectorized: usize,
    /// Number of reduction operations vectorized
    pub reductions_vectorized: usize,
    /// Number of map operations vectorized
    pub maps_vectorized: usize,
    /// Total optimizations applied
    pub total_optimizations: usize,
}

impl SimdStats {
    pub fn add(&mut self, other: &SimdStats) {
        self.loops_vectorized += other.loops_vectorized;
        self.reductions_vectorized += other.reductions_vectorized;
        self.maps_vectorized += other.maps_vectorized;
        self.total_optimizations += other.total_optimizations;
    }
}

/// Perform SIMD vectorization optimization on a program
pub fn optimize_simd_vectorization<'ast>(
    program: &Program<'ast>,
    optimizer: &crate::optimizer::Optimizer,
) -> (Program<'ast>, SimdStats) {
    let mut stats = SimdStats::default();
    let mut new_items = Vec::new();

    for item in &program.items {
        let new_item = match item {
            Item::Function { decl: func, .. } => {
                let (new_func, func_stats) = optimize_function_simd(func, optimizer);
                stats.add(&func_stats);
                Item::Function {
                    decl: new_func,
                    location: None,
                }
            }
            Item::Impl {
                block: impl_block, ..
            } => {
                let (new_impl, impl_stats) = optimize_impl_simd(impl_block, optimizer);
                stats.add(&impl_stats);
                Item::Impl {
                    block: new_impl,
                    location: None,
                }
            }
            _ => item.clone(),
        };
        new_items.push(new_item);
    }

    (Program { items: new_items }, stats)
}

/// Optimize a function with SIMD vectorization
fn optimize_function_simd<'ast>(
    func: &FunctionDecl<'ast>,
    optimizer: &crate::optimizer::Optimizer,
) -> (FunctionDecl<'ast>, SimdStats) {
    let mut stats = SimdStats::default();
    let new_body = optimize_statements_simd(&func.body, &mut stats, optimizer);

    (
        FunctionDecl {
            body: new_body,
            ..func.clone()
        },
        stats,
    )
}

/// Optimize an impl block with SIMD vectorization
fn optimize_impl_simd<'ast>(
    impl_block: &ImplBlock<'ast>,
    optimizer: &crate::optimizer::Optimizer,
) -> (ImplBlock<'ast>, SimdStats) {
    let mut stats = SimdStats::default();
    let mut new_functions = Vec::new();

    for func in &impl_block.functions {
        let (new_func, func_stats) = optimize_function_simd(func, optimizer);
        stats.add(&func_stats);
        new_functions.push(new_func);
    }

    (
        ImplBlock {
            functions: new_functions,
            ..impl_block.clone()
        },
        stats,
    )
}

/// Information about a vectorizable loop
#[derive(Debug, Clone)]
struct VectorizableLoop {
    /// Loop variable name
    _variable: String,
    /// Operation type (map, reduction, etc.)
    operation_type: VectorOperation,
    /// Whether the loop can be safely vectorized
    is_safe: bool,
}

#[derive(Debug, Clone, PartialEq)]
enum VectorOperation {
    /// Map: transform each element (a[i] = f(b[i]))
    Map,
    /// Reduction: accumulate (sum += a[i])
    Reduction,
    /// ElementWise: combine arrays (c[i] = a[i] + b[i])
    #[allow(dead_code)]
    ElementWise,
    /// Unknown or not vectorizable
    Unknown,
}

/// Optimize statements with SIMD vectorization
fn optimize_statements_simd<'ast>(
    stmts: &[&'ast Statement<'ast>],
    stats: &mut SimdStats,
    optimizer: &crate::optimizer::Optimizer,
) -> Vec<&'ast Statement<'ast>> {
    let mut result = Vec::new();

    for stmt in stmts {
        let optimized = optimize_statement_simd(stmt, stats, optimizer);
        result.push(optimized);
    }

    result
}

/// Optimize a single statement with SIMD vectorization
fn optimize_statement_simd<'a: 'ast, 'ast>(
    stmt: &'a Statement<'a>,
    stats: &mut SimdStats,
    optimizer: &crate::optimizer::Optimizer,
) -> &'ast Statement<'ast> {
    match stmt {
        Statement::For {
            pattern,
            iterable,
            body,
            ..
        } => {
            // Only vectorize simple loops with identifier patterns
            if let Pattern::Identifier(variable) = pattern {
                // Check if this loop is vectorizable
                if let Some(vectorizable) = analyze_loop_vectorizability(variable, iterable, body) {
                    if vectorizable.is_safe && is_numeric_operation(&vectorizable.operation_type) {
                        // Mark for vectorization (codegen will handle actual SIMD generation)
                        stats.loops_vectorized += 1;
                        stats.total_optimizations += 1;

                        match vectorizable.operation_type {
                            VectorOperation::Reduction => stats.reductions_vectorized += 1,
                            VectorOperation::Map => stats.maps_vectorized += 1,
                            VectorOperation::ElementWise => stats.maps_vectorized += 1,
                            _ => {}
                        }

                        // Add a decorator to mark this loop as vectorizable
                        // The codegen phase will see this and generate SIMD code
                        return create_vectorized_loop(
                            variable,
                            iterable,
                            body,
                            &vectorizable,
                            optimizer,
                        );
                    }
                }
            }

            // Not vectorizable, recurse into body
            optimizer.alloc_stmt(unsafe {
                std::mem::transmute::<Statement<'_>, Statement<'_>>(Statement::For {
                    pattern: pattern.clone(),
                    iterable,
                    body: optimize_statements_simd(body, stats, optimizer),
                    location: None,
                })
            })
        }
        Statement::If {
            condition,
            then_block,
            else_block,
            ..
        } => optimizer.alloc_stmt(unsafe {
            std::mem::transmute::<Statement<'_>, Statement<'_>>(Statement::If {
                condition,
                then_block: optimize_statements_simd(then_block, stats, optimizer),
                else_block: else_block
                    .as_ref()
                    .map(|stmts| optimize_statements_simd(stmts, stats, optimizer)),
                location: None,
            })
        }),
        Statement::While {
            condition, body, ..
        } => optimizer.alloc_stmt(unsafe {
            std::mem::transmute::<Statement<'_>, Statement<'_>>(Statement::While {
                condition,
                body: optimize_statements_simd(body, stats, optimizer),
                location: None,
            })
        }),
        _ => stmt,
    }
}

/// Analyze if a loop can be vectorized
fn analyze_loop_vectorizability<'ast>(
    variable: &str,
    iterable: &'ast Expression<'ast>,
    body: &[&'ast Statement<'ast>],
) -> Option<VectorizableLoop> {
    // Check if we're iterating over a range or array
    let is_range_or_array = matches!(
        iterable,
        Expression::Range { .. } | Expression::Identifier { .. } | Expression::MethodCall { .. }
    );

    if !is_range_or_array {
        return None;
    }

    // Analyze the loop body to determine operation type
    let operation_type = classify_loop_operation(variable, body);

    // Check for vectorization hazards
    let is_safe = check_vectorization_safety(body);

    Some(VectorizableLoop {
        _variable: variable.to_string(),
        operation_type,
        is_safe,
    })
}

/// Classify what type of vector operation the loop performs
fn classify_loop_operation<'ast>(
    variable: &str,
    body: &[&'ast Statement<'ast>],
) -> VectorOperation {
    // Simple heuristic: look for common patterns
    for stmt in body {
        match stmt {
            Statement::Let { value, .. } | Statement::Const { value, .. }
                // Check for accumulation pattern (sum += ...)
                if contains_compound_assignment(value) => {
                    return VectorOperation::Reduction;
                }
            Statement::Expression { expr, .. }
                // Check for array assignment (a[i] = ...)
                if is_array_assignment(expr, variable) => {
                    return VectorOperation::Map;
                }
            _ => {}
        }
    }

    VectorOperation::Unknown
}

/// Check if vectorization is safe (no loop-carried dependencies, function calls, etc.)
fn check_vectorization_safety<'ast>(body: &[&'ast Statement<'ast>]) -> bool {
    // For safety, we'll be conservative and only vectorize simple loops
    // No function calls, no control flow, no early returns
    for stmt in body {
        match stmt {
            Statement::Return { .. } | Statement::Break { .. } | Statement::Continue { .. } => {
                return false
            }
            Statement::If { .. } | Statement::While { .. } | Statement::For { .. } => return false,
            Statement::Expression { expr, .. } if contains_function_call(expr) => {
                return false;
            }
            _ => {}
        }
    }
    true
}

/// Check if an expression contains a compound assignment (+=, *=, etc.)
fn contains_compound_assignment(expr: &Expression) -> bool {
    matches!(expr, Expression::Binary { op, .. } if matches!(op, BinaryOp::Add | BinaryOp::Mul))
}

/// Check if an expression is an array assignment pattern
fn is_array_assignment(expr: &Expression, _loop_var: &str) -> bool {
    matches!(expr, Expression::Index { .. })
}

/// Check if an expression contains a function call
fn contains_function_call(expr: &Expression) -> bool {
    match expr {
        Expression::Call { .. } => true,
        Expression::MethodCall { .. } => true,
        Expression::Binary { left, right, .. } => {
            contains_function_call(left) || contains_function_call(right)
        }
        Expression::Unary { operand, .. } => contains_function_call(operand),
        _ => false,
    }
}

/// Check if an operation is numeric (can benefit from SIMD)
fn is_numeric_operation(op: &VectorOperation) -> bool {
    matches!(
        op,
        VectorOperation::Map | VectorOperation::Reduction | VectorOperation::ElementWise
    )
}

/// Create a vectorized version of the loop
fn create_vectorized_loop<'ast>(
    variable: &str,
    iterable: &'ast Expression<'ast>,
    body: &[&'ast Statement<'ast>],
    _info: &VectorizableLoop,
    optimizer: &crate::optimizer::Optimizer,
) -> &'ast Statement<'ast> {
    // In the real implementation, codegen would recognize vectorizable patterns
    // and generate SIMD code. For now, we just preserve the loop structure
    // and track it in stats.
    optimizer.alloc_stmt(unsafe {
        std::mem::transmute::<Statement<'_>, Statement<'_>>(Statement::For {
            pattern: Pattern::Identifier(variable.to_string()),
            iterable,
            body: body.to_vec(),
            location: None,
        })
    })
}

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

    #[cfg(test)]
    use crate::parser::{Literal, Type};
    use crate::test_utils::{test_alloc_expr, test_alloc_stmt};

    #[test]
    #[allow(unused_comparisons, clippy::absurd_extreme_comparisons)]
    fn test_simd_reduction_pattern() {
        // Test: for i in 0..n { sum += array[i] }
        let program = Program {
            items: vec![Item::Function {
                decl: FunctionDecl {
                    is_pub: false,
                    is_extern: false,
                    name: "sum_array".to_string(),
                    parameters: vec![],
                    return_type: None,
                    return_decorators: Vec::new(),
                    body: vec![
                        test_alloc_stmt(Statement::Let {
                            pattern: Pattern::Identifier("sum".to_string()),
                            mutable: true,
                            type_: Some(Type::Custom("f64".to_string())),
                            value: test_alloc_expr(Expression::Literal {
                                value: Literal::Float(0.0),
                                location: None,
                            }),
                            else_block: None,
                            location: None,
                        }),
                        test_alloc_stmt(Statement::For {
                            pattern: Pattern::Identifier("i".to_string()),
                            iterable: test_alloc_expr(Expression::Range {
                                start: test_alloc_expr(Expression::Literal {
                                    value: Literal::Int(0),
                                    location: None,
                                }),
                                end: test_alloc_expr(Expression::Identifier {
                                    name: "n".to_string(),
                                    location: None,
                                }),
                                inclusive: false,
                                location: None,
                            }),
                            body: vec![test_alloc_stmt(Statement::Expression {
                                expr: test_alloc_expr(Expression::Binary {
                                    left: test_alloc_expr(Expression::Identifier {
                                        name: "sum".to_string(),
                                        location: None,
                                    }),
                                    op: BinaryOp::Add,
                                    right: test_alloc_expr(Expression::Index {
                                        object: test_alloc_expr(Expression::Identifier {
                                            name: "array".to_string(),
                                            location: None,
                                        }),
                                        index: test_alloc_expr(Expression::Identifier {
                                            name: "i".to_string(),
                                            location: None,
                                        }),
                                        location: None,
                                    }),
                                    location: None,
                                }),
                                location: None,
                            })],
                            location: None,
                        }),
                    ],
                    type_params: vec![],
                    where_clause: vec![],
                    is_async: false,
                    decorators: vec![],
                    parent_type: None,
                    impl_trait: None,
                    doc_comment: None,
                },
                location: None,
            }],
        };

        let optimizer = crate::optimizer::Optimizer::with_defaults();
        let (optimized, stats) = optimize_simd_vectorization(&program, &optimizer);

        // Should attempt to vectorize the reduction loop
        // Note: The current implementation may not vectorize all patterns yet
        // This test verifies the analysis runs without panicking
        assert!(stats.loops_vectorized >= 0);
        assert!(stats.total_optimizations >= 0);

        // Verify structure is preserved
        assert_eq!(optimized.items.len(), 1);
    }

    #[test]
    fn test_simd_unsafe_loop() {
        // Test: loop with function call (should NOT vectorize)
        let program = Program {
            items: vec![Item::Function {
                decl: FunctionDecl {
                    is_pub: false,
                    is_extern: false,
                    name: "complex".to_string(),
                    parameters: vec![],
                    return_type: None,
                    return_decorators: Vec::new(),
                    body: vec![test_alloc_stmt(Statement::For {
                        pattern: Pattern::Identifier("i".to_string()),
                        iterable: test_alloc_expr(Expression::Range {
                            start: test_alloc_expr(Expression::Literal {
                                value: Literal::Int(0),
                                location: None,
                            }),
                            end: test_alloc_expr(Expression::Literal {
                                value: Literal::Int(10),
                                location: None,
                            }),
                            inclusive: false,
                            location: None,
                        }),
                        body: vec![test_alloc_stmt(Statement::Expression {
                            expr: test_alloc_expr(Expression::Call {
                                function: test_alloc_expr(Expression::Identifier {
                                    name: "println".to_string(),
                                    location: None,
                                }),
                                arguments: vec![],
                                location: None,
                            }),
                            location: None,
                        })],
                        location: None,
                    })],
                    type_params: vec![],
                    where_clause: vec![],
                    is_async: false,
                    decorators: vec![],
                    parent_type: None,
                    impl_trait: None,
                    doc_comment: None,
                },
                location: None,
            }],
        };

        let optimizer = crate::optimizer::Optimizer::with_defaults();
        let (_, stats) = optimize_simd_vectorization(&program, &optimizer);

        // Should NOT vectorize (has function call)
        assert_eq!(stats.loops_vectorized, 0);
    }
}