reth_trie/proof_v2/node.rs
1use crate::proof_v2::DeferredValueEncoder;
2use alloy_rlp::Encodable;
3use reth_execution_errors::trie::StateProofError;
4use reth_trie_common::{
5 BranchNodeMasks, BranchNodeV2, LeafNode, LeafNodeRef, Nibbles, ProofTrieNodeV2, RlpNode,
6 TrieMask, TrieNodeV2,
7};
8
9/// A trie node which is the child of a branch in the trie.
10#[derive(Debug)]
11pub(crate) enum ProofTrieBranchChild<RF> {
12 /// A leaf node whose value has yet to be calculated and encoded.
13 Leaf {
14 /// The short key of the leaf.
15 short_key: Nibbles,
16 /// The [`DeferredValueEncoder`] which will encode the leaf's value.
17 value: RF,
18 },
19 /// A branch node whose children have already been flattened into [`RlpNode`]s.
20 Branch {
21 /// The node itself, for use during RLP encoding.
22 node: BranchNodeV2,
23 /// Bitmasks carried over from cached `BranchNodeCompact` values, if any.
24 masks: Option<BranchNodeMasks>,
25 },
26 /// A node whose type is not known, as it has already been converted to an [`RlpNode`].
27 RlpNode(RlpNode),
28}
29
30impl<RF: DeferredValueEncoder> ProofTrieBranchChild<RF> {
31 /// Converts this child into its RLP node representation.
32 ///
33 /// This potentially also returns an `RlpNode` buffer which can be re-used for other
34 /// [`ProofTrieBranchChild`]s.
35 pub(crate) fn into_rlp(
36 self,
37 buf: &mut Vec<u8>,
38 ) -> Result<(RlpNode, Option<Vec<RlpNode>>), StateProofError> {
39 match self {
40 Self::Leaf { short_key, value } => {
41 // RLP encode the value itself
42 value.encode(buf)?;
43 let value_enc_len = buf.len();
44
45 // Determine the required buffer size for the encoded leaf
46 let leaf_enc_len = LeafNodeRef::new(&short_key, buf).length();
47
48 // We want to re-use buf for the encoding of the leaf node as well. To do this we
49 // will keep appending to it, leaving the already encoded value in-place. First we
50 // must ensure the buffer is big enough, then we'll split.
51 buf.resize(value_enc_len + leaf_enc_len, 0);
52
53 // SAFETY we have just resized the above to be greater than `value_enc_len`, so it
54 // must be in-bounds.
55 let (value_buf, mut leaf_buf) =
56 unsafe { buf.split_at_mut_unchecked(value_enc_len) };
57
58 // Encode the leaf into the right side of the split buffer, and return the RlpNode.
59 LeafNodeRef::new(&short_key, value_buf).encode(&mut leaf_buf);
60 Ok((RlpNode::from_rlp(&buf[value_enc_len..]), None))
61 }
62 Self::Branch { node: branch_node, .. } => {
63 branch_node.encode(buf);
64 Ok((RlpNode::from_rlp(buf), Some(branch_node.stack)))
65 }
66 Self::RlpNode(rlp_node) => Ok((rlp_node, None)),
67 }
68 }
69
70 /// Converts this child into a [`ProofTrieNodeV2`] having the given path.
71 ///
72 /// # Panics
73 ///
74 /// If called on a [`Self::RlpNode`].
75 pub(crate) fn into_proof_trie_node(
76 self,
77 path: Nibbles,
78 buf: &mut Vec<u8>,
79 ) -> Result<ProofTrieNodeV2, StateProofError> {
80 let (node, masks) = match self {
81 Self::Leaf { short_key, value } => {
82 value.encode(buf)?;
83 // Counter-intuitively a clone is better here than a `core::mem::take`. If we take
84 // the buffer then future RLP-encodes will need to re-allocate a new one, and
85 // RLP-encodes after those may need a bigger buffer and therefore re-alloc again.
86 //
87 // By cloning here we do a single allocation of exactly the size we need to take
88 // this value, and the passed in buffer can remain with whatever large capacity it
89 // already has.
90 let rlp_val = buf.clone();
91 (TrieNodeV2::Leaf(LeafNode::new(short_key, rlp_val)), None)
92 }
93 Self::Branch { node, masks } => (TrieNodeV2::Branch(node), masks),
94 Self::RlpNode(_) => panic!("Cannot call `into_proof_trie_node` on RlpNode"),
95 };
96
97 Ok(ProofTrieNodeV2 { node, path, masks })
98 }
99
100 /// Returns the short key of the child, if it is a leaf or branch, or empty if its a
101 /// [`Self::RlpNode`].
102 pub(crate) fn short_key(&self) -> &Nibbles {
103 match self {
104 Self::Leaf { short_key, .. } |
105 Self::Branch { node: BranchNodeV2 { key: short_key, .. }, .. } => short_key,
106 Self::RlpNode(_) => {
107 static EMPTY_NIBBLES: Nibbles = Nibbles::new();
108 &EMPTY_NIBBLES
109 }
110 }
111 }
112
113 /// Trims the given number of nibbles off the head of the short key.
114 ///
115 /// If the node is an extension and the given length is the same as its short key length, then
116 /// the node is replaced with its child.
117 ///
118 /// # Panics
119 ///
120 /// - If the given len is longer than the short key
121 /// - If the given len is the same as the length of a leaf's short key
122 /// - If the node is a [`Self::Branch`] or [`Self::RlpNode`]
123 pub(crate) fn trim_short_key_prefix(&mut self, len: usize) {
124 match self {
125 Self::Leaf { short_key, .. } => {
126 *short_key = trim_nibbles_prefix(short_key, len);
127 }
128 Self::Branch { node: BranchNodeV2 { key, branch_rlp_node, .. }, .. } => {
129 *key = trim_nibbles_prefix(key, len);
130 if key.is_empty() {
131 *branch_rlp_node = None;
132 }
133 }
134 Self::RlpNode(_) => {
135 panic!("Cannot call `trim_short_key_prefix` on RlpNode")
136 }
137 }
138 }
139}
140
141/// A single branch in the trie which is under construction. The actual child nodes of the branch
142/// will be tracked as [`ProofTrieBranchChild`]s on a stack.
143#[derive(Debug)]
144pub(crate) struct ProofTrieBranch {
145 /// The length of the parent extension node's short key. If zero then the branch's parent is
146 /// not an extension but instead another branch.
147 pub(crate) ext_len: u8,
148 /// A mask tracking which child nibbles are set on the branch so far. There will be a single
149 /// child on the stack for each set bit.
150 pub(crate) state_mask: TrieMask,
151 /// Bitmasks which are subsets of `state_mask`.
152 pub(crate) masks: Option<BranchNodeMasks>,
153}
154
155/// Trims the first `len` nibbles from the head of the given `Nibbles`.
156///
157/// # Panics
158///
159/// Panics if the given `len` is greater than the length of the `Nibbles`.
160pub(crate) fn trim_nibbles_prefix(n: &Nibbles, len: usize) -> Nibbles {
161 debug_assert!(n.len() >= len);
162 n.slice_unchecked(len, n.len())
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168
169 #[test]
170 fn test_trim_nibbles_prefix_basic() {
171 // Create nibbles [1, 2, 3, 4, 5, 6]
172 let nibbles = Nibbles::from_nibbles([1, 2, 3, 4, 5, 6]);
173
174 // Trim first 2 nibbles
175 let trimmed = trim_nibbles_prefix(&nibbles, 2);
176 assert_eq!(trimmed.len(), 4);
177
178 // Verify the remaining nibbles are [3, 4, 5, 6]
179 assert_eq!(trimmed.get(0), Some(3));
180 assert_eq!(trimmed.get(1), Some(4));
181 assert_eq!(trimmed.get(2), Some(5));
182 assert_eq!(trimmed.get(3), Some(6));
183 }
184
185 #[test]
186 fn test_trim_nibbles_prefix_zero() {
187 // Create nibbles [10, 11, 12, 13]
188 let nibbles = Nibbles::from_nibbles([10, 11, 12, 13]);
189
190 // Trim zero nibbles - should return identical nibbles
191 let trimmed = trim_nibbles_prefix(&nibbles, 0);
192 assert_eq!(trimmed, nibbles);
193 }
194
195 #[test]
196 fn test_trim_nibbles_prefix_all() {
197 // Create nibbles [1, 2, 3, 4]
198 let nibbles = Nibbles::from_nibbles([1, 2, 3, 4]);
199
200 // Trim all nibbles - should return empty
201 let trimmed = trim_nibbles_prefix(&nibbles, 4);
202 assert!(trimmed.is_empty());
203 }
204
205 #[test]
206 fn test_trim_nibbles_prefix_empty() {
207 // Create empty nibbles
208 let nibbles = Nibbles::new();
209
210 // Trim zero from empty - should return empty
211 let trimmed = trim_nibbles_prefix(&nibbles, 0);
212 assert!(trimmed.is_empty());
213 }
214}