1pub mod db;
2pub mod error;
3pub mod logger;
4mod nibbles;
5pub mod node;
6mod node_hash;
7pub mod rkyv_utils;
8mod rlp;
9#[cfg(test)]
10mod test_utils;
11pub mod threadpool;
12mod trie_iter;
13pub mod trie_sorted;
14mod verify_range;
15use ethereum_types::H256;
16use ethrex_crypto::keccak::keccak_hash;
17use ethrex_crypto::{Crypto, NativeCrypto};
18use ethrex_rlp::constants::RLP_NULL;
19use ethrex_rlp::encode::RLPEncode;
20use rustc_hash::FxHashSet;
21use std::collections::BTreeMap;
22use std::sync::{Arc, Mutex};
23
24pub use self::db::{InMemoryTrieDB, TrieDB};
25pub use self::logger::{TrieLogger, TrieWitness};
26pub use self::nibbles::Nibbles;
27pub use self::threadpool::ThreadPool;
28pub use self::verify_range::verify_range;
29pub use self::{
30 node::{Node, NodeRef},
31 node_hash::NodeHash,
32};
33
34pub use self::error::{ExtensionNodeErrorData, InconsistentTreeError, TrieError};
35use self::{node::LeafNode, trie_iter::TrieIterator};
36
37use ethrex_rlp::decode::RLPDecode;
38use lazy_static::lazy_static;
39
40lazy_static! {
41 pub static ref EMPTY_TRIE_HASH: H256 = H256(
43 keccak_hash([RLP_NULL]),
44 );
45}
46
47pub type PathRLP = Vec<u8>;
49pub type ValueRLP = Vec<u8>;
51pub type NodeRLP = Vec<u8>;
53pub type TrieNode = (Nibbles, NodeRLP);
55
56pub struct Trie {
58 db: Box<dyn TrieDB>,
59 pub root: NodeRef,
60 pending_removal: FxHashSet<Nibbles>,
61 dirty: FxHashSet<Nibbles>,
62}
63
64impl Default for Trie {
65 fn default() -> Self {
66 Self::new_temp()
67 }
68}
69
70impl Trie {
71 pub fn new(db: Box<dyn TrieDB>) -> Self {
73 Self {
74 db,
75 root: NodeRef::default(),
76 pending_removal: Default::default(),
77 dirty: Default::default(),
78 }
79 }
80
81 pub fn open(db: Box<dyn TrieDB>, root: H256) -> Self {
83 Self {
84 db,
85 root: if root != *EMPTY_TRIE_HASH {
86 NodeHash::from(root).into()
87 } else {
88 Default::default()
89 },
90 pending_removal: Default::default(),
91 dirty: Default::default(),
92 }
93 }
94
95 pub fn db(&self) -> &dyn TrieDB {
100 self.db.as_ref()
101 }
102
103 pub fn get(&self, pathrlp: &[u8]) -> Result<Option<ValueRLP>, TrieError> {
105 let path = Nibbles::from_bytes(pathrlp);
106
107 if !self.dirty.contains(&path) && self.db().flatkeyvalue_computed(path.clone()) {
108 let Some(value_rlp) = self.db.get(path)? else {
109 return Ok(None);
110 };
111 if value_rlp.is_empty() {
112 return Ok(None);
113 }
114 return Ok(Some(value_rlp));
115 }
116
117 Ok(match self.root {
118 NodeRef::Node(ref node, _) => node.get(self.db.as_ref(), path)?,
119 NodeRef::Hash(hash) if hash.is_valid() => {
120 Node::decode(&self.db.get(Nibbles::default())?.ok_or_else(|| {
121 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(
122 hash.finalize(&NativeCrypto),
123 )))
124 })?)
125 .map_err(TrieError::RLPDecode)?
126 .get(self.db.as_ref(), path)?
127 }
128 _ => None,
129 })
130 }
131
132 pub fn insert(&mut self, path: PathRLP, value: ValueRLP) -> Result<(), TrieError> {
134 let path = Nibbles::from_bytes(&path);
135 self.pending_removal.remove(&path);
136 self.dirty.insert(path.clone());
137
138 if self.root.is_valid() {
139 self.root
141 .get_node_mut(self.db.as_ref(), Nibbles::default())?
142 .ok_or_else(|| {
143 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
144 })?
145 .insert(self.db.as_ref(), path, value)?
146 } else {
147 self.root = Node::from(LeafNode::new(path, value)).into()
149 };
150 self.root.clear_hash();
151
152 Ok(())
153 }
154
155 pub fn remove(&mut self, path: &[u8]) -> Result<Option<ValueRLP>, TrieError> {
158 self.dirty.insert(Nibbles::from_bytes(path));
159 if !self.root.is_valid() {
160 return Ok(None);
161 }
162 self.pending_removal.insert(Nibbles::from_bytes(path));
163
164 let (is_trie_empty, value) = self
166 .root
167 .get_node_mut(self.db.as_ref(), Nibbles::default())?
168 .ok_or_else(|| {
169 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
170 })?
171 .remove(self.db.as_ref(), Nibbles::from_bytes(path))?;
172 if is_trie_empty {
173 self.root = NodeRef::default();
174 } else {
175 self.root.clear_hash();
176 }
177
178 Ok(value)
179 }
180
181 pub fn hash(&mut self, crypto: &dyn Crypto) -> Result<H256, TrieError> {
185 self.commit(crypto)?;
186 Ok(self.hash_no_commit(crypto))
187 }
188
189 pub fn hash_no_commit(&self, crypto: &dyn Crypto) -> H256 {
192 if self.root.is_valid() {
193 let mut buf = Vec::with_capacity(512);
195 self.root
196 .compute_hash_no_alloc(&mut buf, crypto)
197 .finalize(crypto)
198 } else {
199 *EMPTY_TRIE_HASH
200 }
201 }
202
203 pub fn get_root_node(&self, path: Nibbles) -> Result<Arc<Node>, TrieError> {
204 self.root
205 .get_node_checked(self.db.as_ref(), path)?
206 .ok_or_else(|| {
207 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(
208 self.root
209 .compute_hash(&NativeCrypto)
210 .finalize(&NativeCrypto),
211 )))
212 })
213 }
214
215 pub fn collect_changes_since_last_hash(
221 &mut self,
222 crypto: &dyn Crypto,
223 ) -> (H256, Vec<TrieNode>) {
224 let updates = self.commit_without_storing(crypto);
225 let ret_hash = self.hash_no_commit(crypto);
226 (ret_hash, updates)
227 }
228
229 pub fn commit(&mut self, crypto: &dyn Crypto) -> Result<(), TrieError> {
234 let acc = self.commit_without_storing(crypto);
235 self.db.put_batch(acc)?;
236
237 self.db.commit()?;
239
240 Ok(())
241 }
242
243 pub fn commit_without_storing(&mut self, crypto: &dyn Crypto) -> Vec<TrieNode> {
246 let mut acc = Vec::new();
247 if self.root.is_valid() {
248 self.root.commit(Nibbles::default(), &mut acc, crypto);
249 }
250 if self.root.compute_hash(crypto) == NodeHash::Hashed(*EMPTY_TRIE_HASH) {
251 acc.push((Nibbles::default(), vec![RLP_NULL]))
252 }
253 acc.extend(self.pending_removal.drain().map(|nib| (nib, vec![])));
254
255 acc
256 }
257
258 pub fn get_proof(&self, path: &[u8]) -> Result<Vec<NodeRLP>, TrieError> {
266 if self.root.is_valid() {
267 let hash = self.root.compute_hash(&NativeCrypto);
268
269 let mut node_path = Vec::new();
270 if let NodeHash::Inline((data, len)) = hash {
271 node_path.push(data[..len as usize].to_vec());
272 }
273
274 let root = match self
275 .root
276 .get_node_checked(self.db.as_ref(), Nibbles::default())?
277 {
278 Some(x) => x,
279 None => return Ok(Vec::new()),
280 };
281 root.get_path(self.db.as_ref(), Nibbles::from_bytes(path), &mut node_path)?;
282
283 Ok(node_path)
284 } else {
285 Ok(Vec::new())
286 }
287 }
288
289 pub fn get_proofs(
293 &self,
294 paths: &[PathRLP],
295 ) -> Result<(Option<NodeRLP>, Vec<NodeRLP>), TrieError> {
296 if self.root.is_valid() {
297 let encoded_root = self.get_root_node(Nibbles::default())?.encode_to_vec();
298
299 let mut node_path: FxHashSet<_> = Default::default();
300 for path in paths {
301 let mut nodes = self.get_proof(path)?;
302 nodes.swap_remove(0);
303 node_path.extend(nodes);
304 }
305
306 Ok((Some(encoded_root), node_path.into_iter().collect()))
307 } else {
308 Ok((None, Vec::new()))
309 }
310 }
311
312 pub fn empty_in_memory() -> Self {
313 Self::new(Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(
314 BTreeMap::new(),
315 )))))
316 }
317
318 pub fn get_embedded_root(
320 all_nodes: &BTreeMap<H256, Node>,
321 root_hash: H256,
322 ) -> Result<NodeRef, TrieError> {
323 if root_hash == *EMPTY_TRIE_HASH {
325 return Ok(NodeRef::default());
326 }
327
328 let root_rlp = all_nodes.get(&root_hash).ok_or_else(|| {
329 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(root_hash)))
330 })?;
331
332 fn get_embedded_node(
333 all_nodes: &BTreeMap<H256, Node>,
334 cur_node: &Node,
335 ) -> Result<Node, TrieError> {
336 Ok(match cur_node.clone() {
337 Node::Branch(mut node) => {
338 for choice in &mut node.choices {
339 let NodeRef::Hash(hash) = *choice else {
340 continue;
341 };
342
343 if hash.is_valid() {
344 *choice = match all_nodes.get(&hash.finalize(&NativeCrypto)) {
345 Some(node) => get_embedded_node(all_nodes, node)?.into(),
346 None => hash.into(),
347 };
348 }
349 }
350
351 (*node).into()
352 }
353 Node::Extension(mut node) => {
354 let NodeRef::Hash(hash) = node.child else {
355 return Ok(node.into());
356 };
357
358 node.child = match all_nodes.get(&hash.finalize(&NativeCrypto)) {
359 Some(node) => get_embedded_node(all_nodes, node)?.into(),
360 None => hash.into(),
361 };
362
363 node.into()
364 }
365 Node::Leaf(node) => node.into(),
366 })
367 }
368
369 let root = get_embedded_node(all_nodes, root_rlp)?;
370 Ok(root.into())
371 }
372
373 pub fn from_nodes(
381 root_hash: H256,
382 state_nodes: &BTreeMap<H256, Node>,
383 ) -> Result<Self, TrieError> {
384 let mut trie = Trie::new(Box::new(InMemoryTrieDB::default()));
385 let root = Self::get_embedded_root(state_nodes, root_hash)?;
386 trie.root = root;
387
388 Ok(trie)
389 }
390
391 pub fn compute_hash_from_unsorted_iter(
393 iter: impl Iterator<Item = (PathRLP, ValueRLP)>,
394 crypto: &dyn Crypto,
395 ) -> H256 {
396 let mut trie = Trie::stateless();
397 for (path, value) in iter {
398 trie.insert(path, value).unwrap();
400 }
401
402 trie.hash_no_commit(crypto)
403 }
404
405 pub(crate) fn stateless() -> Trie {
408 struct NullTrieDB;
410
411 impl TrieDB for NullTrieDB {
412 fn get(&self, _key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
413 Ok(None)
414 }
415
416 fn put_batch(&self, _key_values: Vec<TrieNode>) -> Result<(), TrieError> {
417 Ok(())
418 }
419 }
420
421 Trie::new(Box::new(NullTrieDB))
422 }
423
424 pub fn get_node(&self, partial_path: &PathRLP) -> Result<Vec<u8>, TrieError> {
427 let partial_path = match partial_path.len() {
429 n if n < 32 => Nibbles::decode_compact(partial_path),
431 32 => Nibbles::from_bytes(partial_path),
433 _ => return Ok(vec![]),
435 };
436
437 fn get_node_inner(
438 db: &dyn TrieDB,
439 current_path: Nibbles,
440 node: &Node,
441 mut partial_path: Nibbles,
442 ) -> Result<Vec<u8>, TrieError> {
443 if partial_path.is_empty() {
445 return Ok(node.encode_to_vec());
446 }
447 match node {
448 Node::Branch(branch_node) => match partial_path.next_choice() {
449 Some(idx) => {
450 let child_ref = &branch_node.choices[idx];
451 if child_ref.is_valid() {
452 let child_path = current_path.append_new(idx as u8);
453 let child_node = child_ref
454 .get_node_checked(db, child_path.clone())?
455 .ok_or_else(|| {
456 TrieError::InconsistentTree(Box::new(
457 InconsistentTreeError::NodeNotFoundOnBranchNode(
458 child_ref
459 .compute_hash(&NativeCrypto)
460 .finalize(&NativeCrypto),
461 branch_node
462 .compute_hash(&NativeCrypto)
463 .finalize(&NativeCrypto),
464 child_path.clone(),
465 ),
466 ))
467 })?;
468 get_node_inner(db, child_path, &child_node, partial_path)
469 } else {
470 Ok(vec![])
471 }
472 }
473 _ => Ok(vec![]),
474 },
475 Node::Extension(extension_node) => {
476 if partial_path.skip_prefix(&extension_node.prefix)
477 && extension_node.child.is_valid()
478 {
479 let child_path = partial_path.concat(&extension_node.prefix);
480 let child_node = extension_node
481 .child
482 .get_node_checked(db, child_path.clone())?
483 .ok_or_else(|| {
484 TrieError::InconsistentTree(Box::new(
485 InconsistentTreeError::ExtensionNodeChildNotFound(
486 ExtensionNodeErrorData {
487 node_hash: extension_node
488 .child
489 .compute_hash(&NativeCrypto)
490 .finalize(&NativeCrypto),
491 extension_node_hash: extension_node
492 .compute_hash(&NativeCrypto)
493 .finalize(&NativeCrypto),
494 extension_node_prefix: extension_node.prefix.clone(),
495 node_path: child_path.clone(),
496 },
497 ),
498 ))
499 })?;
500 get_node_inner(db, child_path, &child_node, partial_path)
501 } else {
502 Ok(vec![])
503 }
504 }
505 Node::Leaf(_) => Ok(vec![]),
506 }
507 }
508
509 if self.root.is_valid() {
511 let root_node = self.get_root_node(Default::default())?;
512 get_node_inner(
513 self.db.as_ref(),
514 Default::default(),
515 &root_node,
516 partial_path,
517 )
518 } else {
519 Ok(Vec::new())
520 }
521 }
522
523 pub fn root_node(&self) -> Result<Option<Arc<Node>>, TrieError> {
524 if self.root.is_valid() {
525 self.root.get_node(self.db.as_ref(), Nibbles::default())
526 } else {
527 Ok(None)
528 }
529 }
530
531 pub fn new_temp() -> Self {
533 let db = InMemoryTrieDB::new(Default::default());
534 Trie::new(Box::new(db))
535 }
536
537 pub fn new_temp_with_root(root: NodeRef) -> Self {
541 let db = InMemoryTrieDB::new(Default::default());
542 let mut trie = Trie::new(Box::new(db));
543 trie.root = root;
544 trie
545 }
546
547 pub fn validate(self) -> Result<(), TrieError> {
554 let mut expected_count = if self.root.is_valid() { 1 } else { 0 };
555 for (_, node) in self.into_iter() {
556 expected_count -= 1;
557 match node {
558 Node::Branch(branch_node) => {
559 expected_count += branch_node
560 .choices
561 .iter()
562 .filter(|child| child.is_valid())
563 .count();
564 }
565 Node::Extension(_) => {
566 expected_count += 1;
567 }
568 Node::Leaf(_) => {}
569 }
570 }
571 if expected_count != 0 {
572 return Err(TrieError::Verify(format!(
573 "Node count mismatch, expected {expected_count} more"
574 )));
575 }
576 Ok(())
577 }
578
579 pub fn validate_parallel(self) -> Result<(), TrieError> {
582 use rayon::prelude::*;
583
584 if !self.root.is_valid() {
585 return Ok(());
586 }
587
588 let db = &*self.db;
589 let root_node = self
590 .root
591 .get_node_checked(db, Nibbles::default())?
592 .ok_or_else(|| TrieError::Verify("Root node not found".to_string()))?;
593
594 match &*root_node {
595 Node::Branch(branch_node) => {
596 let children: Vec<(Nibbles, NodeRef)> = branch_node
597 .choices
598 .iter()
599 .enumerate()
600 .filter(|(_, child)| child.is_valid())
601 .map(|(i, child)| {
602 let path = Nibbles::default().append_new(i as u8);
603 (path, child.clone())
604 })
605 .collect();
606
607 children.par_iter().try_for_each(|(start_path, start_ref)| {
608 validate_subtree(db, start_path.clone(), start_ref.clone())
609 })
610 }
611 _ => {
612 validate_subtree(db, Nibbles::default(), self.root.clone())
614 }
615 }
616 }
617}
618
619fn validate_subtree(
622 db: &dyn TrieDB,
623 start_path: Nibbles,
624 start_ref: NodeRef,
625) -> Result<(), TrieError> {
626 let mut expected_count: isize = 1;
627 let mut stack = vec![(start_path, start_ref)];
628
629 while let Some((path, node_ref)) = stack.pop() {
630 let node = node_ref
631 .get_node_checked(db, path.clone())?
632 .ok_or_else(|| TrieError::Verify(format!("Missing node at path {path:?}")))?;
633
634 expected_count -= 1;
635 match &*node {
636 Node::Branch(branch) => {
637 for (choice, child) in branch.choices.iter().enumerate().rev() {
638 if child.is_valid() {
639 expected_count += 1;
640 stack.push((path.append_new(choice as u8), child.clone()));
641 }
642 }
643 }
644 Node::Extension(ext) => {
645 expected_count += 1;
646 stack.push((path.concat(&ext.prefix), ext.child.clone()));
647 }
648 Node::Leaf(_) => {}
649 }
650 }
651
652 if expected_count != 0 {
653 return Err(TrieError::Verify(format!(
654 "Node count mismatch in subtree, expected {expected_count} more"
655 )));
656 }
657 Ok(())
658}
659
660impl IntoIterator for Trie {
661 type Item = (Nibbles, Node);
662
663 type IntoIter = TrieIterator;
664
665 fn into_iter(self) -> Self::IntoIter {
666 TrieIterator::new(self)
667 }
668}
669
670pub struct ProofTrie(Trie);
671
672impl ProofTrie {
673 pub fn insert(
674 &mut self,
675 partial_path: Nibbles,
676 external_ref: NodeHash,
677 ) -> Result<(), TrieError> {
678 if self.0.root.is_valid() {
679 self.0
681 .root
682 .get_node_mut(self.0.db.as_ref(), Nibbles::default())?
683 .ok_or_else(|| {
684 TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
685 })?
686 .insert(self.0.db.as_ref(), partial_path, external_ref)?;
687 self.0.root.clear_hash();
688 } else {
689 self.0.root = external_ref.into();
690 };
691
692 Ok(())
693 }
694
695 pub fn hash(&self, crypto: &dyn Crypto) -> H256 {
696 self.0.hash_no_commit(crypto)
697 }
698}
699
700impl From<Trie> for ProofTrie {
701 fn from(value: Trie) -> Self {
702 Self(value)
703 }
704}