reth_trie_sparse/arena/
nodes.rs

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

reth_trie_sparse/arena/
nodes.rs
