1#[cfg(not(feature = "std"))]
2use alloc::{boxed::Box, string::ToString, vec::Vec};
3use core::mem;
4
5use ethrex_crypto::{Crypto, NativeCrypto};
6use ethrex_rlp::encode::RLPEncode;
7
8use crate::{
9 InconsistentTreeError, TrieDB, ValueRLP, error::TrieError, nibbles::Nibbles,
10 node::NodeRemoveResult, node_hash::NodeHash,
11};
12
13use super::{ExtensionNode, LeafNode, Node, NodeRef, ValueOrHash};
14
15#[derive(
18 Debug,
19 Clone,
20 PartialEq,
21 Default,
22 serde::Serialize,
23 serde::Deserialize,
24 rkyv::Serialize,
25 rkyv::Deserialize,
26 rkyv::Archive,
27)]
28pub struct BranchNode {
29 pub choices: [NodeRef; 16],
30 pub value: ValueRLP,
31}
32
33impl BranchNode {
34 const EMPTY_REF: NodeRef = NodeRef::Hash(NodeHash::Inline(([0; 31], 0)));
35
36 pub const EMPTY_CHOICES: [NodeRef; 16] = [Self::EMPTY_REF; 16];
38
39 pub fn new(choices: [NodeRef; 16]) -> Self {
41 Self {
42 choices,
43 value: Default::default(),
44 }
45 }
46
47 pub const fn new_with_value(choices: [NodeRef; 16], value: ValueRLP) -> Self {
49 Self { choices, value }
50 }
51
52 pub fn update(&mut self, new_value: ValueRLP) {
54 self.value = new_value;
55 }
56
57 pub fn get(&self, db: &dyn TrieDB, mut path: Nibbles) -> Result<Option<ValueRLP>, TrieError> {
59 if let Some(choice) = path.next_choice() {
62 let child_ref = &self.choices[choice];
64 if child_ref.is_valid() {
65 let child_node = child_ref.get_node(db, path.current())?.ok_or_else(|| {
66 TrieError::InconsistentTree(Box::new(
67 InconsistentTreeError::NodeNotFoundOnBranchNode(
68 child_ref
69 .compute_hash(&NativeCrypto)
70 .finalize(&NativeCrypto),
71 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
72 path.current(),
73 ),
74 ))
75 })?;
76 child_node.get(db, path)
77 } else {
78 Ok(None)
79 }
80 } else {
81 Ok((!self.value.is_empty()).then_some(self.value.clone()))
83 }
84 }
85
86 pub fn insert(
88 &mut self,
89 db: &dyn TrieDB,
90 mut path: Nibbles,
91 value: ValueOrHash,
92 ) -> Result<(), TrieError> {
93 if let Some(choice) = path.next_choice() {
96 match (&mut self.choices[choice], value) {
97 (choice_ref, ValueOrHash::Value(value)) if !choice_ref.is_valid() => {
99 let new_leaf = LeafNode::new(path, value);
100 *choice_ref = Node::from(new_leaf).into()
101 }
102 (choice_ref, ValueOrHash::Value(value)) => {
104 let Some(choice_node) = choice_ref.get_node_mut(db, path.current())? else {
105 return Err(TrieError::InconsistentTree(Box::new(
106 InconsistentTreeError::NodeNotFoundOnBranchNode(
107 choice_ref
108 .compute_hash(&NativeCrypto)
109 .finalize(&NativeCrypto),
110 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
111 path.current(),
112 ),
113 )));
114 };
115
116 choice_node.insert(db, path, value)?;
117 choice_ref.clear_hash();
118 }
119 (choice_ref, value @ ValueOrHash::Hash(hash)) => {
121 if !choice_ref.is_valid() {
122 *choice_ref = hash.into();
123 } else if path.is_empty() {
124 return Err(TrieError::Verify(
125 "attempt to override proof node with external hash".to_string(),
126 ));
127 } else {
128 let Some(choice_node) = choice_ref.get_node_mut(db, path.current())? else {
129 return Err(TrieError::InconsistentTree(Box::new(
130 InconsistentTreeError::NodeNotFoundOnBranchNode(
131 choice_ref
132 .compute_hash(&NativeCrypto)
133 .finalize(&NativeCrypto),
134 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
135 path.current(),
136 ),
137 )));
138 };
139 choice_node.insert(db, path, value)?;
140 choice_ref.clear_hash();
141 }
142 }
143 }
144 } else if let ValueOrHash::Value(value) = value {
145 self.update(value);
147 } else {
148 todo!("handle override case (error?)")
150 }
151
152 Ok(())
153 }
154
155 pub fn remove(
159 &mut self,
160 db: &dyn TrieDB,
161 mut path: Nibbles,
162 ) -> Result<(Option<NodeRemoveResult>, Option<ValueRLP>), TrieError> {
163 let base_path = path.clone();
181
182 let value = if let Some(choice_index) = path.next_choice() {
185 if self.choices[choice_index].is_valid() {
186 let Some(child_node) =
187 self.choices[choice_index].get_node_mut(db, path.current())?
188 else {
189 return Err(TrieError::InconsistentTree(Box::new(
190 InconsistentTreeError::NodeNotFoundOnBranchNode(
191 self.choices[choice_index]
192 .compute_hash(&NativeCrypto)
193 .finalize(&NativeCrypto),
194 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
195 path.current(),
196 ),
197 )));
198 };
199
200 let (empty_trie, old_value) = child_node.remove(db, path.clone())?;
202 if empty_trie {
203 self.choices[choice_index] = NodeHash::default().into();
205 }
206 self.choices[choice_index].clear_hash();
207 old_value
208 } else {
209 None
210 }
211 } else {
212 if !self.value.is_empty() {
214 Some(mem::take(&mut self.value))
215 } else {
216 None
217 }
218 };
219
220 let mut children = self
222 .choices
223 .iter_mut()
224 .enumerate()
225 .filter(|(_, child)| child.is_valid())
226 .collect::<Vec<_>>();
227 let new_node = match (children.len(), !self.value.is_empty()) {
228 (0, true) => NodeRemoveResult::New(
230 LeafNode::new(Nibbles::from_hex(vec![16]), mem::take(&mut self.value)).into(),
231 ),
232 (1, false) => {
234 let (choice_index, child_ref) = children.get_mut(0).unwrap();
235 let Some(child) = child_ref
236 .get_node_mut(db, base_path.current().append_new(*choice_index as u8))?
237 else {
238 return Err(TrieError::InconsistentTree(Box::new(
239 InconsistentTreeError::NodeNotFoundOnBranchNode(
240 child_ref
241 .compute_hash(&NativeCrypto)
242 .finalize(&NativeCrypto),
243 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
244 base_path.current(),
245 ),
246 )));
247 };
248
249 let node = match child {
250 Node::Branch(_) => ExtensionNode::new(
252 Nibbles::from_hex(vec![*choice_index as u8]),
253 child_ref.clone(),
254 )
255 .into(),
256 Node::Extension(extension_node) => {
258 let mut extension_node = extension_node.take();
259 extension_node.prefix.prepend(*choice_index as u8);
260 extension_node.into()
261 }
262 Node::Leaf(leaf) => {
263 let mut leaf = leaf.take();
264 leaf.partial.prepend(*choice_index as u8);
265 leaf.into()
266 }
267 };
268 NodeRemoveResult::New(node)
269 }
270 _ => NodeRemoveResult::Mutated,
272 };
273 Ok((Some(new_node), value))
274 }
275
276 pub fn compute_hash(&self, crypto: &dyn Crypto) -> NodeHash {
278 self.compute_hash_no_alloc(&mut vec![], crypto)
279 }
280
281 pub fn compute_hash_no_alloc(&self, buf: &mut Vec<u8>, crypto: &dyn Crypto) -> NodeHash {
283 buf.clear();
284 self.encode_into_vec(buf);
285 let hash = NodeHash::from_encoded(buf, crypto);
286 buf.clear();
287 hash
288 }
289
290 pub fn get_path(
294 &self,
295 db: &dyn TrieDB,
296 mut path: Nibbles,
297 node_path: &mut Vec<Vec<u8>>,
298 ) -> Result<(), TrieError> {
299 let encoded = self.encode_to_vec();
301 if encoded.len() >= 32 {
302 node_path.push(encoded);
303 };
304 if let Some(choice) = path.next_choice() {
306 let child_ref = &self.choices[choice];
308 if child_ref.is_valid() {
309 let child_node = child_ref.get_node(db, path.current())?.ok_or_else(|| {
310 TrieError::InconsistentTree(Box::new(
311 InconsistentTreeError::NodeNotFoundOnBranchNode(
312 child_ref
313 .compute_hash(&NativeCrypto)
314 .finalize(&NativeCrypto),
315 self.compute_hash(&NativeCrypto).finalize(&NativeCrypto),
316 path.current(),
317 ),
318 ))
319 })?;
320 child_node.get_path(db, path, node_path)?;
321 }
322 }
323 Ok(())
324 }
325}
326
327#[cfg(test)]
328mod test {
329 use ethereum_types::H256;
330 use ethrex_crypto::NativeCrypto;
331 use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
332
333 use super::*;
334
335 use crate::{Trie, pmt_node};
336
337 #[test]
338 fn new() {
339 let node = BranchNode::new({
340 let mut choices = BranchNode::EMPTY_CHOICES;
341
342 choices[2] = NodeHash::Hashed(H256([2; 32])).into();
343 choices[5] = NodeHash::Hashed(H256([5; 32])).into();
344
345 choices
346 });
347
348 assert_eq!(
349 node.choices,
350 [
351 Default::default(),
352 Default::default(),
353 NodeHash::Hashed(H256([2; 32])).into(),
354 Default::default(),
355 Default::default(),
356 NodeHash::Hashed(H256([5; 32])).into(),
357 Default::default(),
358 Default::default(),
359 Default::default(),
360 Default::default(),
361 Default::default(),
362 Default::default(),
363 Default::default(),
364 Default::default(),
365 Default::default(),
366 Default::default(),
367 ],
368 );
369 }
370
371 #[test]
372 fn get_some() {
373 let trie = Trie::new_temp();
374 let node = pmt_node! { @(trie)
375 branch {
376 0 => leaf { vec![0,16] => vec![0x12, 0x34, 0x56, 0x78] },
377 1 => leaf { vec![0,16] => vec![0x34, 0x56, 0x78, 0x9A] },
378 }
379 };
380
381 assert_eq!(
382 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
383 .unwrap(),
384 Some(vec![0x12, 0x34, 0x56, 0x78]),
385 );
386 assert_eq!(
387 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x10]))
388 .unwrap(),
389 Some(vec![0x34, 0x56, 0x78, 0x9A]),
390 );
391 }
392
393 #[test]
394 fn get_none() {
395 let trie = Trie::new_temp();
396 let node = pmt_node! { @(trie)
397 branch {
398 0 => leaf { vec![0,16] => vec![0x12, 0x34, 0x56, 0x78] },
399 1 => leaf { vec![0,16] => vec![0x34, 0x56, 0x78, 0x9A] },
400 }
401 };
402
403 assert_eq!(
404 node.get(trie.db.as_ref(), Nibbles::from_bytes(&[0x20]))
405 .unwrap(),
406 None,
407 );
408 }
409
410 #[test]
411 fn insert_self() {
412 let trie = Trie::new_temp();
413 let mut node = pmt_node! { @(trie)
414 branch {
415 0 => leaf { vec![0, 16] => vec![0x12, 0x34, 0x56, 0x78] },
416 1 => leaf { vec![0, 16] => vec![0x34, 0x56, 0x78, 0x9A] },
417 }
418 };
419 let path = Nibbles::from_bytes(&[2]);
420 let value = vec![0x3];
421
422 node.insert(trie.db.as_ref(), path.clone(), value.clone().into())
423 .unwrap();
424
425 assert_eq!(node.get(trie.db.as_ref(), path).unwrap(), Some(value));
426 }
427
428 #[test]
429 fn insert_choice() {
430 let trie = Trie::new_temp();
431 let mut node = pmt_node! { @(trie)
432 branch {
433 0 => leaf { vec![0, 16] => vec![0x12, 0x34, 0x56, 0x78] },
434 1 => leaf { vec![0, 16] => vec![0x34, 0x56, 0x78, 0x9A] },
435 }
436 };
437
438 let path = Nibbles::from_bytes(&[0x20]);
439 let value = vec![0x21];
440
441 node.insert(trie.db.as_ref(), path.clone(), value.clone().into())
442 .unwrap();
443
444 assert_eq!(node.get(trie.db.as_ref(), path).unwrap(), Some(value));
445 }
446
447 #[test]
448 fn insert_passthrough() {
449 let trie = Trie::new_temp();
450 let node = pmt_node! { @(trie)
451 branch {
452 0 => leaf { vec![0, 16] => vec![0x12, 0x34, 0x56, 0x78] },
453 1 => leaf { vec![0, 16] => vec![0x34, 0x56, 0x78, 0x9A] },
454 }
455 };
456
457 let path = Nibbles::from_bytes(&[0x00]).offset(2);
459 let value = vec![0x1];
460
461 let mut new_node = node.clone();
462 new_node
463 .insert(trie.db.as_ref(), path, value.clone().into())
464 .unwrap();
465
466 assert_eq!(new_node.choices, node.choices);
467 assert_eq!(new_node.value, value);
468 }
469
470 #[test]
471 fn remove_choice_into_inner() {
472 let trie = Trie::new_temp();
473 let mut node = pmt_node! { @(trie)
474 branch {
475 0 => leaf { vec![0, 16] => vec![0x00] },
476 1 => leaf { vec![0, 16] => vec![0x10] },
477 }
478 };
479
480 let (node, value) = node
481 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
482 .unwrap();
483
484 assert!(matches!(node, Some(NodeRemoveResult::New(Node::Leaf(_)))));
485 assert_eq!(value, Some(vec![0x00]));
486 }
487
488 #[test]
489 fn remove_choice() {
490 let trie = Trie::new_temp();
491 let mut node = pmt_node! { @(trie)
492 branch {
493 0 => leaf { vec![0, 16] => vec![0x00] },
494 1 => leaf { vec![0, 16] => vec![0x10] },
495 2 => leaf { vec![0, 16] => vec![0x10] },
496 }
497 };
498
499 let (node, value) = node
500 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
501 .unwrap();
502
503 assert!(matches!(node, Some(NodeRemoveResult::Mutated)));
504 assert_eq!(value, Some(vec![0x00]));
505 }
506
507 #[test]
508 fn remove_choice_into_value() {
509 let trie = Trie::new_temp();
510 let mut node = pmt_node! { @(trie)
511 branch {
512 0 => leaf { vec![0, 16] => vec![0x00] },
513 } with_leaf { &[0x01] => vec![0xFF] }
514 };
515
516 let (node, value) = node
517 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[0x00]))
518 .unwrap();
519
520 assert!(matches!(node, Some(NodeRemoveResult::New(Node::Leaf(_)))));
521 assert_eq!(value, Some(vec![0x00]));
522 }
523
524 #[test]
525 fn remove_value_into_inner() {
526 let trie = Trie::new_temp();
527 let mut node = pmt_node! { @(trie)
528 branch {
529 0 => leaf { vec![0, 16] => vec![0x00] },
530 } with_leaf { &[0x1] => vec![0xFF] }
531 };
532
533 let (node, value) = node
534 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[]))
535 .unwrap();
536
537 assert!(matches!(node, Some(NodeRemoveResult::New(Node::Leaf(_)))));
538 assert_eq!(value, Some(vec![0xFF]));
539 }
540
541 #[test]
542 fn remove_value() {
543 let trie = Trie::new_temp();
544 let mut node = pmt_node! { @(trie)
545 branch {
546 0 => leaf { vec![0, 16] => vec![0x00] },
547 1 => leaf { vec![0, 16] => vec![0x10] },
548 } with_leaf { &[0x1] => vec![0xFF] }
549 };
550
551 let (node, value) = node
552 .remove(trie.db.as_ref(), Nibbles::from_bytes(&[]))
553 .unwrap();
554
555 assert!(matches!(node, Some(NodeRemoveResult::Mutated)));
556 assert_eq!(value, Some(vec![0xFF]));
557 }
558
559 #[test]
560 fn compute_hash_two_choices() {
561 let node = pmt_node! { @(trie)
562 branch {
563 2 => leaf { vec![0, 16] => vec![0x20] },
564 4 => leaf { vec![0, 16] => vec![0x40] },
565 }
566 };
567
568 assert_eq!(
569 node.compute_hash(&NativeCrypto).as_ref(),
570 &[
571 0xD5, 0x80, 0x80, 0xC2, 0x30, 0x20, 0x80, 0xC2, 0x30, 0x40, 0x80, 0x80, 0x80, 0x80,
572 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
573 ],
574 );
575 }
576
577 #[test]
578 fn compute_hash_all_choices() {
579 let node = pmt_node! { @(trie)
580 branch {
581 0x0 => leaf { vec![0, 16] => vec![0x00] },
582 0x1 => leaf { vec![0, 16] => vec![0x10] },
583 0x2 => leaf { vec![0, 16] => vec![0x20] },
584 0x3 => leaf { vec![0, 16] => vec![0x30] },
585 0x4 => leaf { vec![0, 16] => vec![0x40] },
586 0x5 => leaf { vec![0, 16] => vec![0x50] },
587 0x6 => leaf { vec![0, 16] => vec![0x60] },
588 0x7 => leaf { vec![0, 16] => vec![0x70] },
589 0x8 => leaf { vec![0, 16] => vec![0x80] },
590 0x9 => leaf { vec![0, 16] => vec![0x90] },
591 0xA => leaf { vec![0, 16] => vec![0xA0] },
592 0xB => leaf { vec![0, 16] => vec![0xB0] },
593 0xC => leaf { vec![0, 16] => vec![0xC0] },
594 0xD => leaf { vec![0, 16] => vec![0xD0] },
595 0xE => leaf { vec![0, 16] => vec![0xE0] },
596 0xF => leaf { vec![0, 16] => vec![0xF0] },
597 }
598 };
599
600 assert_eq!(
601 node.compute_hash(&NativeCrypto).as_ref(),
602 &[
603 0x0A, 0x3C, 0x06, 0x2D, 0x4A, 0xE3, 0x61, 0xEC, 0xC4, 0x82, 0x07, 0xB3, 0x2A, 0xDB,
604 0x6A, 0x3A, 0x3F, 0x3E, 0x98, 0x33, 0xC8, 0x9C, 0x9A, 0x71, 0x66, 0x3F, 0x4E, 0xB5,
605 0x61, 0x72, 0xD4, 0x9D,
606 ],
607 );
608 }
609
610 #[test]
611 fn compute_hash_one_choice_with_value() {
612 let node = pmt_node! { @(trie)
613 branch {
614 2 => leaf { vec![0, 16] => vec![0x20] },
615 4 => leaf { vec![0, 16] => vec![0x40] },
616 } with_leaf { &[0x1] => vec![0x1] }
617 };
618
619 assert_eq!(
620 node.compute_hash(&NativeCrypto).as_ref(),
621 &[
622 0xD5, 0x80, 0x80, 0xC2, 0x30, 0x20, 0x80, 0xC2, 0x30, 0x40, 0x80, 0x80, 0x80, 0x80,
623 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01,
624 ],
625 );
626 }
627
628 #[test]
629 fn compute_hash_all_choices_with_value() {
630 let node = pmt_node! { @(trie)
631 branch {
632 0x0 => leaf { vec![0, 16] => vec![0x00] },
633 0x1 => leaf { vec![0, 16] => vec![0x10] },
634 0x2 => leaf { vec![0, 16] => vec![0x20] },
635 0x3 => leaf { vec![0, 16] => vec![0x30] },
636 0x4 => leaf { vec![0, 16] => vec![0x40] },
637 0x5 => leaf { vec![0, 16] => vec![0x50] },
638 0x6 => leaf { vec![0, 16] => vec![0x60] },
639 0x7 => leaf { vec![0, 16] => vec![0x70] },
640 0x8 => leaf { vec![0, 16] => vec![0x80] },
641 0x9 => leaf { vec![0, 16] => vec![0x90] },
642 0xA => leaf { vec![0, 16] => vec![0xA0] },
643 0xB => leaf { vec![0, 16] => vec![0xB0] },
644 0xC => leaf { vec![0, 16] => vec![0xC0] },
645 0xD => leaf { vec![0, 16] => vec![0xD0] },
646 0xE => leaf { vec![0, 16] => vec![0xE0] },
647 0xF => leaf { vec![0, 16] => vec![0xF0] },
648 } with_leaf { &[0x1] => vec![0x1] }
649 };
650
651 assert_eq!(
652 node.compute_hash(&NativeCrypto).as_ref(),
653 &[
654 0x2A, 0x85, 0x67, 0xC5, 0x63, 0x4A, 0x87, 0xBA, 0x19, 0x6F, 0x2C, 0x65, 0x15, 0x16,
655 0x66, 0x37, 0xE0, 0x9A, 0x34, 0xE6, 0xC9, 0xB0, 0x4D, 0xA5, 0x6F, 0xC4, 0x70, 0x4E,
656 0x38, 0x61, 0x7D, 0x8E
657 ],
658 );
659 }
660
661 #[test]
662 fn symmetric_encoding_a() {
663 let node: Node = pmt_node! { @(trie)
664 branch {
665 0 => leaf { vec![0,16] => vec![0x12, 0x34, 0x56, 0x78] },
666 1 => leaf { vec![0,16] => vec![0x34, 0x56, 0x78, 0x9A] },
667 }
668 }
669 .into();
670 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
671 }
672
673 #[test]
674 fn symmetric_encoding_b() {
675 let node: Node = pmt_node! { @(trie)
676 branch {
677 0 => leaf { vec![0, 16] => vec![0x00] },
678 1 => leaf { vec![0, 16] => vec![0x10] },
679 3 => extension { [0], branch {
680 0 => leaf { vec![16] => vec![0x01, 0x00] },
681 1 => leaf { vec![16] => vec![0x01, 0x01] },
682 } },
683 }
684 }
685 .into();
686 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
687 }
688
689 #[test]
690 fn symmetric_encoding_c() {
691 let node: Node = pmt_node! { @(trie)
692 branch {
693 0x0 => leaf { vec![0, 16] => vec![0x00] },
694 0x1 => leaf { vec![0, 16] => vec![0x10] },
695 0x2 => leaf { vec![0, 16] => vec![0x20] },
696 0x3 => leaf { vec![0, 16] => vec![0x30] },
697 0x4 => leaf { vec![0, 16] => vec![0x40] },
698 0x5 => leaf { vec![0, 16] => vec![0x50] },
699 0x6 => leaf { vec![0, 16] => vec![0x60] },
700 0x7 => leaf { vec![0, 16] => vec![0x70] },
701 0x8 => leaf { vec![0, 16] => vec![0x80] },
702 0x9 => leaf { vec![0, 16] => vec![0x90] },
703 0xA => leaf { vec![0, 16] => vec![0xA0] },
704 0xB => leaf { vec![0, 16] => vec![0xB0] },
705 0xC => leaf { vec![0, 16] => vec![0xC0] },
706 0xD => leaf { vec![0, 16] => vec![0xD0] },
707 0xE => leaf { vec![0, 16] => vec![0xE0] },
708 0xF => leaf { vec![0, 16] => vec![0xF0] },
709 } with_leaf { &[0x1] => vec![0x1] }
710 }
711 .into();
712 assert_eq!(Node::decode(&node.encode_to_vec()).unwrap(), node)
713 }
714}