ethrex_trie/
logger.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex},
4};
5
6use ethrex_crypto::NativeCrypto;
7use ethrex_rlp::decode::RLPDecode;
8
9use crate::{Nibbles, Node, NodeHash, Trie, TrieDB, TrieError};
10
11pub type TrieWitness = Arc<Mutex<HashMap<NodeHash, Node>>>;
12
13pub struct TrieLogger {
14    inner_db: Box<dyn TrieDB>,
15    witness: TrieWitness,
16}
17
18impl TrieLogger {
19    pub fn get_witness(&self) -> Result<HashMap<NodeHash, Node>, TrieError> {
20        let lock = self.witness.lock().map_err(|_| TrieError::LockError)?;
21        Ok(lock.clone())
22    }
23
24    pub fn open_trie(trie: Trie) -> (TrieWitness, Trie) {
25        let root = trie.hash_no_commit(&NativeCrypto);
26        let db = trie.db;
27        let witness = Arc::new(Mutex::new(HashMap::new()));
28        let logger = TrieLogger {
29            inner_db: db,
30            witness: witness.clone(),
31        };
32        (witness, Trie::open(Box::new(logger), root))
33    }
34}
35
36impl TrieDB for TrieLogger {
37    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
38        let result = self.inner_db.get(key)?;
39        if let Some(result) = result.as_ref()
40            && let Ok(decoded) = Node::decode(result)
41        {
42            let mut lock = self.witness.lock().map_err(|_| TrieError::LockError)?;
43            lock.insert(decoded.compute_hash(&NativeCrypto), decoded);
44        }
45        Ok(result)
46    }
47
48    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
49        self.inner_db.put(key, value)
50    }
51
52    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
53        self.inner_db.put_batch(key_values)
54    }
55}
ethrex_trie/logger.rs

ethrex_trie/
logger.rs
