gobject_ast/model/statement/

mod.rs

mod break_stmt;
mod compound_stmt;
mod continue_stmt;
mod for_stmt;
mod goto_stmt;
mod if_stmt;
mod labeled_stmt;
mod return_stmt;
mod switch_stmt;
mod variable_decl;
mod while_stmt;

pub use break_stmt::BreakStatement;
pub use compound_stmt::CompoundStatement;
pub use continue_stmt::ContinueStatement;
pub use for_stmt::{ForInit, ForStatement};
pub use goto_stmt::GotoStatement;
pub use if_stmt::IfStatement;
pub use labeled_stmt::LabeledStatement;
pub use return_stmt::ReturnStatement;
use serde::Serialize;
pub use switch_stmt::{CaseLabel, SwitchCase, SwitchStatement};
pub use variable_decl::VariableDecl;
pub use while_stmt::{DoWhileStatement, WhileStatement};

use crate::model::{
    Assignment, CallExpression, Comment, Expression, SourceLocation,
    top_level::PreprocessorDirective,
};

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Statement {
    Declaration(Box<VariableDecl>),
    Expression(Box<Expression>),
    If(IfStatement),
    Return(ReturnStatement),
    Goto(GotoStatement),
    Labeled(LabeledStatement),
    Compound(CompoundStatement),
    Switch(SwitchStatement),
    For(Box<ForStatement>),
    While(WhileStatement),
    DoWhile(DoWhileStatement),
    Break(BreakStatement),
    Continue(ContinueStatement),
    Preprocessor(PreprocessorDirective),
    Comment(Comment),
}

