hakorune/src/mir/loop_canonicalizer/pattern_recognizer.rs

//! Pattern Recognition Helpers
//!
//! This module contains pattern detection logic for specific loop structures.
//! Currently supports the skip_whitespace pattern, with room for future patterns.

use crate::ast::ASTNode;

// ============================================================================
// Skip Whitespace Pattern
// ============================================================================

/// Try to extract skip_whitespace pattern from loop
///
/// Pattern structure:
/// ```
/// loop(cond) {
///     // ... optional body statements (Body)
///     if check_cond {
///         carrier = carrier + const
///     } else {
///         break
///     }
/// }
/// ```
///
/// Returns (carrier_name, delta, body_stmts) if pattern matches.
pub fn try_extract_skip_whitespace_pattern(
    body: &[ASTNode],
) -> Option<(String, i64, Vec<ASTNode>)> {
    if body.is_empty() {
        return None;
    }

    // Last statement must be if-else with break
    let last_stmt = &body[body.len() - 1];

    let (then_body, else_body) = match last_stmt {
        ASTNode::If {
            then_body,
            else_body: Some(else_body),
            ..
        } => (then_body, else_body),
        _ => return None,
    };

    // Then branch must be single assignment: carrier = carrier + const
    if then_body.len() != 1 {
        return None;
    }

    let (carrier_name, delta) = match &then_body[0] {
        ASTNode::Assignment { target, value, .. } => {
            // Extract target variable name
            let target_name = match target.as_ref() {
                ASTNode::Variable { name, .. } => name.clone(),
                _ => return None,
            };

            // Value must be: target + const
            match value.as_ref() {
                ASTNode::BinaryOp {
                    operator: crate::ast::BinaryOperator::Add,
                    left,
                    right,
                    ..
                } => {
                    // Left must be same variable
                    let left_name = match left.as_ref() {
                        ASTNode::Variable { name, .. } => name,
                        _ => return None,
                    };

                    if left_name != &target_name {
                        return None;
                    }

                    // Right must be integer literal
                    let delta = match right.as_ref() {
                        ASTNode::Literal {
                            value: crate::ast::LiteralValue::Integer(n),
                            ..
                        } => *n,
                        _ => return None,
                    };

                    (target_name, delta)
                }
                _ => return None,
            }
        }
        _ => return None,
    };

    // Else branch must be single break
    if else_body.len() != 1 {
        return None;
    }

    match &else_body[0] {
        ASTNode::Break { .. } => {
            // Success! Extract body statements (all except last if)
            let body_stmts = body[..body.len() - 1].to_vec();
            Some((carrier_name, delta, body_stmts))
        }
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{BinaryOperator, LiteralValue, Span};

    #[test]
    fn test_skip_whitespace_basic_pattern() {
        // Build: if is_ws { p = p + 1 } else { break }
        let body = vec![ASTNode::If {
            condition: Box::new(ASTNode::Variable {
                name: "is_ws".to_string(),
                span: Span::unknown(),
            }),
            then_body: vec![ASTNode::Assignment {
                target: Box::new(ASTNode::Variable {
                    name: "p".to_string(),
                    span: Span::unknown(),
                }),
                value: Box::new(ASTNode::BinaryOp {
                    operator: BinaryOperator::Add,
                    left: Box::new(ASTNode::Variable {
                        name: "p".to_string(),
                        span: Span::unknown(),
                    }),
                    right: Box::new(ASTNode::Literal {
                        value: LiteralValue::Integer(1),
                        span: Span::unknown(),
                    }),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }],
            else_body: Some(vec![ASTNode::Break {
                span: Span::unknown(),
            }]),
            span: Span::unknown(),
        }];

        let result = try_extract_skip_whitespace_pattern(&body);
        assert!(result.is_some());

        let (carrier_name, delta, body_stmts) = result.unwrap();
        assert_eq!(carrier_name, "p");
        assert_eq!(delta, 1);
        assert_eq!(body_stmts.len(), 0);
    }

    #[test]
    fn test_skip_whitespace_with_body() {
        // Build: local ch = get_char(p); if is_ws { p = p + 1 } else { break }
        let body = vec![
            ASTNode::Assignment {
                target: Box::new(ASTNode::Variable {
                    name: "ch".to_string(),
                    span: Span::unknown(),
                }),
                value: Box::new(ASTNode::FunctionCall {
                    name: "get_char".to_string(),
                    arguments: vec![ASTNode::Variable {
                        name: "p".to_string(),
                        span: Span::unknown(),
                    }],
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            },
            ASTNode::If {
                condition: Box::new(ASTNode::Variable {
                    name: "is_ws".to_string(),
                    span: Span::unknown(),
                }),
                then_body: vec![ASTNode::Assignment {
                    target: Box::new(ASTNode::Variable {
                        name: "p".to_string(),
                        span: Span::unknown(),
                    }),
                    value: Box::new(ASTNode::BinaryOp {
                        operator: BinaryOperator::Add,
                        left: Box::new(ASTNode::Variable {
                            name: "p".to_string(),
                            span: Span::unknown(),
                        }),
                        right: Box::new(ASTNode::Literal {
                            value: LiteralValue::Integer(1),
                            span: Span::unknown(),
                        }),
                        span: Span::unknown(),
                    }),
                    span: Span::unknown(),
                }],
                else_body: Some(vec![ASTNode::Break {
                    span: Span::unknown(),
                }]),
                span: Span::unknown(),
            },
        ];

        let result = try_extract_skip_whitespace_pattern(&body);
        assert!(result.is_some());

        let (carrier_name, delta, body_stmts) = result.unwrap();
        assert_eq!(carrier_name, "p");
        assert_eq!(delta, 1);
        assert_eq!(body_stmts.len(), 1); // The assignment before the if
    }

    #[test]
    fn test_skip_whitespace_rejects_no_else() {
        // Build: if is_ws { p = p + 1 } (no else)
        let body = vec![ASTNode::If {
            condition: Box::new(ASTNode::Variable {
                name: "is_ws".to_string(),
                span: Span::unknown(),
            }),
            then_body: vec![ASTNode::Assignment {
                target: Box::new(ASTNode::Variable {
                    name: "p".to_string(),
                    span: Span::unknown(),
                }),
                value: Box::new(ASTNode::BinaryOp {
                    operator: BinaryOperator::Add,
                    left: Box::new(ASTNode::Variable {
                        name: "p".to_string(),
                        span: Span::unknown(),
                    }),
                    right: Box::new(ASTNode::Literal {
                        value: LiteralValue::Integer(1),
                        span: Span::unknown(),
                    }),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }],
            else_body: None,
            span: Span::unknown(),
        }];

        let result = try_extract_skip_whitespace_pattern(&body);
        assert!(result.is_none());
    }
}