ethrex_trie/node/
leaf.rs

1#[cfg(not(feature = "std"))]
2use alloc::{string::ToString, vec::Vec};
3use core::mem;
4
5use ethrex_crypto::Crypto;
6use ethrex_rlp::encode::RLPEncode;
7
8use crate::{
9    ValueRLP,
10    error::TrieError,
11    nibbles::Nibbles,
12    node::{BranchNode, NodeRemoveResult},
13    node_hash::NodeHash,
14};
15
16use super::{ExtensionNode, Node, ValueOrHash};
17/// Leaf Node of an an Ethereum Compatible Patricia Merkle Trie
18/// Contains the node's hash, value & path
19#[derive(
20    Debug,
21    Clone,
22    Default,
23    PartialEq,
24    Eq,
25    serde::Serialize,
26    serde::Deserialize,
27    rkyv::Deserialize,
28    rkyv::Serialize,
29    rkyv::Archive,
30)]
31pub struct LeafNode {
32    pub partial: Nibbles,
33    pub value: ValueRLP,
34}
35
36impl LeafNode {
37    /// Creates a new leaf node and stores the given (path, value) pair
38    pub const fn new(partial: Nibbles, value: ValueRLP) -> Self {
39        Self { partial, value }
40    }
41
42    /// Returns the stored value if the given path matches the stored path
43    pub fn get(&self, path: Nibbles) -> Result<Option<ValueRLP>, TrieError> {
44        if self.partial == path {
45            Ok(Some(self.value.clone()))
46        } else {
47            Ok(None)
48        }
49    }
50
51    /// Stores the received value and returns the new root of the subtrie previously consisting of self
52    pub fn insert(&mut self, path: Nibbles, value: ValueOrHash) -> Result<Option<Node>, TrieError> {
53        /* Possible flow paths:
54            Leaf { SelfValue } -> Leaf { Value }
55            Leaf { SelfValue } -> Extension { Branch { [Self,...] Value } }
56            Leaf { SelfValue } -> Extension { Branch { [ Leaf { Value } , ... ], SelfValue} }
57            Leaf { SelfValue } -> Branch { [ Leaf { Value }, Self, ... ], None}
58        */
59        // If the path matches the stored path, update the value and return self
60        if self.partial == path {
61            match value {
62                ValueOrHash::Value(value) => self.value = value,
63                ValueOrHash::Hash(_) => {
64                    return Err(TrieError::Verify(
65                        "attempt to override proof node with external hash".to_string(),
66                    ));
67                }
68            }
69            Ok(None)
70        } else {
71            let match_index = path.count_prefix(&self.partial);
72            let self_choice_idx = self.partial.at(match_index);
73            let new_leaf_choice_idx = path.at(match_index);
74            self.partial = self.partial.offset(match_index + 1);
75
76            let branch_node = if self_choice_idx == 16 {
77                // Create a new leaf node and store the value in it
78                // Create a new branch node with the leaf as a child and store self's value
79                // Branch { [ Leaf { Value } , ... ], SelfValue}
80                let mut choices = BranchNode::EMPTY_CHOICES;
81                choices[new_leaf_choice_idx] = match value {
82                    ValueOrHash::Value(value) => {
83                        Node::from(LeafNode::new(path.offset(match_index + 1), value)).into()
84                    }
85                    ValueOrHash::Hash(hash) => hash.into(),
86                };
87                BranchNode::new_with_value(choices, mem::take(&mut self.value))
88            } else if new_leaf_choice_idx == 16 {
89                // Create a branch node with self as a child and store the value in the branch node
90                // Branch { [Self,...], Value }
91                let mut choices = BranchNode::EMPTY_CHOICES;
92                let child: Node = self.take().into();
93                choices[self_choice_idx] = child.into();
94                BranchNode::new_with_value(
95                    choices,
96                    match value {
97                        ValueOrHash::Value(value) => value,
98                        // Value in branches don't happen in our use-case.
99                        ValueOrHash::Hash(_) => todo!("handle override case (error?)"),
100                    },
101                )
102            } else {
103                // Create a new leaf node and store the path and value in it
104                // Create a new branch node with the leaf and self as children
105                // Branch { [ Leaf { Path, Value }, Self, ... ], None, None}
106                let mut choices = BranchNode::EMPTY_CHOICES;
107                choices[new_leaf_choice_idx] = match value {
108                    ValueOrHash::Value(value) => {
109                        Node::from(LeafNode::new(path.offset(match_index + 1), value)).into()
110                    }
111                    ValueOrHash::Hash(hash) => hash.into(),
112                };
113                let child: Node = self.take().into();
114                choices[self_choice_idx] = child.into();
115                BranchNode::new(choices)
116            };
117
118            let final_node = if match_index == 0 {
119                branch_node.into()
120            } else {
121                // Create an extension node with the branch node as child
122                // Extension { BranchNode }
123                ExtensionNode::new(path.slice(0, match_index), Node::from(branch_node).into())
124                    .into()
125            };
126
127            Ok(Some(final_node))
128        }
129    }
130
131    /// Removes own value if the path matches own path and returns self and the value if it was removed
132    pub fn remove(
133        &mut self,
134        path: Nibbles,
135    ) -> Result<(Option<NodeRemoveResult>, Option<ValueRLP>), TrieError> {
136        Ok(if self.partial == path {
137            (None, Some(mem::take(&mut self.value)))
138        } else {
139            (Some(NodeRemoveResult::Mutated), None)
140        })
141    }
142
143    /// Computes the node's hash
144    pub fn compute_hash(&self, crypto: &dyn Crypto) -> NodeHash {
145        self.compute_hash_no_alloc(&mut vec![], crypto)
146    }
147
148    /// Computes the node's hash, using the provided buffer
149    pub fn compute_hash_no_alloc(&self, buf: &mut Vec<u8>, crypto: &dyn Crypto) -> NodeHash {
150        buf.clear();
151        self.encode(buf);
152        let hash = NodeHash::from_encoded(buf, crypto);
153        buf.clear();
154        hash
155    }
156
157    /// Encodes the node and appends it to `node_path` if the encoded node is 32 or more bytes long
158    pub fn get_path(&self, node_path: &mut Vec<Vec<u8>>) -> Result<(), TrieError> {
159        let encoded = self.encode_to_vec();
160        if encoded.len() >= 32 {
161            node_path.push(encoded);
162        }
163        Ok(())
164    }
165
166    /// Creates a new node by emptying `self` partial and its value
167    ///
168    /// This is a way to "consume" the node when we just have a mutable reference to it
169    pub fn take(&mut self) -> Self {
170        LeafNode {
171            partial: self.partial.take(),
172            value: mem::take(&mut self.value),
173        }
174    }
175}
176
177#[cfg(test)]
178mod test {
179    use ethrex_crypto::NativeCrypto;
180    use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
181
182    use super::*;
183    use crate::{Trie, pmt_node};
184
185    #[test]
186    fn new() {
187        let node = LeafNode::new(Default::default(), Default::default());
188        assert_eq!(node.value, ValueRLP::default());
189    }
190
191    #[test]
192    fn get_some() {
193        let node = pmt_node! { @(trie)
194            leaf { vec![1, 2, 16] => vec![0x12, 0x34, 0x56, 0x78] }
195        };
196
197        assert_eq!(
198            node.get(Nibbles::from_bytes(&[0x12])).unwrap(),
199            Some(vec![0x12, 0x34, 0x56, 0x78]),
200        );
201    }
202
203    #[test]
204    fn get_none() {
205        let node = pmt_node! { @(trie)
206            leaf { vec![1,2,16] => vec![0x12, 0x34, 0x56, 0x78] }
207        };
208
209        assert!(node.get(Nibbles::from_bytes(&[0x34])).unwrap().is_none());
210    }
211
212    #[test]
213    fn insert_replace() {
214        let mut node = pmt_node! { @(trie)
215            leaf { vec![1,2,16] => vec![0x12, 0x34, 0x56, 0x78] }
216        };
217
218        assert!(
219            node.insert(Nibbles::from_bytes(&[0x12]), vec![0x13].into())
220                .unwrap()
221                .is_none()
222        );
223
224        assert_eq!(node.value, vec![0x13]);
225    }
226
227    #[test]
228    fn insert_branch() {
229        let trie = Trie::new_temp();
230        let mut node = pmt_node! { @(trie)
231            leaf { vec![1,2,16] => vec![0x12, 0x34, 0x56, 0x78] }
232        };
233        let path = Nibbles::from_bytes(&[0x22]);
234        let value = vec![0x23];
235        let node = node.insert(path.clone(), value.clone().into()).unwrap();
236        let node = match node {
237            Some(Node::Branch(x)) => x,
238            _ => panic!("expected a branch node"),
239        };
240        assert_eq!(node.get(trie.db.as_ref(), path).unwrap(), Some(value));
241    }
242
243    #[test]
244    fn insert_extension_branch() {
245        let trie = Trie::new_temp();
246        let mut node = pmt_node! { @(trie)
247            leaf { vec![1,2,16] => vec![0x12, 0x34, 0x56, 0x78] }
248        };
249
250        let path = Nibbles::from_bytes(&[0x13]);
251        let value = vec![0x15];
252
253        let node = node.insert(path.clone(), value.clone().into()).unwrap();
254
255        assert!(matches!(node, Some(Node::Extension(_))));
256        assert_eq!(
257            node.unwrap().get(trie.db.as_ref(), path).unwrap(),
258            Some(value)
259        );
260    }
261
262    #[test]
263    fn insert_extension_branch_value_self() {
264        let trie = Trie::new_temp();
265        let mut node = pmt_node! { @(trie)
266            leaf { vec![1,2,16] => vec![0x12, 0x34, 0x56, 0x78] }
267        };
268
269        let path = Nibbles::from_bytes(&[0x12, 0x34]);
270        let value = vec![0x17];
271
272        let node = node.insert(path.clone(), value.clone().into()).unwrap();
273
274        assert!(matches!(node, Some(Node::Extension(_))));
275        assert_eq!(
276            node.unwrap().get(trie.db.as_ref(), path).unwrap(),
277            Some(value)
278        );
279    }
280
281    #[test]
282    fn insert_extension_branch_value_other() {
283        let trie = Trie::new_temp();
284        let mut node = pmt_node! { @(trie)
285            leaf { vec![1, 2, 3, 4, 16] => vec![0x12, 0x34, 0x56, 0x78] }
286        };
287
288        let path = Nibbles::from_bytes(&[0x12]);
289        let value = vec![0x17];
290
291        let node = node.insert(path.clone(), value.clone().into()).unwrap();
292
293        assert!(matches!(node, Some(Node::Extension(_))));
294        assert_eq!(
295            node.unwrap().get(trie.db.as_ref(), path).unwrap(),
296            Some(value)
297        );
298    }
299
300    // An insertion that returns branch [value=(x)] -> leaf (y) is not possible because of the path
301    // restrictions: nibbles come in pairs. If the first nibble is different, the node will be a
302    // branch but it cannot have a value. If the second nibble is different, then it'll be an
303    // extension followed by a branch with value and a child.
304
305    // Because of that, the two tests that would check those cases are neither necessary nor
306    // possible.
307
308    #[test]
309    fn remove_self() {
310        let mut node = LeafNode::new(
311            Nibbles::from_bytes(&[0x12, 0x34]),
312            vec![0x12, 0x34, 0x56, 0x78],
313        );
314        let (node, value) = node.remove(Nibbles::from_bytes(&[0x12, 0x34])).unwrap();
315
316        assert!(node.is_none());
317        assert_eq!(value, Some(vec![0x12, 0x34, 0x56, 0x78]));
318    }
319
320    #[test]
321    fn remove_none() {
322        let mut node = LeafNode::new(
323            Nibbles::from_bytes(&[0x12, 0x34]),
324            vec![0x12, 0x34, 0x56, 0x78],
325        );
326
327        let (node, value) = node.remove(Nibbles::from_bytes(&[0x12])).unwrap();
328
329        assert!(node.is_some());
330        assert_eq!(value, None);
331    }
332
333    #[test]
334    fn compute_hash_x() {
335        let node = LeafNode::new(Nibbles::from_bytes(b"key".as_ref()), b"value".to_vec());
336        let node_hash_ref = node.compute_hash(&NativeCrypto);
337        assert_eq!(
338            node_hash_ref.as_ref(),
339            &[
340                0xCB, 0x84, 0x20, 0x6B, 0x65, 0x79, 0x85, 0x76, 0x61, 0x6C, 0x75, 0x65
341            ],
342        );
343    }
344
345    #[test]
346    fn compute_hash_long() {
347        let node = LeafNode::new(
348            Nibbles::from_bytes(b"key".as_ref()),
349            b"a comparatively long value".to_vec(),
350        );
351
352        let node_hash_ref = node.compute_hash(&NativeCrypto);
353        assert_eq!(
354            node_hash_ref.as_ref(),
355            &[
356                0xEB, 0x92, 0x75, 0xB3, 0xAE, 0x09, 0x3A, 0x17, 0x75, 0x7C, 0xFB, 0x42, 0xF7, 0xD5,
357                0x57, 0xF9, 0xE5, 0x77, 0xBD, 0x5B, 0xEB, 0x86, 0xA8, 0x68, 0x49, 0x91, 0xA6, 0x5B,
358                0x87, 0x5F, 0x80, 0x7A,
359            ],
360        );
361    }
362
363    #[test]
364    fn symmetric_encoding_a() {
365        let node: Node = LeafNode::new(
366            Nibbles::from_bytes(b"key".as_ref()),
367            b"a comparatively long value".to_vec(),
368        )
369        .into();
370        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
371    }
372
373    #[test]
374    fn symmetric_encoding_b() {
375        let node: Node = LeafNode::new(
376            Nibbles::from_bytes(&[0x12, 0x34]),
377            vec![0x12, 0x34, 0x56, 0x78],
378        )
379        .into();
380        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
381    }
382
383    #[test]
384    fn symmetric_encoding_c() {
385        let node: Node = LeafNode::new(
386            Nibbles::from_bytes(&[]),
387            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 20],
388        )
389        .into();
390        assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
391    }
392}
ethrex_trie/node/leaf.rs

ethrex_trie/node/
leaf.rs