impl Statement {
    /// Recursively visit all nested statements. The closure receives a
    /// `&'s Statement` tied to `self`'s lifetime, so references extracted
    /// inside the closure can be stored in an outer `Vec<&'s T>`.
    pub fn walk<'s, F>(&'s self, f: &mut F)
    where
        F: FnMut(&'s Self),
    {
        f(self);
        match self {
            Self::If(if_stmt) => {
                for stmt in &if_stmt.then_body {
                    stmt.walk(f);
                }
                if let Some(else_body) = &if_stmt.else_body {
                    for stmt in else_body {
                        stmt.walk(f);
                    }
                }
            }
            Self::Compound(compound) => {
                for stmt in &compound.statements {
                    stmt.walk(f);
                }
            }
            Self::Labeled(labeled) => {
                labeled.statement.walk(f);
            }
            Self::Switch(switch) => {
                for case in &switch.cases {
                    for stmt in &case.body {
                        stmt.walk(f);
                    }
                }
            }
            Self::For(for_stmt) => {
                for stmt in &for_stmt.body {
                    stmt.walk(f);
                }
            }
            Self::While(while_stmt) => {
                for stmt in &while_stmt.body {
                    stmt.walk(f);
                }
            }
            Self::DoWhile(do_while) => {
                for stmt in &do_while.body {
                    stmt.walk(f);
                }
            }
            Self::Declaration(_)
            | Self::Expression(_)
            | Self::Return(_)
            | Self::Goto(_)
            | Self::Break(_)
            | Self::Continue(_)
            | Self::Preprocessor(_)
            | Self::Comment(_) => {}
        }
    }

    /// Call `f` with each immediate child block (a `&[Statement]` slice) of
    /// this statement. Covers all container variants exhaustively; leaf
    /// variants produce no calls. Use this instead of a manual `match` when
    /// you need to apply a slice-level operation (e.g. pair detection) to
    /// every nested body.
    pub fn for_each_child_block<F>(&self, mut f: F)
    where
        F: FnMut(&[Self]),
    {
        match self {
            Self::If(if_stmt) => {
                f(&if_stmt.then_body);
                if let Some(else_body) = &if_stmt.else_body {
                    f(else_body);
                }
            }
            Self::Compound(c) => f(&c.statements),
            Self::Labeled(l) => f(std::slice::from_ref(&l.statement)),
            Self::For(for_stmt) => f(&for_stmt.body),
            Self::While(w) => f(&w.body),
            Self::DoWhile(d) => f(&d.body),
            Self::Switch(sw) => {
                for case in &sw.cases {
                    f(&case.body);
                }
            }
            Self::Declaration(_)
            | Self::Expression(_)
            | Self::Return(_)
            | Self::Goto(_)
            | Self::Break(_)
            | Self::Continue(_)
            | Self::Preprocessor(_)
            | Self::Comment(_) => {}
        }
    }

    /// Visit all direct expressions contained in this statement
    /// (non-recursive). Includes conditions, initialisers, and all other
    /// expressions that are immediate children of this statement node.
    /// Zero-allocation alternative to collecting into a Vec.
    pub fn visit_expressions<'s>(&'s self, f: &mut impl FnMut(&'s Expression)) {
        match self {
            Self::Expression(expr_stmt) => f(expr_stmt),
            Self::Return(ret) => {
                if let Some(v) = &ret.value {
                    f(v);
                }
            }
            Self::Declaration(decl) => {
                if let Some(init) = &decl.initializer {
                    f(init);
                }
                if let Some(size) = &decl.array_size {
                    f(size);
                }
            }
            Self::If(if_stmt) => f(&if_stmt.condition),
            Self::Switch(switch) => {
                f(&switch.condition);
                for case in &switch.cases {
                    if let Some(label_expr) = &case.label.value {
                        f(label_expr);
                    }
                }
            }
            Self::For(for_stmt) => {
                match &for_stmt.initializer {
                    Some(ForInit::Expr(init)) => f(init),
                    Some(ForInit::Decl(decl)) => {
                        if let Some(init) = &decl.initializer {
                            f(init);
                        }
                    }
                    None => {}
                }
                if let Some(cond) = &for_stmt.condition {
                    f(cond);
                }
                if let Some(update) = &for_stmt.update {
                    f(update);
                }
            }
            Self::While(while_stmt) => f(&while_stmt.condition),
            Self::DoWhile(do_while) => f(&do_while.condition),
            Self::Goto(_)
            | Self::Labeled(_)
            | Self::Compound(_)
            | Self::Break(_)
            | Self::Continue(_)
            | Self::Preprocessor(_)
            | Self::Comment(_) => {}
        }
    }

    pub fn location(&self) -> &SourceLocation {
        match self {
            Self::Declaration(d) => &d.location,
            Self::Expression(e) => e.location(),
            Self::If(i) => &i.location,
            Self::Return(r) => &r.location,
            Self::Goto(g) => &g.location,
            Self::Labeled(l) => &l.location,
            Self::Compound(c) => &c.location,
            Self::Switch(s) => &s.location,
            Self::For(f) => &f.location,
            Self::While(w) => &w.location,
            Self::DoWhile(d) => &d.location,
            Self::Break(b) => &b.location,
            Self::Continue(c) => &c.location,
            Self::Preprocessor(p) => p.location(),
            Self::Comment(c) => &c.location,
        }
    }

    /// Recursively walk all expressions in this statement tree.
    /// Visits the direct expression of each statement (not sub-expressions —
    /// call `Expression::walk` on the result if you need nested expressions).
    pub fn walk_expressions<'s, F>(&'s self, f: &mut F)
    where
        F: FnMut(&'s Expression),
    {
        self.walk(&mut |s| {
            s.visit_expressions(f);
        });
    }

    /// Iterator over all switch statements in this statement tree (recursive)
    pub fn iter_switches<'s>(&'s self) -> impl Iterator<Item = &'s SwitchStatement> + 's {
        let mut results: Vec<&'s SwitchStatement> = Vec::new();
        self.walk(&mut |s| {
            if let Self::Switch(sw) = s {
                results.push(sw);
            }
        });
        results.into_iter()
    }

    /// Iterator over all if statements in this statement tree (recursive)
    pub fn iter_if_statements<'s>(&'s self) -> impl Iterator<Item = &'s IfStatement> + 's {
        let mut results: Vec<&'s IfStatement> = Vec::new();
        self.walk(&mut |s| {
            if let Self::If(if_stmt) = s {
                results.push(if_stmt);
            }
        });
        results.into_iter()
    }

    /// Iterator over all variable declarations in this statement tree
    /// (recursive)
    pub fn iter_declarations<'s>(&'s self) -> impl Iterator<Item = &'s VariableDecl> + 's {
        let mut results: Vec<&'s VariableDecl> = Vec::new();
        self.walk(&mut |s| {
            if let Self::Declaration(decl) = s {
                results.push(decl);
            }
            if let Self::For(for_stmt) = s
                && let Some(ForInit::Decl(decl)) = &for_stmt.initializer
            {
                results.push(decl);
            }
        });
        results.into_iter()
    }

    /// Iterator over all return statements in this statement tree (recursive)
    pub fn iter_returns<'s>(&'s self) -> impl Iterator<Item = &'s ReturnStatement> + 's {
        let mut results: Vec<&'s ReturnStatement> = Vec::new();
        self.walk(&mut |s| {
            if let Self::Return(ret) = s {
                results.push(ret);
            }
        });
        results.into_iter()
    }

    /// Iterator over all top-level assignment statements in this statement tree
    /// (recursive). Only yields assignments that are the entire expression
    /// statement, not assignments nested inside other expressions.
    pub fn iter_assignments<'s>(&'s self) -> impl Iterator<Item = &'s Assignment> + 's {
        let mut results: Vec<&'s Assignment> = Vec::new();
        self.walk(&mut |s| {
            if let Self::Expression(expr_stmt) = s
                && let Expression::Assignment(assign) = expr_stmt.as_ref()
            {
                results.push(assign);
            }
        });
        results.into_iter()
    }

    /// Iterator over all call expressions in this statement tree (recursive).
    /// Includes calls at all nesting levels within expressions (e.g. nested
    /// arguments).
    pub fn iter_calls<'s>(&'s self) -> impl Iterator<Item = &'s CallExpression> + 's {
        // Two-step: first collect top-level expressions (avoids nested-closure
        // invariance issue with &mut Vec), then walk each for nested calls.
        let mut exprs: Vec<&'s Expression> = Vec::new();
        self.walk_expressions(&mut |expr| exprs.push(expr));

        let mut results: Vec<&'s CallExpression> = Vec::new();
        for expr in exprs {
            expr.walk(&mut |e| {
                if let Expression::Call(call) = e {
                    results.push(call);
                }
            });
        }
        results.into_iter()
    }

    /// Extract the call expression if this is an expression statement with a
    /// call
    pub fn extract_call(&self) -> Option<&CallExpression> {
        if let Self::Expression(expr_stmt) = self
            && let Expression::Call(call) = expr_stmt.as_ref()
        {
            return Some(call);
        }
        None
    }

    /// Check if this statement assigns a value matching the predicate to the
    /// target variable
    pub fn is_assignment_to<F>(&self, target_var: &str, value_check: F, source: &[u8]) -> bool
    where
        F: Fn(&Expression) -> bool,
    {
        if let Self::Expression(expr_stmt) = self
            && let Expression::Assignment(assign) = expr_stmt.as_ref()
        {
            let lhs_text = assign.lhs.location().as_str(source).unwrap_or("");
            return lhs_text.trim() == target_var.trim() && value_check(&assign.rhs);
        }
        false
    }

    /// Extract the assignment expression if this is an assignment statement
    pub fn extract_assignment(&self) -> Option<&Assignment> {
        if let Self::Expression(expr_stmt) = self
            && let Expression::Assignment(assign) = expr_stmt.as_ref()
        {
            return Some(assign);
        }
        None
    }

    /// Check if this statement assigns NULL to the target variable
    pub fn is_null_assignment_to(&self, var_name: &str, source: &[u8]) -> bool {
        self.is_assignment_to(var_name, Expression::is_null, source)
    }

    fn non_comments(stmts: &[Self]) -> Vec<&Self> {
        stmts
            .iter()
            .filter(|s| !matches!(s, Self::Comment(_)))
            .collect()
    }

    /// Iterate over consecutive pairs of statements (skipping comments)
    pub fn for_each_pair<F>(statements: &[Self], mut f: F)
    where
        F: FnMut(&Self, &Self),
    {
        let nc = Self::non_comments(statements);
        for w in nc.windows(2) {
            f(w[0], w[1]);
        }
    }

    /// Iterate over consecutive triples of statements (skipping comments)
    pub fn for_each_triple<F>(statements: &[Self], mut f: F)
    where
        F: FnMut(&Self, &Self, &Self),
    {
        let nc = Self::non_comments(statements);
        for w in nc.windows(3) {
            f(w[0], w[1], w[2]);
        }
    }

    /// Recursively visit consecutive pairs of statements at every nesting level
    /// (if bodies, else bodies, compound blocks, loops, switch cases, labeled
    /// statements). Comments are skipped so they don't break pair detection.
    pub fn walk_pairs<F>(stmts: &[Self], f: &mut F)
    where
        F: FnMut(&Self, &Self),
    {
        let nc = Self::non_comments(stmts);
        for w in nc.windows(2) {
            f(w[0], w[1]);
        }
        for stmt in stmts {
            match stmt {
                Self::If(if_stmt) => {
                    Self::walk_pairs(&if_stmt.then_body, f);
                    if let Some(else_body) = &if_stmt.else_body {
                        Self::walk_pairs(else_body, f);
                    }
                }
                Self::Compound(c) => Self::walk_pairs(&c.statements, f),
                Self::Labeled(l) => Self::walk_pairs(std::slice::from_ref(&l.statement), f),
                Self::For(for_stmt) => Self::walk_pairs(&for_stmt.body, f),
                Self::While(w) => Self::walk_pairs(&w.body, f),
                Self::DoWhile(d) => Self::walk_pairs(&d.body, f),
                Self::Switch(sw) => {
                    for case in &sw.cases {
                        Self::walk_pairs(&case.body, f);
                    }
                }
                Self::Declaration(_)
                | Self::Expression(_)
                | Self::Return(_)
                | Self::Goto(_)
                | Self::Break(_)
                | Self::Continue(_)
                | Self::Preprocessor(_)
                | Self::Comment(_) => {}
            }
        }
    }

    /// Recursively visit consecutive triples of statements at every nesting
    /// level. Comments are skipped so they don't break triple detection.
    pub fn walk_triples<F>(stmts: &[Self], f: &mut F)
    where
        F: FnMut(&Self, &Self, &Self),
    {
        let nc = Self::non_comments(stmts);
        for w in nc.windows(3) {
            f(w[0], w[1], w[2]);
        }
        for stmt in stmts {
            match stmt {
                Self::If(if_stmt) => {
                    Self::walk_triples(&if_stmt.then_body, f);
                    if let Some(else_body) = &if_stmt.else_body {
                        Self::walk_triples(else_body, f);
                    }
                }
                Self::Compound(c) => Self::walk_triples(&c.statements, f),
                Self::Labeled(l) => Self::walk_triples(std::slice::from_ref(&l.statement), f),
                Self::For(for_stmt) => Self::walk_triples(&for_stmt.body, f),
                Self::While(w) => Self::walk_triples(&w.body, f),
                Self::DoWhile(d) => Self::walk_triples(&d.body, f),
                Self::Switch(sw) => {
                    for case in &sw.cases {
                        Self::walk_triples(&case.body, f);
                    }
                }
                Self::Declaration(_)
                | Self::Expression(_)
                | Self::Return(_)
                | Self::Goto(_)
                | Self::Break(_)
                | Self::Continue(_)
                | Self::Preprocessor(_)
                | Self::Comment(_) => {}
            }
        }
    }
}