1use solar_ast::{self as ast, ItemId, visit::Visit, yul};
2use solar_data_structures::{
3 Never,
4 map::{FxHashMap, FxHashSet},
5};
6use std::ops::ControlFlow;
7
8struct NodeStats {
9 count: usize,
10 size: usize,
11}
12
13impl NodeStats {
14 fn new() -> Self {
15 Self { count: 0, size: 0 }
16 }
17
18 fn accum_size(&self) -> usize {
19 self.count * self.size
20 }
21}
22
23struct Node {
24 stats: NodeStats,
25 subnodes: FxHashMap<&'static str, NodeStats>,
26}
27
28impl Node {
29 fn new() -> Self {
30 Self { stats: NodeStats::new(), subnodes: FxHashMap::default() }
31 }
32}
33
34struct StatCollector {
36 nodes: FxHashMap<&'static str, Node>,
37 seen: FxHashSet<ItemId>,
38}
39
40pub fn print_ast_stats<'ast>(ast: &'ast ast::SourceUnit<'ast>, title: &str, prefix: &str) {
41 let mut collector = StatCollector { nodes: FxHashMap::default(), seen: FxHashSet::default() };
42 let _ = collector.visit_source_unit(ast);
43 collector.print(title, prefix)
44}
45
46impl StatCollector {
47 fn record<T: ?Sized>(&mut self, label: &'static str, id: Option<ItemId>, val: &T) {
49 self.record_inner(label, None, id, val);
50 }
51
52 fn record_variant<T: ?Sized>(
54 &mut self,
55 label1: &'static str,
56 label2: &'static str,
57 id: Option<ItemId>,
58 val: &T,
59 ) {
60 self.record_inner(label1, Some(label2), id, val);
61 }
62
63 fn record_inner<T: ?Sized>(
64 &mut self,
65 label1: &'static str,
66 label2: Option<&'static str>,
67 id: Option<ItemId>,
68 val: &T,
69 ) {
70 if id.is_some_and(|x| !self.seen.insert(x)) {
71 return;
72 }
73
74 let node = self.nodes.entry(label1).or_insert(Node::new());
75 node.stats.count += 1;
76 node.stats.size = size_of_val(val);
77
78 if let Some(label2) = label2 {
79 let subnode = node.subnodes.entry(label2).or_insert(NodeStats::new());
80 subnode.count += 1;
81 subnode.size = size_of_val(val);
82 }
83 }
84
85 fn print(&self, title: &str, prefix: &str) {
86 let mut nodes: Vec<_> = self.nodes.iter().collect();
87 nodes.sort_by_cached_key(|(label, node)| (node.stats.accum_size(), label.to_string()));
88
89 let total_size = nodes.iter().map(|(_, node)| node.stats.accum_size()).sum();
90
91 eprintln!("{prefix} {title}");
92 eprintln!(
93 "{} {:<18}{:>18}{:>14}{:>14}",
94 prefix, "Name", "Accumulated Size", "Count", "Item Size"
95 );
96 eprintln!("{prefix} ----------------------------------------------------------------");
97
98 let percent = |m, n| (m * 100) as f64 / n as f64;
99
100 for (label, node) in nodes {
101 let size = node.stats.accum_size();
102 eprintln!(
103 "{} {:<18}{:>10} ({:4.1}%){:>14}{:>14}",
104 prefix,
105 label,
106 to_readable_str(size),
107 percent(size, total_size),
108 to_readable_str(node.stats.count),
109 to_readable_str(node.stats.size)
110 );
111 if !node.subnodes.is_empty() {
112 let mut subnodes: Vec<_> = node.subnodes.iter().collect();
113 subnodes.sort_by_cached_key(|(label, subnode)| {
114 (subnode.accum_size(), label.to_string())
115 });
116
117 for (label, subnode) in subnodes {
118 let size = subnode.accum_size();
119 eprintln!(
120 "{} - {:<16}{:>10} ({:4.1}%){:>14}",
121 prefix,
122 label,
123 to_readable_str(size),
124 percent(size, total_size),
125 to_readable_str(subnode.count),
126 );
127 }
128 }
129 }
130 eprintln!("{prefix} ----------------------------------------------------------------");
131 eprintln!("{} {:<18}{:>10}", prefix, "Total", to_readable_str(total_size));
132 eprintln!("{prefix}");
133 }
134}
135
136macro_rules! record_variants {
138 (
139 ($self:ident, $val:expr, $kind:expr, $id:expr, $mod:ident, $ty:ty, $tykind:ident),
140 [$($variant:ident),*]
141 ) => {
142 match $kind {
143 $(
144 $mod::$tykind::$variant { .. } => {
145 $self.record_variant(stringify!($ty), stringify!($variant), $id, $val)
146 }
147 )*
148 }
149 };
150}
151
152impl<'ast> Visit<'ast> for StatCollector {
153 type BreakValue = Never;
154
155 fn visit_source_unit(
156 &mut self,
157 source_unit: &'ast ast::SourceUnit<'ast>,
158 ) -> ControlFlow<Self::BreakValue> {
159 self.record("SourceUnit", None, source_unit);
160 self.walk_source_unit(source_unit)
161 }
162
163 fn visit_item(&mut self, item: &'ast ast::Item<'ast>) -> ControlFlow<Self::BreakValue> {
164 record_variants!(
165 (self, item, item.kind, None, ast, Item, ItemKind),
166 [Pragma, Import, Using, Contract, Function, Variable, Struct, Enum, Udvt, Error, Event]
167 );
168 self.walk_item(item)
169 }
170
171 fn visit_pragma_directive(
172 &mut self,
173 pragma: &'ast ast::PragmaDirective<'ast>,
174 ) -> ControlFlow<Self::BreakValue> {
175 self.record("PragmaDirective", None, pragma);
176 self.walk_pragma_directive(pragma)
177 }
178
179 fn visit_import_directive(
180 &mut self,
181 import: &'ast ast::ImportDirective<'ast>,
182 ) -> ControlFlow<Self::BreakValue> {
183 self.record("ImportDirective", None, import);
184 self.walk_import_directive(import)
185 }
186
187 fn visit_using_directive(
188 &mut self,
189 using: &'ast ast::UsingDirective<'ast>,
190 ) -> ControlFlow<Self::BreakValue> {
191 self.record("UsingDirective", None, using);
192 match &using.list {
193 ast::UsingList::Single(path) => {
194 self.visit_path(path)?;
195 }
196 ast::UsingList::Multiple(paths) => {
197 for (path, _) in paths.iter() {
198 self.visit_path(path)?;
199 }
200 }
201 }
202 ControlFlow::Continue(())
204 }
205
206 fn visit_item_contract(
207 &mut self,
208 contract: &'ast ast::ItemContract<'ast>,
209 ) -> ControlFlow<Self::BreakValue> {
210 self.record("ItemContract", None, contract);
211 if let Some(layout) = &contract.layout {
212 self.visit_expr(layout.slot)?;
213 }
214 for base in contract.bases.iter() {
215 self.visit_modifier(base)?;
216 }
217 for item in contract.body.iter() {
218 self.visit_item(item)?;
219 }
220 ControlFlow::Continue(())
222 }
223
224 fn visit_item_function(
225 &mut self,
226 function: &'ast ast::ItemFunction<'ast>,
227 ) -> ControlFlow<Self::BreakValue> {
228 self.record("ItemFunction", None, function);
229 if let Some(body) = &function.body {
230 self.visit_block(body)?;
231 }
232 ControlFlow::Continue(())
234 }
235
236 fn visit_item_struct(
237 &mut self,
238 strukt: &'ast ast::ItemStruct<'ast>,
239 ) -> ControlFlow<Self::BreakValue> {
240 self.record("ItemStruct", None, strukt);
241 for field in strukt.fields.iter() {
242 self.visit_variable_definition(field)?;
243 }
244 ControlFlow::Continue(())
246 }
247
248 fn visit_item_enum(
249 &mut self,
250 enum_: &'ast ast::ItemEnum<'ast>,
251 ) -> ControlFlow<Self::BreakValue> {
252 self.record("ItemEnum", None, enum_);
253 for variant in enum_.variants.iter() {
254 self.visit_ident(variant)?;
255 }
256 ControlFlow::Continue(())
258 }
259
260 fn visit_item_udvt(
261 &mut self,
262 udvt: &'ast ast::ItemUdvt<'ast>,
263 ) -> ControlFlow<Self::BreakValue> {
264 self.record("ItemUdvt", None, udvt);
265 ControlFlow::Continue(())
267 }
268
269 fn visit_item_error(
270 &mut self,
271 error: &'ast ast::ItemError<'ast>,
272 ) -> ControlFlow<Self::BreakValue> {
273 self.record("ItemError", None, error);
274 self.visit_parameter_list(&error.parameters)?;
275 ControlFlow::Continue(())
277 }
278
279 fn visit_item_event(
280 &mut self,
281 event: &'ast ast::ItemEvent<'ast>,
282 ) -> ControlFlow<Self::BreakValue> {
283 self.record("ItemEvent", None, event);
284 self.visit_parameter_list(&event.parameters)?;
285 ControlFlow::Continue(())
287 }
288
289 fn visit_variable_definition(
290 &mut self,
291 var: &'ast ast::VariableDefinition<'ast>,
292 ) -> ControlFlow<Self::BreakValue> {
293 self.record("VariableDefinition", None, var);
294 if let Some(initializer) = &var.initializer {
295 self.visit_expr(initializer)?;
296 }
297 ControlFlow::Continue(())
299 }
300
301 fn visit_ty(&mut self, ty: &'ast ast::Type<'ast>) -> ControlFlow<Self::BreakValue> {
302 record_variants!(
303 (self, ty, ty.kind, None, ast, Type, TypeKind),
304 [Elementary, Array, Function, Mapping, Custom]
305 );
306 self.walk_ty(ty)
307 }
308
309 fn visit_function_header(
310 &mut self,
311 header: &'ast ast::FunctionHeader<'ast>,
312 ) -> ControlFlow<Self::BreakValue> {
313 self.record("FunctionHeader", None, header);
314 self.visit_parameter_list(&header.parameters)?;
315 for modifier in header.modifiers.iter() {
316 self.visit_modifier(modifier)?;
317 }
318 if let Some(returns) = &header.returns {
319 self.visit_parameter_list(returns)?;
320 }
321 ControlFlow::Continue(())
323 }
324
325 fn visit_modifier(
326 &mut self,
327 modifier: &'ast ast::Modifier<'ast>,
328 ) -> ControlFlow<Self::BreakValue> {
329 self.record("Modifier", None, modifier);
330 ControlFlow::Continue(())
332 }
333
334 fn visit_call_args(
335 &mut self,
336 args: &'ast ast::CallArgs<'ast>,
337 ) -> ControlFlow<Self::BreakValue> {
338 self.record("CallArgs", None, args);
339 self.walk_call_args(args)
340 }
341
342 fn visit_named_args(
343 &mut self,
344 args: &'ast ast::NamedArgList<'ast>,
345 ) -> ControlFlow<Self::BreakValue> {
346 self.record("NamedArgList", None, args);
347 self.walk_named_args(args)
348 }
349
350 fn visit_stmt(&mut self, stmt: &'ast ast::Stmt<'ast>) -> ControlFlow<Self::BreakValue> {
351 record_variants!(
352 (self, stmt, stmt.kind, None, ast, Stmt, StmtKind),
353 [
354 Assembly,
355 DeclSingle,
356 DeclMulti,
357 Block,
358 Break,
359 Continue,
360 DoWhile,
361 Emit,
362 Expr,
363 For,
364 If,
365 Return,
366 Revert,
367 Try,
368 UncheckedBlock,
369 While,
370 Placeholder
371 ]
372 );
373 self.walk_stmt(stmt)
374 }
375
376 fn visit_stmt_assembly(
377 &mut self,
378 assembly: &'ast ast::StmtAssembly<'ast>,
379 ) -> ControlFlow<Self::BreakValue> {
380 self.record("StmtAssembly", None, assembly);
381 self.walk_stmt_assembly(assembly)
382 }
383
384 fn visit_stmt_try(&mut self, try_: &'ast ast::StmtTry<'ast>) -> ControlFlow<Self::BreakValue> {
385 self.record("StmtTry", None, try_);
386 self.walk_stmt_try(try_)
387 }
388
389 fn visit_try_catch_clause(
390 &mut self,
391 catch: &'ast ast::TryCatchClause<'ast>,
392 ) -> ControlFlow<Self::BreakValue> {
393 self.record("TryCatchClause", None, catch);
394 self.visit_parameter_list(&catch.args)?;
395 self.visit_block(&catch.block)?;
396 ControlFlow::Continue(())
398 }
399
400 fn visit_block(&mut self, block: &'ast ast::Block<'ast>) -> ControlFlow<Self::BreakValue> {
401 self.record("Block", None, block);
402 self.walk_block(block)
403 }
404
405 fn visit_expr(&mut self, expr: &'ast ast::Expr<'ast>) -> ControlFlow<Self::BreakValue> {
406 record_variants!(
407 (self, expr, expr.kind, None, ast, Expr, ExprKind),
408 [
409 Array,
410 Assign,
411 Binary,
412 Call,
413 CallOptions,
414 Delete,
415 Ident,
416 Index,
417 Lit,
418 Member,
419 New,
420 Payable,
421 Ternary,
422 Tuple,
423 TypeCall,
424 Type,
425 Unary
426 ]
427 );
428 self.walk_expr(expr)
429 }
430
431 fn visit_parameter_list(
432 &mut self,
433 list: &'ast ast::ParameterList<'ast>,
434 ) -> ControlFlow<Self::BreakValue> {
435 self.record("ParameterList", None, list);
436 self.walk_parameter_list(list)
437 }
438
439 fn visit_lit(&mut self, lit: &'ast ast::Lit<'_>) -> ControlFlow<Self::BreakValue> {
440 self.record("Lit", None, lit);
441 ControlFlow::Continue(())
443 }
444
445 fn visit_yul_stmt(&mut self, stmt: &'ast yul::Stmt<'ast>) -> ControlFlow<Self::BreakValue> {
446 record_variants!(
447 (self, stmt, stmt.kind, None, yul, Stmt, StmtKind),
448 [
449 Block,
450 AssignSingle,
451 AssignMulti,
452 Expr,
453 If,
454 For,
455 Switch,
456 Leave,
457 Break,
458 Continue,
459 FunctionDef,
460 VarDecl
461 ]
462 );
463 self.walk_yul_stmt(stmt)
464 }
465
466 fn visit_yul_block(&mut self, block: &'ast yul::Block<'ast>) -> ControlFlow<Self::BreakValue> {
467 self.record("YulBlock", None, block);
468 self.walk_yul_block(block)
469 }
470
471 fn visit_yul_stmt_switch(
472 &mut self,
473 switch: &'ast yul::StmtSwitch<'ast>,
474 ) -> ControlFlow<Self::BreakValue> {
475 self.record("YulStmtSwitch", None, switch);
476 for case in switch.cases.iter() {
478 self.visit_yul_stmt_case(case)?;
479 }
480 ControlFlow::Continue(())
481 }
482
483 fn visit_yul_stmt_case(
484 &mut self,
485 case: &'ast yul::StmtSwitchCase<'ast>,
486 ) -> ControlFlow<Self::BreakValue> {
487 self.record("YulStmtSwitchCase", None, case);
488 self.visit_yul_block(&case.body)?;
490 ControlFlow::Continue(())
491 }
492
493 fn visit_yul_function(
494 &mut self,
495 function: &'ast yul::Function<'ast>,
496 ) -> ControlFlow<Self::BreakValue> {
497 self.record("YulFunction", None, function);
498 for ident in function.parameters.iter() {
500 self.visit_ident(ident)?;
501 }
502 for ident in function.returns.iter() {
503 self.visit_ident(ident)?;
504 }
505 self.visit_yul_block(&function.body)?;
506 ControlFlow::Continue(())
507 }
508
509 fn visit_yul_expr(&mut self, expr: &'ast yul::Expr<'ast>) -> ControlFlow<Self::BreakValue> {
510 record_variants!((self, expr, expr.kind, None, yul, Expr, ExprKind), [Path, Call, Lit]);
511 self.walk_yul_expr(expr)
512 }
513
514 fn visit_yul_expr_call(
515 &mut self,
516 call: &'ast yul::ExprCall<'ast>,
517 ) -> ControlFlow<Self::BreakValue> {
518 self.record("YulExprCall", None, call);
519 for arg in call.arguments.iter() {
521 self.visit_yul_expr(arg)?;
522 }
523 ControlFlow::Continue(())
524 }
525
526 fn visit_doc_comments(
527 &mut self,
528 doc_comments: &'ast ast::DocComments<'ast>,
529 ) -> ControlFlow<Self::BreakValue> {
530 self.record("DocComments", None, doc_comments);
531 self.walk_doc_comments(doc_comments)
532 }
533
534 fn visit_doc_comment(
535 &mut self,
536 doc_comment: &'ast ast::DocComment,
537 ) -> ControlFlow<Self::BreakValue> {
538 self.record("DocComment", None, doc_comment);
539 ControlFlow::Continue(())
541 }
542
543 fn visit_path(&mut self, path: &'ast ast::PathSlice) -> ControlFlow<Self::BreakValue> {
544 self.record("PathSlice", None, path);
545 self.walk_path(path)
546 }
547
548 fn visit_ident(&mut self, ident: &'ast ast::Ident) -> ControlFlow<Self::BreakValue> {
549 self.record("Ident", None, ident);
550 ControlFlow::Continue(())
552 }
553
554 fn visit_span(&mut self, span: &'ast ast::Span) -> ControlFlow<Self::BreakValue> {
555 self.record("Span", None, span);
556 ControlFlow::Continue(())
557 }
558}
559
560pub fn to_readable_str(mut val: usize) -> String {
561 let mut groups = vec![];
562 loop {
563 let group = val % 1000;
564 val /= 1000;
565 if val == 0 {
566 groups.push(group.to_string());
567 break;
568 } else {
569 groups.push(format!("{group:03}"));
570 }
571 }
572 groups.reverse();
573 groups.join("_")
574}