1use crate::{Block, DeclKind, Expr, ExprKind, Program};
13
14impl Expr {
19 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(), ExprKind::IfExpr { .. } => 1, 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, ExprKind::Match { arms, .. } => 1 + arms.len(), ExprKind::VariantLit { fields, .. } => fields.len(),
53 }
54 }
55
56 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 pub fn is_place(&self) -> bool {
74 matches!(
75 &self.kind,
76 ExprKind::Ident(_) | ExprKind::Field { .. } | ExprKind::Index { .. }
77 )
78 }
79
80 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
96impl Block {
101 pub fn is_empty(&self) -> bool {
103 self.stmts.is_empty() && self.expr.is_none()
104 }
105
106 pub fn stmt_count(&self) -> usize {
108 self.stmts.len()
109 }
110
111 pub fn has_trailing_expr(&self) -> bool {
113 self.expr.is_some()
114 }
115}
116
117impl Program {
122 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 pub fn struct_count(&self) -> usize {
137 self.declarations
138 .iter()
139 .filter(|d| matches!(&d.kind, DeclKind::Struct(_)))
140 .count()
141 }
142
143 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#[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}