miden_rowan/green/
node_cache.rs

1use alloc::vec::Vec;
2use core::hash::{BuildHasherDefault, Hash, Hasher};
3use hashbrown::hash_map::RawEntryMut;
4use rustc_hash::FxHasher;
5
6use crate::{
7    GreenNode, GreenNodeData, GreenToken, GreenTokenData, NodeOrToken, SyntaxKind,
8    green::GreenElementRef,
9};
10
11use super::element::GreenElement;
12
13type HashMap<K, V> = hashbrown::HashMap<K, V, BuildHasherDefault<FxHasher>>;
14
15#[derive(Debug)]
16struct NoHash<T>(T);
17
18/// Interner for GreenTokens and GreenNodes
19// XXX: the impl is a bit tricky. As usual when writing interners, we want to
20// store all values in one HashSet.
21//
22// However, hashing trees is fun: hash of the tree is recursively defined. We
23// maintain an invariant -- if the tree is interned, then all of its children
24// are interned as well.
25//
26// That means that computing the hash naively is wasteful -- we just *know*
27// hashes of children, and we can re-use those.
28//
29// So here we use *raw* API of hashbrown and provide the hashes manually,
30// instead of going via a `Hash` impl. Our manual `Hash` and the
31// `#[derive(Hash)]` are actually different! At some point we had a fun bug,
32// where we accidentally mixed the two hashes, which made the cache much less
33// efficient.
34//
35// To fix that, we additionally wrap the data in `NoHash` wrapper, to make sure
36// we don't accidentally use the wrong hash!
37#[derive(Default, Debug)]
38pub struct NodeCache {
39    nodes: HashMap<NoHash<GreenNode>, ()>,
40    tokens: HashMap<NoHash<GreenToken>, ()>,
41}
42
43fn token_hash(token: &GreenTokenData) -> u64 {
44    let mut h = FxHasher::default();
45    token.kind().hash(&mut h);
46    token.text().hash(&mut h);
47    h.finish()
48}
49
50fn node_hash(node: &GreenNodeData) -> u64 {
51    let mut h = FxHasher::default();
52    node.kind().hash(&mut h);
53    for child in node.children() {
54        match child {
55            NodeOrToken::Node(it) => node_hash(it),
56            NodeOrToken::Token(it) => token_hash(it),
57        }
58        .hash(&mut h)
59    }
60    h.finish()
61}
62
63fn element_id(elem: GreenElementRef<'_>) -> *const () {
64    match elem {
65        NodeOrToken::Node(it) => it as *const GreenNodeData as *const (),
66        NodeOrToken::Token(it) => it as *const GreenTokenData as *const (),
67    }
68}
69
70impl NodeCache {
71    pub(crate) fn node(
72        &mut self,
73        kind: SyntaxKind,
74        children: &mut Vec<(u64, GreenElement)>,
75        first_child: usize,
76    ) -> (u64, GreenNode) {
77        let build_node = move |children: &mut Vec<(u64, GreenElement)>| {
78            GreenNode::new(kind, children.drain(first_child..).map(|(_, it)| it))
79        };
80
81        let children_ref = &children[first_child..];
82        if children_ref.len() > 3 {
83            let node = build_node(children);
84            return (0, node);
85        }
86
87        let hash = {
88            let mut h = FxHasher::default();
89            kind.hash(&mut h);
90            for &(hash, _) in children_ref {
91                if hash == 0 {
92                    let node = build_node(children);
93                    return (0, node);
94                }
95                hash.hash(&mut h);
96            }
97            h.finish()
98        };
99
100        // Green nodes are fully immutable, so it's ok to deduplicate them.
101        // This is the same optimization that Roslyn does
102        // https://github.com/KirillOsenkov/Bliki/wiki/Roslyn-Immutable-Trees
103        //
104        // For example, all `#[inline]` in this file share the same green node!
105        // For `libsyntax/parse/parser.rs`, measurements show that deduping saves
106        // 17% of the memory for green nodes!
107        let entry = self.nodes.raw_entry_mut().from_hash(hash, |node| {
108            node.0.kind() == kind && node.0.children().len() == children_ref.len() && {
109                let lhs = node.0.children();
110                let rhs = children_ref.iter().map(|(_, it)| it.as_deref());
111
112                let lhs = lhs.map(element_id);
113                let rhs = rhs.map(element_id);
114
115                lhs.eq(rhs)
116            }
117        });
118
119        let node = match entry {
120            RawEntryMut::Occupied(entry) => {
121                drop(children.drain(first_child..));
122                entry.key().0.clone()
123            }
124            RawEntryMut::Vacant(entry) => {
125                let node = build_node(children);
126                entry.insert_with_hasher(hash, NoHash(node.clone()), (), |n| node_hash(&n.0));
127                node
128            }
129        };
130
131        (hash, node)
132    }
133
134    pub(crate) fn token(&mut self, kind: SyntaxKind, text: &str) -> (u64, GreenToken) {
135        let hash = {
136            let mut h = FxHasher::default();
137            kind.hash(&mut h);
138            text.hash(&mut h);
139            h.finish()
140        };
141
142        let entry = self
143            .tokens
144            .raw_entry_mut()
145            .from_hash(hash, |token| token.0.kind() == kind && token.0.text() == text);
146
147        let token = match entry {
148            RawEntryMut::Occupied(entry) => entry.key().0.clone(),
149            RawEntryMut::Vacant(entry) => {
150                let token = GreenToken::new(kind, text);
151                entry.insert_with_hasher(hash, NoHash(token.clone()), (), |t| token_hash(&t.0));
152                token
153            }
154        };
155
156        (hash, token)
157    }
158}
miden_rowan/green/node_cache.rs

miden_rowan/green/
node_cache.rs
