Skip to main content

cjc_ast/
node_utils.rs

1//! AST Node Utilities — Pure query methods on existing types
2//!
3//! Adds convenience methods to `Expr`, `Block`, and `Program` via new `impl`
4//! blocks.  These are read-only, side-effect-free query methods.
5//!
6//! ## Design decisions
7//!
8//! - **New `impl` blocks** — safe in Rust; does not modify existing impls
9//! - **No mutation** — all methods return computed values
10//! - **No dependencies** — uses only types from this crate
11
12use crate::{Block, DeclKind, Expr, ExprKind, Program};
13
14// ---------------------------------------------------------------------------
15// Expr utilities
16// ---------------------------------------------------------------------------
17
18impl Expr {
19    /// Number of direct child expressions.
20    pub fn child_count(&self) -> usize {
21        match &self.kind {
22            ExprKind::IntLit(_)
23            | ExprKind::FloatLit(_)
24            | ExprKind::StringLit(_)
25            | ExprKind::ByteStringLit(_)
26            | ExprKind::ByteCharLit(_)
27            | ExprKind::RawStringLit(_)
28            | ExprKind::RawByteStringLit(_)
29            | ExprKind::RegexLit { .. }
30            | ExprKind::BoolLit(_)
31            | ExprKind::Ident(_)
32            | ExprKind::Col(_) => 0,
33
34            ExprKind::FStringLit(segs) => segs.iter().filter(|(_, e)| e.is_some()).count(),
35            ExprKind::TensorLit { rows } => rows.iter().map(|r| r.len()).sum(),
36            ExprKind::Unary { .. } | ExprKind::Try(_) => 1,
37            ExprKind::Binary { .. }
38            | ExprKind::Assign { .. }
39            | ExprKind::CompoundAssign { .. }
40            | ExprKind::Pipe { .. }
41            | ExprKind::Index { .. } => 2,
42            ExprKind::Field { .. } => 1,
43            ExprKind::MultiIndex { object: _, indices } => 1 + indices.len(),
44            ExprKind::Call { args, .. } => 1 + args.len(), // callee + args
45            ExprKind::IfExpr { .. } => 1,                  // condition
46            ExprKind::Block(_) => 0,
47            ExprKind::StructLit { fields, .. } => fields.len(),
48            ExprKind::ArrayLit(elems) => elems.len(),
49            ExprKind::TupleLit(elems) => elems.len(),
50            ExprKind::Lambda { .. } => 1,                  // body
51            ExprKind::Match { arms, .. } => 1 + arms.len(), // scrutinee + arms
52            ExprKind::VariantLit { fields, .. } => fields.len(),
53        }
54    }
55
56    /// Returns true if this expression is a literal value.
57    pub fn is_literal(&self) -> bool {
58        matches!(
59            &self.kind,
60            ExprKind::IntLit(_)
61                | ExprKind::FloatLit(_)
62                | ExprKind::StringLit(_)
63                | ExprKind::ByteStringLit(_)
64                | ExprKind::ByteCharLit(_)
65                | ExprKind::RawStringLit(_)
66                | ExprKind::RawByteStringLit(_)
67                | ExprKind::BoolLit(_)
68                | ExprKind::RegexLit { .. }
69        )
70    }
71
72    /// Returns true if this expression is a valid assignment target (place expression).
73    pub fn is_place(&self) -> bool {
74        matches!(
75            &self.kind,
76            ExprKind::Ident(_) | ExprKind::Field { .. } | ExprKind::Index { .. }
77        )
78    }
79
80    /// Returns true if this expression is a compound (non-leaf) expression.
81    pub fn is_compound(&self) -> bool {
82        matches!(
83            &self.kind,
84            ExprKind::Binary { .. }
85                | ExprKind::Unary { .. }
86                | ExprKind::Call { .. }
87                | ExprKind::Match { .. }
88                | ExprKind::IfExpr { .. }
89                | ExprKind::Block(_)
90                | ExprKind::Pipe { .. }
91                | ExprKind::Lambda { .. }
92        )
93    }
94}
95
96// ---------------------------------------------------------------------------
97// Block utilities
98// ---------------------------------------------------------------------------
99
100impl Block {
101    /// Returns true if the block has no statements and no trailing expression.
102    pub fn is_empty(&self) -> bool {
103        self.stmts.is_empty() && self.expr.is_none()
104    }
105
106    /// Number of statements in this block.
107    pub fn stmt_count(&self) -> usize {
108        self.stmts.len()
109    }
110
111    /// Returns true if the block has a trailing expression.
112    pub fn has_trailing_expr(&self) -> bool {
113        self.expr.is_some()
114    }
115}
116
117// ---------------------------------------------------------------------------
118// Program utilities
119// ---------------------------------------------------------------------------
120
121impl Program {
122    /// Number of function declarations (including nested in impls).
123    pub fn function_count(&self) -> usize {
124        let mut count = 0;
125        for decl in &self.declarations {
126            match &decl.kind {
127                DeclKind::Fn(_) => count += 1,
128                DeclKind::Impl(i) => count += i.methods.len(),
129                _ => {}
130            }
131        }
132        count
133    }
134
135    /// Number of struct declarations.
136    pub fn struct_count(&self) -> usize {
137        self.declarations
138            .iter()
139            .filter(|d| matches!(&d.kind, DeclKind::Struct(_)))
140            .count()
141    }
142
143    /// Returns true if there is a function named "main".
144    pub fn has_main_function(&self) -> bool {
145        self.declarations.iter().any(|d| {
146            if let DeclKind::Fn(f) = &d.kind {
147                f.name.name == "main"
148            } else {
149                false
150            }
151        })
152    }
153}
154
155// ---------------------------------------------------------------------------
156// Tests
157// ---------------------------------------------------------------------------
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use crate::*;
163
164    fn dummy_expr(kind: ExprKind) -> Expr {
165        Expr {
166            kind,
167            span: Span::dummy(),
168        }
169    }
170
171    #[test]
172    fn test_expr_child_count() {
173        assert_eq!(dummy_expr(ExprKind::IntLit(1)).child_count(), 0);
174        assert_eq!(
175            dummy_expr(ExprKind::Binary {
176                op: BinOp::Add,
177                left: Box::new(dummy_expr(ExprKind::IntLit(1))),
178                right: Box::new(dummy_expr(ExprKind::IntLit(2))),
179            })
180            .child_count(),
181            2
182        );
183    }
184
185    #[test]
186    fn test_expr_is_literal() {
187        assert!(dummy_expr(ExprKind::IntLit(1)).is_literal());
188        assert!(dummy_expr(ExprKind::FloatLit(1.0)).is_literal());
189        assert!(dummy_expr(ExprKind::BoolLit(true)).is_literal());
190        assert!(!dummy_expr(ExprKind::Ident(Ident::dummy("x"))).is_literal());
191    }
192
193    #[test]
194    fn test_expr_is_place() {
195        assert!(dummy_expr(ExprKind::Ident(Ident::dummy("x"))).is_place());
196        assert!(!dummy_expr(ExprKind::IntLit(1)).is_place());
197    }
198
199    #[test]
200    fn test_expr_is_compound() {
201        assert!(dummy_expr(ExprKind::Binary {
202            op: BinOp::Add,
203            left: Box::new(dummy_expr(ExprKind::IntLit(1))),
204            right: Box::new(dummy_expr(ExprKind::IntLit(2))),
205        })
206        .is_compound());
207        assert!(!dummy_expr(ExprKind::IntLit(1)).is_compound());
208    }
209
210    #[test]
211    fn test_block_utils() {
212        let empty = Block {
213            stmts: vec![],
214            expr: None,
215            span: Span::dummy(),
216        };
217        assert!(empty.is_empty());
218        assert_eq!(empty.stmt_count(), 0);
219        assert!(!empty.has_trailing_expr());
220
221        let with_expr = Block {
222            stmts: vec![Stmt {
223                kind: StmtKind::Expr(dummy_expr(ExprKind::IntLit(1))),
224                span: Span::dummy(),
225            }],
226            expr: Some(Box::new(dummy_expr(ExprKind::IntLit(2)))),
227            span: Span::dummy(),
228        };
229        assert!(!with_expr.is_empty());
230        assert_eq!(with_expr.stmt_count(), 1);
231        assert!(with_expr.has_trailing_expr());
232    }
233
234    #[test]
235    fn test_program_utils() {
236        let program = Program {
237            declarations: vec![
238                Decl {
239                    kind: DeclKind::Fn(FnDecl {
240                        name: Ident::dummy("main"),
241                        type_params: vec![],
242                        params: vec![],
243                        return_type: None,
244                        body: Block {
245                            stmts: vec![],
246                            expr: None,
247                            span: Span::dummy(),
248                        },
249                        is_nogc: false,
250                        effect_annotation: None,
251                        decorators: vec![],
252                        vis: Visibility::Private,
253                    }),
254                    span: Span::dummy(),
255                },
256                Decl {
257                    kind: DeclKind::Struct(StructDecl {
258                        name: Ident::dummy("Point"),
259                        type_params: vec![],
260                        fields: vec![],
261                        vis: Visibility::Private,
262                    }),
263                    span: Span::dummy(),
264                },
265            ],
266        };
267        assert_eq!(program.function_count(), 1);
268        assert_eq!(program.struct_count(), 1);
269        assert!(program.has_main_function());
270    }
271}