1#[cfg(not(feature = "std"))]
2use alloc::{boxed::Box, vec::Vec};
3use ethrex_crypto::{Crypto, NativeCrypto};
4use ethrex_rlp::encode::RLPEncode;
5
6use crate::ValueRLP;
7use crate::nibbles::Nibbles;
8use crate::node::NodeRemoveResult;
9use crate::node_hash::NodeHash;
10use crate::{
11 TrieDB,
12 error::{ExtensionNodeErrorData, InconsistentTreeError, TrieError},
13};
14
15use super::{BranchNode, Node, NodeRef, ValueOrHash};
16
17#[derive(
20 Debug,
21 Clone,
22 PartialEq,
23 serde::Serialize,
24 serde::Deserialize,
25 rkyv::Serialize,
26 rkyv::Deserialize,
27 rkyv::Archive,
28)]
29pub struct ExtensionNode {
30 pub prefix: Nibbles,
31 pub child: NodeRef,
32}
33
34impl ExtensionNode {
35 pub const fn new(prefix: Nibbles, child: NodeRef) -> Self {
37 Self { prefix, child }
38 }
39
40 pub fn get(&self, db: &dyn TrieDB, mut path: Nibbles) -> Result<Option<ValueRLP>, TrieError> {
42 if path.skip_prefix(&self.prefix) {
45 let child_node = self.child.get_node(db, path.current())?.ok_or_else(|| {
46 TrieError::InconsistentTree(Box::new(
47 InconsistentTreeError::ExtensionNodeChildNotFound(ExtensionNodeErrorData {
48 node_hash: self
49 .child
50 .compute_hash(&NativeCrypto)
51 .finalize(&NativeCrypto),
52 extension_node_hash: self
53 .compute_hash(&NativeCrypto)
54 .finalize(&NativeCrypto),
55 extension_node_prefix: self.prefix.clone(),
56 node_path: path.current(),
57 }),
58 ))
59 })?;
60
61 child_node.get(db, path)
62 } else {
63 Ok(None)
64 }
65 }
66
67 pub fn insert(
70 &mut self,
71 db: &dyn TrieDB,
72 path: Nibbles,
73 value: ValueOrHash,
74 ) -> Result<Option<Node>, TrieError> {
75 let match_index = path.count_prefix(&self.prefix);
86 if match_index == self.prefix.len() {
87 let path = path.offset(match_index);
88 let Some(child_node) = self.child.get_node_mut(db, path.current())? else {
90 return Err(TrieError::InconsistentTree(Box::new(
91 InconsistentTreeError::ExtensionNodeChildNotFound(ExtensionNodeErrorData {
92 node_hash: self
93 .child
94 .compute_hash(&NativeCrypto)
95 .finalize(&NativeCrypto),
96 extension_node_hash: self
97 .compute_hash(&NativeCrypto)
98 .finalize(&NativeCrypto),
99 extension_node_prefix: self.prefix.clone(),
100 node_path: path.current(),
101 }),
102 )));
103 };
104 child_node.insert(db, path, value)?;
105 self.child.clear_hash();
106 Ok(None)
107 } else if match_index == 0 {
108 let mut new_node = if self.prefix.len() == 1 {
109 self.child.clone()
110 } else {
111 Node::from(ExtensionNode::new(
112 self.prefix.offset(1),
113 self.child.clone(),
114 ))
115 .into()
116 };
117 let mut choices = BranchNode::EMPTY_CHOICES;
118 let mut branch_node = if self.prefix.at(0) == 16 {
119 match new_node.get_node_mut(db, path.current())? {
120 Some(Node::Leaf(leaf)) => {
121 BranchNode::new_with_value(choices, leaf.value.clone())
122 }
123 Some(_) => {
124 return Err(TrieError::InconsistentTree(Box::new(
125 InconsistentTreeError::ExtensionNodeChildDiffers(
126 ExtensionNodeErrorData {
127 node_hash: new_node
128 .compute_hash(&NativeCrypto)
129 .finalize(&NativeCrypto),
130 extension_node_hash: self
131 .compute_hash(&NativeCrypto)
132 .finalize(&NativeCrypto),
133 extension_node_prefix: self.prefix.clone(),
134 node_path: path.current(),
135 },
136 ),
137 )));
138 }
139 None => {
140 return Err(TrieError::InconsistentTree(Box::new(
141 InconsistentTreeError::ExtensionNodeChildNotFound(
142 ExtensionNodeErrorData {
143 node_hash: new_node
144 .compute_hash(&NativeCrypto)
145 .finalize(&NativeCrypto),
146 extension_node_hash: self
147 .compute_hash(&NativeCrypto)
148 .finalize(&NativeCrypto),
149 extension_node_prefix: self.prefix.clone(),
150 node_path: path.current(),
151 },
152 ),
153 )));
154 }
155 }
156 } else {
157 choices[self.prefix.at(0)] = new_node;
158 BranchNode::new(choices)
159 };
160 branch_node.insert(db, path, value)?;
161 Ok(Some(branch_node.into()))
162 } else {
163 let mut new_extension =
164 ExtensionNode::new(self.prefix.offset(match_index), self.child.clone());
165 let new_node = new_extension
166 .insert(db, path.offset(match_index), value)?
167 .unwrap_or(new_extension.into());
168 self.prefix = self.prefix.slice(0, match_index);
169 self.child = new_node.into();
170 Ok(None)
171 }
172 }
173
174 pub fn remove(
175 &mut self,
176 db: &dyn TrieDB,
177 mut path: Nibbles,
178 ) -> Result<(Option<NodeRemoveResult>, Option<ValueRLP>), TrieError> {
179 if path.skip_prefix(&self.prefix) {
189 let Some(child_node) = self.child.get_node_mut(db, path.current())? else {
190 return Err(TrieError::InconsistentTree(Box::new(
191 InconsistentTreeError::ExtensionNodeChildNotFound(ExtensionNodeErrorData {
192 node_hash: self
193 .child
194 .compute_hash(&NativeCrypto)
195 .finalize(&NativeCrypto),
196 extension_node_hash: self
197 .compute_hash(&NativeCrypto)
198 .finalize(&NativeCrypto),
199 extension_node_prefix: self.prefix.clone(),
200 node_path: path.current(),
201 }),
202 )));
203 };
204 let (empty_trie, old_value) = child_node.remove(db, path)?;
206 let result = if empty_trie {
208 Ok((None, old_value))
209 } else {
210 let node = match child_node {
211 branch_node @ Node::Branch(_) => {
213 self.child = (*branch_node).clone().into();
214 NodeRemoveResult::Mutated
215 }
216 Node::Extension(extension_node) => {
218 let mut extension_node = extension_node.take();
219 let mut self_node = self.take();
220 self_node.prefix.extend(&extension_node.prefix);
221 extension_node.prefix = self_node.prefix;
222 NodeRemoveResult::New(extension_node.into())
223 }
224 Node::Leaf(leaf_node) => {
226 let mut leaf_node = leaf_node.take();
227 let mut self_node = self.take();
228 self_node.prefix.extend(&leaf_node.partial);
229 leaf_node.partial = self_node.prefix;
230 NodeRemoveResult::New(leaf_node.into())
231 }
232 };
233 Ok((Some(node), old_value))
234 };
235 self.child.clear_hash();
236 result
237 } else {
238 Ok((Some(NodeRemoveResult::Mutated), None))
239 }
240 }
241
242 pub fn compute_hash(&self, crypto: &dyn Crypto) -> NodeHash {
244 self.compute_hash_no_alloc(&mut vec![], crypto)
245 }
246
247 pub fn compute_hash_no_alloc(&self, buf: &mut Vec<u8>, crypto: &dyn Crypto) -> NodeHash {
249 buf.clear();
250 self.encode(buf);
251 let hash = NodeHash::from_encoded(buf, crypto);
252 buf.clear();
253 hash
254 }
255
256 pub fn get_path(
260 &self,
261 db: &dyn TrieDB,
262 mut path: Nibbles,
263 node_path: &mut Vec<Vec<u8>>,
264 ) -> Result<(), TrieError> {
265 let encoded = self.encode_to_vec();
267 if encoded.len() >= 32 {
268 node_path.push(encoded);
269 };
270 if path.skip_prefix(&self.prefix) {
272 let child_node = self.child.get_node(db, path.current())?.ok_or_else(|| {
273 TrieError::InconsistentTree(Box::new(
274 InconsistentTreeError::ExtensionNodeChildNotFound(ExtensionNodeErrorData {
275 node_hash: self
276 .child
277 .compute_hash(&NativeCrypto)
278 .finalize(&NativeCrypto),
279 extension_node_hash: self
280 .compute_hash(&NativeCrypto)
281 .finalize(&NativeCrypto),
282 extension_node_prefix: self.prefix.clone(),
283 node_path: path.current(),
284 }),
285 ))
286 })?;
287 child_node.get_path(db, path, node_path)?;
288 }
289 Ok(())
290 }
291
292 pub fn take(&mut self) -> Self {
296 ExtensionNode {
297 prefix: self.prefix.take(),
298 child: self.child.clone(),
299 }
300 }
301}
302
303#[cfg(test)]
304mod test {
305 use ethrex_crypto::NativeCrypto;
306 use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
307
308 use super::*;
309 use crate::{Trie, node::LeafNode, pmt_node};
310
311 #[test]
312 fn new() {
313 let node = ExtensionNode::new(Nibbles::default(), Default::default());
314
315 assert_eq!(node.prefix.len(), 0);
316 assert_eq!(node.child, Default::default());
317 }
318
319 #[test]
320 fn get_some() {
321 let trie = Trie::new_temp();
322 let node = pmt_node! { @(trie)
323 extension { [0], branch {
324 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
325 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
326 } }
327 };
328
329 assert_eq!(
330 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
331 .unwrap(),
332 Some(vec![0x12, 0x34, 0x56, 0x78]),
333 );
334 assert_eq!(
335 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x01]))
336 .unwrap(),
337 Some(vec![0x34, 0x56, 0x78, 0x9A]),
338 );
339 }
340
341 #[test]
342 fn get_none() {
343 let trie = Trie::new_temp();
344 let node = pmt_node! { @(trie)
345 extension { [0], branch {
346 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
347 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
348 } }
349 };
350
351 assert_eq!(
352 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x02]))
353 .unwrap(),
354 None,
355 );
356 }
357
358 #[test]
359 fn insert_passthrough() {
360 let trie = Trie::new_temp();
361 let mut node = pmt_node! { @(trie)
362 extension { [0], branch {
363 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
364 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
365 } }
366 };
367
368 let none = node
369 .insert(
370 trie.db.as_ref(),
371 Nibbles::from_bytes(&[0x02]),
372 Vec::new().into(),
373 )
374 .unwrap();
375 assert!(none.is_none());
376
377 assert_eq!(node.prefix.as_ref(), &[0]);
378 }
379
380 #[test]
381 fn insert_branch() {
382 let trie = Trie::new_temp();
383 let mut node = pmt_node! { @(trie)
384 extension { [0], branch {
385 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
386 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
387 } }
388 };
389
390 let node = node
391 .insert(
392 trie.db.as_ref(),
393 Nibbles::from_bytes(&[0x10]),
394 vec![0x20].into(),
395 )
396 .unwrap();
397 let node = match node {
398 Some(Node::Branch(x)) => x,
399 _ => panic!("expected a branch node"),
400 };
401 assert_eq!(
402 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x10]))
403 .unwrap(),
404 Some(vec![0x20])
405 );
406 }
407
408 #[test]
409 fn insert_branch_extension() {
410 let trie = Trie::new_temp();
411 let mut node = pmt_node! { @(trie)
412 extension { [0, 0], branch {
413 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
414 1 => leaf { vec![16]=> vec![0x34, 0x56, 0x78, 0x9A] },
415 } }
416 };
417
418 let node = node
419 .insert(
420 trie.db.as_ref(),
421 Nibbles::from_bytes(&[0x10]),
422 vec![0x20].into(),
423 )
424 .unwrap();
425 let node = match node {
426 Some(Node::Branch(x)) => x,
427 _ => panic!("expected a branch node"),
428 };
429 assert_eq!(
430 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x10]))
431 .unwrap(),
432 Some(vec![0x20])
433 );
434 }
435
436 #[test]
437 fn insert_extension_branch() {
438 let trie = Trie::new_temp();
439 let mut node = pmt_node! { @(trie)
440 extension { [0, 0], branch {
441 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
442 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
443 } }
444 };
445
446 let path = Nibbles::from_bytes(&[0x01]);
447 let value = vec![0x02];
448
449 let none = node
450 .insert(trie.db.as_ref(), path.clone(), value.clone().into())
451 .unwrap();
452
453 assert!(none.is_none());
454 assert_eq!(node.get(trie.db.as_ref(), path).unwrap(), Some(value));
455 }
456
457 #[test]
458 fn insert_extension_branch_extension() {
459 let trie = Trie::new_temp();
460 let mut node = pmt_node! { @(trie)
461 extension { [0, 0], branch {
462 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
463 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
464 } }
465 };
466
467 let path = Nibbles::from_bytes(&[0x01]);
468 let value = vec![0x04];
469
470 let none = node
471 .insert(trie.db.as_ref(), path.clone(), value.clone().into())
472 .unwrap();
473
474 assert!(none.is_none());
475 assert_eq!(node.get(trie.db.as_ref(), path).unwrap(), Some(value));
476 }
477
478 #[test]
479 fn remove_none() {
480 let trie = Trie::new_temp();
481 let mut node = pmt_node! { @(trie)
482 extension { [0], branch {
483 0 => leaf { vec![16] => vec![0x00] },
484 1 => leaf { vec![16] => vec![0x01] },
485 } }
486 };
487
488 let (node, value) = node
489 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x02]))
490 .unwrap();
491
492 assert!(matches!(node, Some(NodeRemoveResult::Mutated)));
493 assert_eq!(value, None);
494 }
495
496 #[test]
497 fn remove_into_leaf() {
498 let trie = Trie::new_temp();
499 let mut node = pmt_node! { @(trie)
500 extension { [0], branch {
501 0 => leaf { vec![16] => vec![0x00] },
502 1 => leaf { vec![16] => vec![0x01] },
503 } }
504 };
505
506 let (node, value) = node
507 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x01]))
508 .unwrap();
509
510 assert!(matches!(node, Some(NodeRemoveResult::New(Node::Leaf(_)))));
511 assert_eq!(value, Some(vec![0x01]));
512 }
513
514 #[test]
515 fn remove_into_extension() {
516 let trie = Trie::new_temp();
517 let mut node = pmt_node! { @(trie)
518 extension { [0], branch {
519 0 => leaf { vec![16] => vec![0x00] },
520 1 => extension { [0], branch {
521 0 => leaf { vec![16] => vec![0x01, 0x00] },
522 1 => leaf { vec![16] => vec![0x01, 0x01] },
523 } },
524 } }
525 };
526
527 let (node, value) = node
528 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
529 .unwrap();
530
531 assert!(matches!(
532 node,
533 Some(NodeRemoveResult::New(Node::Extension(_)))
534 ));
535 assert_eq!(value, Some(vec![0x00]));
536 }
537
538 #[test]
539 fn compute_hash() {
540 let leaf_node_a = LeafNode::new(Nibbles::from_hex(vec![0, 16]), vec![0x12, 0x34]);
550 let leaf_node_b = LeafNode::new(Nibbles::from_hex(vec![0, 16]), vec![0x56, 0x78]);
551 let mut choices = BranchNode::EMPTY_CHOICES;
552 choices[0] = leaf_node_a.compute_hash(&NativeCrypto).into();
553 choices[1] = leaf_node_b.compute_hash(&NativeCrypto).into();
554 let branch_node = BranchNode::new(choices);
555 let node = ExtensionNode::new(
556 Nibbles::from_hex(vec![0, 0]),
557 branch_node.compute_hash(&NativeCrypto).into(),
558 );
559
560 assert_eq!(
561 node.compute_hash(&NativeCrypto).as_ref(),
562 &[
563 0xDD, 0x82, 0x00, 0x00, 0xD9, 0xC4, 0x30, 0x82, 0x12, 0x34, 0xC4, 0x30, 0x82, 0x56,
564 0x78, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
565 0x80, 0x80,
566 ],
567 );
568 }
569
570 #[test]
571 fn compute_hash_long() {
572 let leaf_node_a = LeafNode::new(
582 Nibbles::from_hex(vec![0, 16]),
583 vec![0x12, 0x34, 0x56, 0x78, 0x9A],
584 );
585 let leaf_node_b = LeafNode::new(
586 Nibbles::from_hex(vec![0, 16]),
587 vec![0x34, 0x56, 0x78, 0x9A, 0xBC],
588 );
589 let mut choices = BranchNode::EMPTY_CHOICES;
590 choices[0] = leaf_node_a.compute_hash(&NativeCrypto).into();
591 choices[1] = leaf_node_b.compute_hash(&NativeCrypto).into();
592 let branch_node = BranchNode::new(choices);
593 let node = ExtensionNode::new(
594 Nibbles::from_hex(vec![0, 0]),
595 branch_node.compute_hash(&NativeCrypto).into(),
596 );
597
598 assert_eq!(
599 node.compute_hash(&NativeCrypto).as_ref(),
600 &[
601 0xFA, 0xBA, 0x42, 0x79, 0xB3, 0x9B, 0xCD, 0xEB, 0x7C, 0x53, 0x0F, 0xD7, 0x6E, 0x5A,
602 0xA3, 0x48, 0xD3, 0x30, 0x76, 0x26, 0x14, 0x84, 0x55, 0xA0, 0xAE, 0xFE, 0x0F, 0x52,
603 0x89, 0x5F, 0x36, 0x06,
604 ],
605 );
606 }
607
608 #[test]
609 fn symmetric_encoding_a() {
610 let node: Node = pmt_node! { @(trie)
611 extension { [0], branch {
612 0 => leaf { vec![16] => vec![0x12, 0x34, 0x56, 0x78] },
613 1 => leaf { vec![16] => vec![0x34, 0x56, 0x78, 0x9A] },
614 } }
615 }
616 .into();
617 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
618 }
619
620 #[test]
621 fn symmetric_encoding_b() {
622 let node: Node = pmt_node! { @(trie)
623 extension { [0], branch {
624 0 => leaf { vec![16] => vec![0x00] },
625 1 => extension { [0], branch {
626 0 => leaf { vec![16] => vec![0x01, 0x00] },
627 1 => leaf { vec![16] => vec![0x01, 0x01] },
628 } },
629 } }
630 }
631 .into();
632
633 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
634 }
635
636 #[test]
637 fn symmetric_encoding_c() {
638 let node: Node = pmt_node! { @(trie)
639 extension { [0], branch {
640 0 => leaf { vec![16] => vec![0x00] },
641 1 => extension { [0], branch {
642 0 => leaf { vec![16] => vec![0x01, 0x00] },
643 1 => leaf { vec![16] => vec![0x01, 0x01] },
644 2 => leaf { vec![16] => vec![0x01, 0x00] },
645 3 => leaf { vec![16] => vec![0x03, 0x01] },
646 4 => leaf { vec![16] => vec![0x04, 0x00] },
647 5 => leaf { vec![16] => vec![0x05, 0x01] },
648 6 => leaf { vec![16] => vec![0x06, 0x00] },
649 7 => leaf { vec![16] => vec![0x07, 0x01] },
650 8 => leaf { vec![16] => vec![0x08, 0x00] },
651 9 => leaf { vec![16] => vec![0x09, 0x01] },
652 10 => leaf { vec![16] => vec![0x10, 0x00] },
653 11 => leaf { vec![16] => vec![0x11, 0x01] },
654 } },
655 } }
656 }
657 .into();
658 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
659 }
660}