ethrex_trie/
db.rs

1use ethereum_types::H256;
2use ethrex_rlp::encode::RLPEncode;
3
4use crate::{Nibbles, Node, Trie, error::TrieError};
5use std::{
6    collections::BTreeMap,
7    sync::{Arc, Mutex},
8};
9
10// Nibbles -> encoded node
11pub type NodeMap = Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>;
12
13pub trait TrieDB: Send + Sync {
14    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError>;
15    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError>;
16    // TODO: replace putbatch with this function.
17    fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
18        self.put_batch(
19            key_values
20                .iter()
21                .map(|node| (node.0.clone(), node.1.encode_to_vec()))
22                .collect(),
23        )
24    }
25    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
26        self.put_batch(vec![(key, value)])
27    }
28    /// Commits any pending changes to the underlying storage
29    /// For read-only or in-memory implementations, this is a no-op
30    fn commit(&self) -> Result<(), TrieError> {
31        Ok(())
32    }
33
34    fn flatkeyvalue_computed(&self, _key: Nibbles) -> bool {
35        false
36    }
37}
38
39// TODO: we should replace this with BackendTrieDB
40/// InMemory implementation for the TrieDB trait, with get and put operations.
41#[derive(Default)]
42pub struct InMemoryTrieDB {
43    inner: NodeMap,
44    prefix: Option<Nibbles>,
45}
46
47impl InMemoryTrieDB {
48    pub const fn new(map: NodeMap) -> Self {
49        Self {
50            inner: map,
51            prefix: None,
52        }
53    }
54
55    pub const fn new_with_prefix(map: NodeMap, prefix: Nibbles) -> Self {
56        Self {
57            inner: map,
58            prefix: Some(prefix),
59        }
60    }
61
62    pub fn new_empty() -> Self {
63        Self {
64            inner: Default::default(),
65            prefix: None,
66        }
67    }
68
69    // Do not remove or make private as we use this in ethrex-replay
70    pub fn from_nodes(
71        root_hash: H256,
72        state_nodes: &BTreeMap<H256, Node>,
73    ) -> Result<Self, TrieError> {
74        let mut embedded_root = Trie::get_embedded_root(state_nodes, root_hash)?;
75        let mut hashed_nodes = vec![];
76        embedded_root.commit(
77            Nibbles::default(),
78            &mut hashed_nodes,
79            &ethrex_crypto::NativeCrypto,
80        );
81
82        let hashed_nodes = hashed_nodes
83            .into_iter()
84            .map(|(k, v)| (k.into_vec(), v))
85            .collect();
86
87        let in_memory_trie = Arc::new(Mutex::new(hashed_nodes));
88        Ok(Self::new(in_memory_trie))
89    }
90
91    fn apply_prefix(&self, path: Nibbles) -> Nibbles {
92        match &self.prefix {
93            Some(prefix) => prefix.concat(&path),
94            None => path,
95        }
96    }
97
98    // Do not remove or make private as we use this in ethrex-replay
99    pub fn inner(&self) -> NodeMap {
100        Arc::clone(&self.inner)
101    }
102}
103
104impl TrieDB for InMemoryTrieDB {
105    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
106        Ok(self
107            .inner
108            .lock()
109            .map_err(|_| TrieError::LockError)?
110            .get(self.apply_prefix(key).as_ref())
111            .cloned())
112    }
113
114    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
115        let mut db = self.inner.lock().map_err(|_| TrieError::LockError)?;
116
117        for (key, value) in key_values {
118            let prefixed_key = self.apply_prefix(key);
119            db.insert(prefixed_key.into_vec(), value);
120        }
121
122        Ok(())
123    }
124}
ethrex_trie/db.rs

ethrex_trie/
db.rs
