ethrex_trie/node/
leaf.rs

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

ethrex_trie/node/
leaf.rs
