solar_sema/
stats.rs

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
34/// Stat collector.
35struct 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    // Record a top-level node.
48    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    // Record a two-level entry, with a top-level enum type and a variant.
53    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
136// Used to avoid boilerplate for types with many variants.
137macro_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        // Don't visit ty field since it isn't boxed
203        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        // Don't visit name field since it isn't boxed
221        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        // Don't visit header field since it isn't boxed
233        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        // Don't visit name field since it isn't boxed
245        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        // Don't visit name field since it isn't boxed
257        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        // Don't visit name or ty field since they aren't boxed
266        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        // Don't visit name field since it isn't boxed
276        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        // Don't visit name field since it isn't boxed
286        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        // Don't visit span, ty, name, or initializer since they aren't boxed
298        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        // Don't visit ident field since it isn't boxed
322        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        // Don't visit name or arguments field since they aren't boxed
331        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        // Don't visit name field since it isn't boxed
397        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        // Don't visit span field since it isn't boxed
442        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        // Don't visit selector field since it isn't boxed
477        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        // Don't visit lit field since it isn't boxed
489        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        // Don't visit ident field since it isn't boxed
499        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        // Don't visit name field since it isn't boxed
520        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        // Don't visit span field since it isn't boxed
540        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        // Don't visit span field since it isn't boxed
551        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}