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

reth_trie_sparse/arena/
nodes.rs
