Skip to main content

walk_ast/
walk_ast.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8//! Walk a parsed AST with the read-only `Visitor` and print a node-kind
9//! histogram.
10//!
11//! ```text
12//! cargo run -p hermes-parser --example walk_ast
13//! ```
14
15use std::collections::HashMap;
16
17// The AST crate is re-exported as `hermes_parser::ast`, so this example needs
18// no second dependency.
19use hermes_parser::ast::node::{Node, NodeKind};
20use hermes_parser::ast::visitor::Visitor;
21use hermes_parser::{parse, ParseFlags};
22
23/// Counts how many nodes of each kind the tree contains.
24struct Histogram {
25    counts: HashMap<NodeKind, usize>,
26}
27
28impl<'gc> Visitor<'gc> for Histogram {
29    fn visit_node(&mut self, node: &'gc Node<'gc>) {
30        *self.counts.entry(node.kind()).or_default() += 1;
31        // The default `visit_node` does exactly this; recursion is ours to
32        // control, so a visitor can prune subtrees by not calling it.
33        node.visit_children(self);
34    }
35}
36
37const SOURCE: &str = r#"
38class Counter {
39  #n = 0;
40  increment(by = 1) {
41    this.#n += by;
42    return this.#n;
43  }
44}
45
46const c = new Counter();
47for (const step of [1, 2, 3]) {
48  console.log(`${step} -> ${c.increment(step)}`);
49}
50"#;
51
52fn main() {
53    let flags = ParseFlags::default();
54    let mut parsed = parse(SOURCE, flags).expect("snippet must parse");
55
56    // The AST is only reachable while the arena is locked, so collect owned
57    // data inside the closure and return it.
58    let counts = parsed.with_program(|_gc, program| {
59        let mut hist = Histogram {
60            counts: HashMap::new(),
61        };
62        hist.visit_node(program);
63        hist.counts
64    });
65
66    let mut rows: Vec<(NodeKind, usize)> = counts.into_iter().collect();
67    // Most frequent first; ties broken by kind name for a stable listing.
68    rows.sort_by(|a, b| {
69        b.1.cmp(&a.1)
70            .then_with(|| format!("{:?}", a.0).cmp(&format!("{:?}", b.0)))
71    });
72
73    let total: usize = rows.iter().map(|(_, n)| n).sum();
74    println!("{total} nodes, {} distinct kinds", rows.len());
75    for (kind, n) in rows {
76        println!("{:<24} {:>3} {}", format!("{kind:?}"), n, "#".repeat(n));
77    }
78}