turbopack_ecmascript/
ast_path_trie.rs

1//! A compact, shared representation of [`AstParentKind`] paths.
2//!
3//! Code generation identifies the AST node it wants to patch by the path of
4//! [`AstParentKind`]s leading to it from the program root. Stored flat, those paths can be
5//! quadratic in expression nesting depth: a left-leaning `a() + b() + c() + ...` chain of
6//! `N` terms produces `N` paths of lengths `2, 4, ..., 2N`, because every additional term
7//! nests one level deeper. Pathological, often generated, code can hit this.  This module removes
8//! the redundancy by storing
9//!
10//! Even in typical code shared prefixes are common:
11//! ```text
12//! import {a,b} from '...'
13//! ...
14//! export function foo() {
15//!   if (a()) return b();
16//! }
17//! ```
18//!
19//! Silly but the binding usage location for `a` and `b` share a common prefix.  That prefix can get
20//! arbitrarily long and this structure fixes.
21//!
22//! Paths are added through [`AstPathTrieBuilder`], which holds the index needed to
23//! deduplicate them, and it is then frozen into an immutable [`AstPathTrie`]. Only
24//! requested paths are interned, so the trie stays sparse: it holds the handful of nodes
25//! leading to code-generated locations, not every node in the file.
26
27use auto_hash_map::AutoMap;
28use bincode::{Decode, Encode};
29use rustc_hash::FxBuildHasher;
30use swc_core::ecma::visit::AstParentKind;
31use turbo_tasks::NonLocalValue;
32
33/// A reference to a path interned in an [`AstPathTrie`].
34///
35/// Only meaningful together with the trie that produced it.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode)]
37pub struct AstPathId(u32);
38
39impl AstPathId {
40    /// The empty path, i.e. the program root. Present in every trie without interning.
41    pub const ROOT: AstPathId = AstPathId(0);
42
43    pub fn is_root(self) -> bool {
44        self == Self::ROOT
45    }
46
47    /// Index into the node list. The root is implicit, so ids are offset by one.
48    fn index(self) -> usize {
49        debug_assert!(!self.is_root(), "the root has no node");
50        self.0 as usize - 1
51    }
52}
53
54// `AstPathId` is a plain index, so it holds no `Vc`s to trace and is trivially non-local.
55unsafe impl NonLocalValue for AstPathId {}
56
57#[derive(Debug, Clone, PartialEq, Eq, Hash, Encode, Decode)]
58struct Node {
59    parent: AstPathId,
60    #[bincode(with_serde)]
61    kind: AstParentKind,
62}
63
64/// Accumulates [`AstParentKind`] paths, sharing their common prefixes.
65///
66/// Holds the child index that deduplicating requires; [`AstPathTrieBuilder::build`] drops
67/// it and yields the immutable [`AstPathTrie`] that code generation reads.
68#[derive(Debug)]
69pub struct AstPathTrieBuilder {
70    nodes: Vec<Node>,
71    /// The children of each node, keyed by the kind that reaches them, so interning the
72    /// same path twice yields the same id. Effects are processed out of AST order, so there
73    /// is no walk position to extend and children have to be looked up.
74    ///
75    /// Indexed the same way as `nodes`, with slot 0 holding the root's children. Almost
76    /// every node has one or two children, so `AutoMap` keeps those inline and only spills
77    /// to a hash map for the occasional wide node (a module body, say).
78    children: Vec<AutoMap<AstParentKind, AstPathId, FxBuildHasher, 4>>,
79}
80
81impl Default for AstPathTrieBuilder {
82    fn default() -> Self {
83        Self {
84            nodes: Vec::new(),
85            // Slot 0 holds the root's children.
86            children: vec![AutoMap::default()],
87        }
88    }
89}
90
91impl AstPathTrieBuilder {
92    pub fn new() -> Self {
93        Self::default()
94    }
95
96    /// Interns a path, returning a handle to its last element.
97    pub fn intern(&mut self, path: impl IntoIterator<Item = AstParentKind>) -> AstPathId {
98        let parent = AstPathId::ROOT;
99        let mut id = parent;
100        for kind in path {
101            id = self.push(id, kind);
102        }
103        id
104    }
105
106    /// Extends the path `parent` by one element.
107    fn push(&mut self, parent: AstPathId, kind: AstParentKind) -> AstPathId {
108        // The root's children live in slot 0, so every node's children are one slot past
109        // its id; that keeps the root from needing a case of its own.
110        let slot = parent.0 as usize;
111        if let Some(&existing) = self.children[slot].get(&kind) {
112            return existing;
113        }
114
115        let id = AstPathId(u32::try_from(self.nodes.len() + 1).expect("too many ast path nodes"));
116        self.nodes.push(Node { parent, kind });
117        self.children.push(AutoMap::default());
118        self.children[slot].insert(kind, id);
119        id
120    }
121
122    /// The path with its last element removed, or the root when already there.
123    pub fn parent_or_root(&self, id: AstPathId) -> AstPathId {
124        if id.is_root() {
125            AstPathId::ROOT
126        } else {
127            self.nodes[id.index()].parent
128        }
129    }
130
131    /// Freezes the paths interned so far.
132    pub fn build(self) -> AstPathTrie {
133        AstPathTrie {
134            nodes: self.nodes.into_boxed_slice(),
135        }
136    }
137
138    /// The number of interned nodes, excluding the implicit root.
139    pub fn node_count(&self) -> usize {
140        self.nodes.len()
141    }
142}
143
144/// An immutable arena of [`AstParentKind`] paths that share their common prefixes.
145///
146/// Built by [`AstPathTrieBuilder`] and read back via [`AstPathId`]. Node 0 is implicit and
147/// represents the empty path, so every trie contains [`AstPathId::ROOT`].
148#[derive(Debug, Default, Clone, PartialEq, Eq, Hash, Encode, Decode)]
149pub struct AstPathTrie {
150    nodes: Box<[Node]>,
151}
152
153unsafe impl NonLocalValue for AstPathTrie {}
154
155impl AstPathTrie {
156    fn node(&self, id: AstPathId) -> Option<&Node> {
157        if id.is_root() {
158            // `AstPathId::ROOT` is implicit and has no node.
159            return None;
160        }
161        Some(self.nodes.get(id.index()).unwrap_or_else(|| {
162            // An id only means anything against the trie that minted it. Reading one from a
163            // different trie is a wiring bug, so say that rather than index out of bounds.
164            panic!(
165                "{id:?} does not belong to this trie ({} nodes); it was interned into a different \
166                 one",
167                self.nodes.len(),
168            )
169        }))
170    }
171
172    /// The last element of the path, or `None` for the root.
173    pub fn get(&self, id: AstPathId) -> Option<AstParentKind> {
174        self.node(id).map(|n| n.kind)
175    }
176
177    /// The path with its last element removed, or `None` for the root.
178    pub fn get_parent(&self, id: AstPathId) -> Option<AstPathId> {
179        self.node(id).map(|n| n.parent)
180    }
181
182    /// The path with its last element removed, or the root when already there.
183    pub fn parent_or_root(&self, id: AstPathId) -> AstPathId {
184        self.get_parent(id).unwrap_or(AstPathId::ROOT)
185    }
186
187    /// The last element of the path together with the path below it, or `None` for the
188    /// root. One node read where [`AstPathTrie::get`] plus [`AstPathTrie::get_parent`]
189    /// would be two.
190    pub fn split_last(&self, id: AstPathId) -> Option<(AstParentKind, AstPathId)> {
191        self.node(id).map(|n| (n.kind, n.parent))
192    }
193
194    /// Walks from `id` towards the root, yielding each element in reverse order.
195    pub fn iter_rev(&self, id: AstPathId) -> impl Iterator<Item = AstParentKind> + '_ {
196        let mut current = id;
197        std::iter::from_fn(move || {
198            let node = self.node(current)?;
199            current = node.parent;
200            Some(node.kind)
201        })
202    }
203
204    /// The innermost ancestor of `id` (or `id` itself) whose last element matches `f`.
205    ///
206    /// Returns `None` when nothing up to the root matches. Callers wanting the node *above*
207    /// the match can follow with [`AstPathTrie::parent_or_root`].
208    pub fn find_last(
209        &self,
210        id: AstPathId,
211        mut f: impl FnMut(&AstParentKind) -> bool,
212    ) -> Option<AstPathId> {
213        let mut current = id;
214        while let Some(node) = self.node(current) {
215            if f(&node.kind) {
216                return Some(current);
217            }
218            current = node.parent;
219        }
220        None
221    }
222
223    /// The number of interned nodes, excluding the implicit root.
224    pub fn node_count(&self) -> usize {
225        self.nodes.len()
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use swc_core::ecma::visit::{
232        AstParentKind,
233        fields::{BinExprField, ExprField, ModuleField, ModuleItemField},
234    };
235
236    use super::*;
237
238    fn expr_bin() -> AstParentKind {
239        AstParentKind::Expr(ExprField::Bin)
240    }
241    fn bin_left() -> AstParentKind {
242        AstParentKind::BinExpr(BinExprField::Left)
243    }
244    fn module_body(i: usize) -> AstParentKind {
245        AstParentKind::Module(ModuleField::Body(i))
246    }
247    fn module_item() -> AstParentKind {
248        AstParentKind::ModuleItem(ModuleItemField::Stmt)
249    }
250
251    /// Interns one path and freezes the trie.
252    fn trie_of(path: &[AstParentKind]) -> (AstPathTrie, AstPathId) {
253        let mut builder = AstPathTrieBuilder::new();
254        let id = builder.intern(path.iter().copied());
255        (builder.build(), id)
256    }
257
258    fn to_vec(trie: &AstPathTrie, id: AstPathId) -> Vec<AstParentKind> {
259        let mut vec: Vec<_> = trie.iter_rev(id).collect();
260        vec.reverse();
261        vec
262    }
263
264    #[test]
265    fn root_is_empty() {
266        let trie = AstPathTrieBuilder::new().build();
267        assert_eq!(trie.get(AstPathId::ROOT), None);
268        assert_eq!(trie.get_parent(AstPathId::ROOT), None);
269        assert_eq!(to_vec(&trie, AstPathId::ROOT), vec![]);
270        assert_eq!(trie.node_count(), 0);
271    }
272
273    #[test]
274    fn roundtrips_a_path() {
275        let path = vec![module_body(3), module_item(), expr_bin()];
276        let (trie, id) = trie_of(&path);
277        assert_eq!(to_vec(&trie, id), path);
278        assert_eq!(trie.get(id), Some(expr_bin()));
279    }
280
281    #[test]
282    fn interning_is_idempotent() {
283        let mut builder = AstPathTrieBuilder::new();
284        let path = [module_body(0), module_item()];
285        assert_eq!(
286            builder.intern(path.iter().copied()),
287            builder.intern(path.iter().copied()),
288        );
289        assert_eq!(builder.node_count(), 2);
290    }
291
292    #[test]
293    fn shares_prefixes() {
294        let mut builder = AstPathTrieBuilder::new();
295        let a = builder.intern([module_body(0), module_item(), expr_bin()]);
296        let b = builder.intern([module_body(0), module_item(), bin_left()]);
297        // The 2-element common prefix is stored once; only the last elements differ.
298        assert_eq!(builder.node_count(), 4);
299        let trie = builder.build();
300        assert_eq!(trie.get_parent(a), trie.get_parent(b));
301        assert_ne!(a, b);
302    }
303
304    #[test]
305    fn distinct_indices_are_distinct_nodes() {
306        let mut builder = AstPathTrieBuilder::new();
307        let a = builder.intern([module_body(0)]);
308        let b = builder.intern([module_body(1)]);
309        assert_ne!(a, b, "Body(0) and Body(1) address different children");
310    }
311
312    /// The motivating case: a left-leaning `a + b + c + ...` chain. Flat storage is
313    /// quadratic in the number of terms; the trie must stay linear.
314    #[test]
315    fn left_leaning_chain_is_linear() {
316        const TERMS: usize = 500;
317        let mut builder = AstPathTrieBuilder::new();
318
319        // Build the spine, interning the path to each term as we descend.
320        let mut spine = Vec::new();
321        let mut flat_elements = 0;
322        for _ in 0..TERMS {
323            spine.push(expr_bin());
324            spine.push(bin_left());
325            builder.intern(spine.iter().copied());
326            flat_elements += spine.len();
327        }
328
329        assert_eq!(builder.node_count(), TERMS * 2);
330        assert_eq!(flat_elements, TERMS * (TERMS + 1));
331        assert!(
332            (builder.node_count() as f64) < (flat_elements as f64) / 100.0,
333            "expected a large saving, got {} nodes for {flat_elements} elements",
334            builder.node_count(),
335        );
336    }
337
338    #[test]
339    fn walks_up_from_a_leaf() {
340        let (trie, id) = trie_of(&[module_body(0), module_item(), expr_bin()]);
341
342        let parent = trie.get_parent(id).unwrap();
343        assert_eq!(to_vec(&trie, parent), vec![module_body(0), module_item()]);
344        assert_eq!(
345            trie.iter_rev(id).collect::<Vec<_>>(),
346            vec![expr_bin(), module_item(), module_body(0)],
347        );
348    }
349
350    #[test]
351    fn finds_the_innermost_match() {
352        let (trie, id) = trie_of(&[module_body(0), expr_bin(), bin_left(), bin_left()]);
353
354        // The innermost `Expr`, skipping the two `BinExpr`s below it.
355        let found = trie
356            .find_last(id, |k| matches!(k, AstParentKind::Expr(_)))
357            .unwrap();
358        assert_eq!(to_vec(&trie, found), vec![module_body(0), expr_bin()]);
359
360        // `id` itself can be the match.
361        assert_eq!(trie.find_last(id, |_| true), Some(id));
362        // Nothing matches.
363        assert_eq!(trie.find_last(id, |_| false), None);
364    }
365
366    #[test]
367    fn encodes_and_decodes() {
368        let (trie, id) = trie_of(&[module_body(0), module_item(), expr_bin()]);
369
370        let config = bincode::config::standard();
371        let bytes = bincode::encode_to_vec(&trie, config).unwrap();
372        let (decoded, _): (AstPathTrie, _) = bincode::decode_from_slice(&bytes, config).unwrap();
373
374        assert_eq!(decoded, trie);
375        assert_eq!(to_vec(&decoded, id), to_vec(&trie, id));
376        assert_eq!(decoded.node_count(), trie.node_count());
377    }
378}
turbopack_ecmascript/ast_path_trie.rs

turbopack_ecmascript/
ast_path_trie.rs
