reth_trie_sparse/arena/
nodes.rs

1use super::{
2    branch_child_idx::{BranchChildIdx, BranchChildIter},
3    ArenaSparseSubtrie, BranchChild, Index, NodeArena,
4};
5use alloc::{boxed::Box, vec::Vec};
6use alloy_primitives::{keccak256, B256};
7use alloy_trie::{BranchNodeCompact, TrieMask};
8use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNodeV2, RlpNode, TrieNodeV2};
9use smallvec::SmallVec;
10use strum::AsRefStr;
11
12use crate::TrieNodeEpoch;
13
14/// Tracks whether a node's RLP encoding is cached or needs recomputation.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub(super) enum ArenaSparseNodeState {
17    /// The node has been revealed but its RLP encoding is not cached.
18    Revealed,
19    /// The node has a cached RLP encoding that is still valid.
20    Cached {
21        /// The cached RLP-encoded representation of the node.
22        rlp_node: RlpNode,
23        /// The newest tracked modification epoch for this node or its descendants.
24        epoch: TrieNodeEpoch,
25    },
26    /// The node has been modified and its RLP encoding needs recomputation.
27    Dirty,
28}
29
30impl ArenaSparseNodeState {
31    /// Converts into a [`Self::Dirty`] if it's not already.
32    pub(super) const fn to_dirty(&self) -> Self {
33        Self::Dirty
34    }
35
36    /// Returns the [`RlpNode`] cached on the state, if there is one.
37    pub(super) const fn cached_rlp_node(&self) -> Option<&RlpNode> {
38        match self {
39            Self::Cached { rlp_node, .. } => Some(rlp_node),
40            _ => None,
41        }
42    }
43
44    /// Returns the cached epoch, if this node is cached.
45    pub(super) const fn cached_epoch(&self) -> Option<TrieNodeEpoch> {
46        match self {
47            Self::Cached { epoch, .. } => Some(*epoch),
48            _ => None,
49        }
50    }
51}
52
53/// The branch-specific data stored in an [`ArenaSparseNode::Branch`].
54#[derive(Debug, Clone)]
55pub(super) struct ArenaSparseNodeBranch {
56    /// Cached or dirty state of this node.
57    pub(super) state: ArenaSparseNodeState,
58    /// Revealed or blinded children, packed densely. The `state_mask` tracks which
59    /// nibble positions have entries in this `SmallVec`, which is sized to hold the maximum
60    /// of 16 children inline and therefore never spills.
61    pub(super) children: SmallVec<[BranchChild; 16]>,
62    /// Bitmask indicating which of the 16 child slots are occupied (have an entry
63    /// in `children`).
64    pub(super) state_mask: TrieMask,
65    /// Tree mask and hash mask for database persistence (`TrieUpdates`).
66    pub(super) branch_masks: BranchNodeMasks,
67}
68
69impl ArenaSparseNodeBranch {
70    /// Unsets the bit for `nibble` in `state_mask`, `hash_mask`, and `tree_mask`.
71    pub(super) const fn unset_child_bit(&mut self, nibble: u8) {
72        self.state_mask.unset_bit(nibble);
73    }
74
75    /// Inserts a child at `nibble`, updating the state mask, children array, and marking the
76    /// branch as dirty.
77    pub(super) fn set_child(&mut self, nibble: u8, child: BranchChild) {
78        let insert_pos = BranchChildIdx::insertion_point(self.state_mask, nibble);
79        self.state_mask.set_bit(nibble);
80        self.children.insert(insert_pos.get(), child);
81        self.state = ArenaSparseNodeState::Dirty;
82    }
83
84    /// Removes the child at `nibble`, updating the state mask, children array, and marking the
85    /// branch as dirty.
86    ///
87    /// # Panics
88    ///
89    /// Panics if `nibble` is not set in the state mask.
90    pub(super) fn remove_child(&mut self, nibble: u8) {
91        let child_idx =
92            BranchChildIdx::new(self.state_mask, nibble).expect("nibble not found in state_mask");
93        self.children.remove(child_idx.get());
94        self.unset_child_bit(nibble);
95        self.state = ArenaSparseNodeState::Dirty;
96    }
97
98    /// Returns the sibling child in a branch with exactly 2 children.
99    ///
100    /// # Panics
101    ///
102    /// Panics (debug) if the branch does not have exactly 2 children, or if `nibble` is not set.
103    pub(super) fn sibling_child(&self, nibble: u8) -> BranchChild {
104        debug_assert_eq!(
105            self.state_mask.count_bits(),
106            2,
107            "sibling_child requires exactly 2 children"
108        );
109        let child_idx =
110            BranchChildIdx::new(self.state_mask, nibble).expect("nibble not found in state_mask");
111        // With exactly 2 children the dense array has indices 0 and 1.
112        self.children[1 - child_idx.get()]
113    }
114
115    /// Iterates over `(nibble, child)` pairs in nibble order.
116    pub(super) fn child_iter(&self) -> impl Iterator<Item = (u8, BranchChild)> + '_ {
117        BranchChildIter::new(self.state_mask).map(|(idx, nibble)| (nibble, self.children[idx]))
118    }
119
120    /// Returns a [`BranchNodeCompact`] from this branch's masks and children hashes.
121    pub(super) fn branch_node_compact(&self, arena: &NodeArena) -> BranchNodeCompact {
122        let mut hashes = Vec::with_capacity(self.branch_masks.hash_mask.count_bits() as usize);
123        for (nibble, child) in self.child_iter() {
124            if self.branch_masks.hash_mask.is_bit_set(nibble) {
125                let hash = match child.revealed_index() {
126                    Some(child_idx) => arena[child_idx].cached_hash(),
127                    None => arena.blinded(child).as_hash().expect("blinded child must be a hash"),
128                };
129                hashes.push(hash);
130            }
131        }
132        BranchNodeCompact::new(
133            self.state_mask,
134            self.branch_masks.tree_mask,
135            self.branch_masks.hash_mask,
136            hashes,
137            None,
138        )
139    }
140}
141
142/// A node in the arena-based sparse trie.
143#[derive(Debug, Clone, AsRefStr)]
144pub(super) enum ArenaSparseNode {
145    /// Indicates a trie with no nodes.
146    EmptyRoot {
147        /// Cached or dirty state of this node. Its cached RLP is always the empty root hash.
148        state: ArenaSparseNodeState,
149    },
150    /// A branch node with up to 16 children.
151    Branch(ArenaSparseNodeBranch),
152    /// An extension node with a revealed branch child.
153    Extension {
154        /// The extension key.
155        key: Nibbles,
156        /// The extension's branch child in the arena.
157        child: Index,
158    },
159    /// A leaf node containing a value.
160    Leaf {
161        /// Cached or dirty state of this node.
162        state: ArenaSparseNodeState,
163        /// The RLP-encoded leaf value.
164        value: Vec<u8>,
165        /// The remaining key suffix for this leaf.
166        key: Nibbles,
167    },
168    /// A subtrie that can be taken for parallel processing.
169    Subtrie(Box<ArenaSparseSubtrie>),
170    /// Placeholder for a subtrie that has been temporarily taken for parallel operations.
171    TakenSubtrie,
172    /// Tombstone for a vacated slot. [`NodeArena::remove`](super::NodeArena::remove) also puts
173    /// the slot on the arena's free list, [`NodeArena::drain_node`](super::NodeArena::drain_node)
174    /// does not.
175    Free,
176}
177
178impl ArenaSparseNode {
179    /// Returns the state of an `EmptyRoot`, `Branch`, `Leaf`, or `Subtrie` root node, or `None` for
180    /// other types.
181    pub(super) fn state_ref(&self) -> Option<&ArenaSparseNodeState> {
182        match self {
183            Self::EmptyRoot { state } | Self::Leaf { state, .. } => Some(state),
184            Self::Branch(b) => Some(&b.state),
185            Self::Subtrie(s) => {
186                let root = &s.arena[s.root];
187                root.extension_child()
188                    .and_then(|child| s.arena[child].state_ref())
189                    .or_else(|| root.state_ref())
190            }
191            _ => None,
192        }
193    }
194
195    /// Returns a mutable reference to the state of an `EmptyRoot`, `Branch`, or `Leaf` node.
196    ///
197    /// # Panics
198    ///
199    /// Panics if called on a non-EmptyRoot/Branch/Leaf node.
200    pub(super) fn state_mut(&mut self) -> &mut ArenaSparseNodeState {
201        match self {
202            Self::EmptyRoot { state } | Self::Leaf { state, .. } => state,
203            Self::Branch(b) => &mut b.state,
204            _ => panic!("state_mut called on non-EmptyRoot/Branch/Leaf node"),
205        }
206    }
207
208    /// Returns `true` if this node's RLP encoding is cached.
209    pub(super) fn is_cached(&self) -> bool {
210        self.state_ref().is_some_and(|s| matches!(s, ArenaSparseNodeState::Cached { .. }))
211    }
212
213    /// Returns the short key of the extension or leaf, or None.
214    pub(super) const fn short_key(&self) -> Option<&Nibbles> {
215        match self {
216            Self::Extension { key, .. } | Self::Leaf { key, .. } => Some(key),
217            _ => None,
218        }
219    }
220
221    /// Returns the extension key, if this is an extension node.
222    pub(super) const fn extension_key(&self) -> Option<&Nibbles> {
223        match self {
224            Self::Extension { key, .. } => Some(key),
225            _ => None,
226        }
227    }
228
229    /// Returns the branch child of this extension node, if any.
230    pub(super) const fn extension_child(&self) -> Option<Index> {
231        match self {
232            Self::Extension { child, .. } => Some(*child),
233            _ => None,
234        }
235    }
236
237    /// Returns a reference to the branch data.
238    ///
239    /// # Panics
240    ///
241    /// Panics if this is not a `Branch` node.
242    pub(super) fn branch_ref(&self) -> &ArenaSparseNodeBranch {
243        match self {
244            Self::Branch(b) => b,
245            _ => panic!("branch_ref called on non-Branch node {self:?}"),
246        }
247    }
248
249    /// Returns a mutable reference to the branch data.
250    ///
251    /// # Panics
252    ///
253    /// Panics if this is not a `Branch` node.
254    pub(super) fn branch_mut(&mut self) -> &mut ArenaSparseNodeBranch {
255        match self {
256            Self::Branch(b) => b,
257            _ => panic!("branch_mut called on non-Branch node {self:?}"),
258        }
259    }
260
261    /// Returns a reference to the subtrie if this is a `Subtrie` node, or `None`.
262    #[cfg(all(debug_assertions, not(feature = "fuzz")))]
263    pub(super) const fn as_subtrie(&self) -> Option<&ArenaSparseSubtrie> {
264        match self {
265            Self::Subtrie(s) => Some(s),
266            _ => None,
267        }
268    }
269
270    /// Returns `true` if this node should contribute a set bit in its parent's `hash_mask`.
271    ///
272    /// That is, if the node is a branch without a parent extension whose cached RLP is a hash
273    /// (>= 32 bytes). Small branches whose RLP is embedded don't get a `hash_mask` bit.
274    pub(super) fn hash_mask_bit(&self) -> bool {
275        match self {
276            Self::Branch(b) => {
277                b.state.cached_rlp_node().expect("branch's RlpNode must be cached").is_hash()
278            }
279            Self::Subtrie(s) => s.arena[s.root].hash_mask_bit(),
280            _ => false,
281        }
282    }
283
284    /// Returns the cached hash of this node. Panics if the node's state is not `Cached`.
285    ///
286    /// If the `RlpNode` is already a hash (>= 32 bytes encoded), returns it directly.
287    /// Otherwise keccak-hashes the RLP encoding to produce the hash. This handles the
288    /// case where a branch's RLP is small enough to be embedded rather than hashed.
289    pub(super) fn cached_hash(&self) -> B256 {
290        let rlp_node = match self {
291            Self::Branch(ArenaSparseNodeBranch { state, .. }) | Self::Leaf { state, .. } => state
292                .cached_rlp_node()
293                .expect("cached_hash called on non-Cached branch or leaf: {self:?}"),
294            Self::Subtrie(s) => return s.arena[s.root].cached_hash(),
295            _ => panic!("cached_hash called on {self:?}"),
296        };
297        rlp_node.as_hash().unwrap_or_else(|| keccak256(rlp_node.as_slice()))
298    }
299}
300
301impl ArenaSparseNode {
302    /// Converts a [`ProofTrieNodeV2`] into an [`ArenaSparseNode`], storing the RLP of a branch's
303    /// unrevealed children in `arena`'s blinded side table.
304    ///
305    /// # Panics
306    ///
307    /// Panics if the node is an `Extension`, which should have been merged into a branch
308    /// by [`TrieNodeV2`].
309    pub(super) fn from_proof_node(arena: &mut NodeArena, proof_node: ProofTrieNodeV2) -> Self {
310        let ProofTrieNodeV2 { node, masks, .. } = proof_node;
311        match node {
312            TrieNodeV2::EmptyRoot => Self::EmptyRoot { state: ArenaSparseNodeState::Revealed },
313            TrieNodeV2::Leaf(leaf) => Self::Leaf {
314                state: ArenaSparseNodeState::Revealed,
315                key: leaf.key,
316                value: leaf.value,
317            },
318            TrieNodeV2::Branch(branch) => {
319                assert!(branch.key.is_empty(), "branches with a short key must be split first");
320                let children = branch.stack[..branch.state_mask.count_bits() as usize]
321                    .iter()
322                    .map(|rlp| arena.insert_blinded(rlp.clone()))
323                    .collect();
324                Self::Branch(ArenaSparseNodeBranch {
325                    state: ArenaSparseNodeState::Revealed,
326                    children,
327                    state_mask: branch.state_mask,
328                    branch_masks: masks.unwrap_or_default(),
329                })
330            }
331            TrieNodeV2::Extension(_) => {
332                panic!("Extension nodes should be merged into branches by TrieNodeV2")
333            }
334        }
335    }
336}
reth_trie_sparse/arena/nodes.rs

reth_trie_sparse/arena/
nodes.rs
