ethrex_trie/
db.rs

1use ethrex_rlp::encode::RLPEncode;
2
3use crate::{Nibbles, Node, error::TrieError};
4use alloc::collections::BTreeMap;
5use alloc::sync::Arc;
6#[cfg(not(feature = "std"))]
7use alloc::vec::Vec;
8
9// `InMemoryTrieDB` is available in both builds. On `std` its map is guarded by
10// `std::sync::Mutex`; on `no_std` (the zkVM guest) by `spin::Mutex`, which never
11// contends because the guest is single-threaded. Only `from_nodes` stays host-only,
12// by design rather than necessity: the guest never builds a trie from a node dump.
13#[cfg(feature = "std")]
14use crate::Trie;
15#[cfg(feature = "std")]
16use ethereum_types::H256;
17#[cfg(not(feature = "std"))]
18use spin::Mutex;
19#[cfg(feature = "std")]
20use std::sync::Mutex;
21
22// Nibbles -> encoded node
23pub type NodeMap = Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>;
24
25// Guard returned by locking a `NodeMap`; the backend mutex differs per build.
26#[cfg(feature = "std")]
27type NodeMapGuard<'a> = std::sync::MutexGuard<'a, BTreeMap<Vec<u8>, Vec<u8>>>;
28#[cfg(not(feature = "std"))]
29type NodeMapGuard<'a> = spin::MutexGuard<'a, BTreeMap<Vec<u8>, Vec<u8>>>;
30
31pub trait TrieDB: Send + Sync {
32    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError>;
33    /// Retrieves multiple values by key.
34    /// Returns results in the same order as the input keys.
35    /// Backends that support batched reads should override this for better
36    /// throughput. Default impl is equivalent to N independent `get` calls.
37    fn multi_get(&self, keys: &[Nibbles]) -> Vec<Result<Option<Vec<u8>>, TrieError>> {
38        keys.iter().map(|k| self.get(k.clone())).collect()
39    }
40    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError>;
41    // TODO: replace putbatch with this function.
42    fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
43        self.put_batch(
44            key_values
45                .iter()
46                .map(|node| (node.0.clone(), node.1.encode_to_vec()))
47                .collect(),
48        )
49    }
50    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
51        self.put_batch(vec![(key, value)])
52    }
53    /// Commits any pending changes to the underlying storage
54    /// For read-only or in-memory implementations, this is a no-op
55    fn commit(&self) -> Result<(), TrieError> {
56        Ok(())
57    }
58
59    fn flatkeyvalue_computed(&self, _key: Nibbles) -> bool {
60        false
61    }
62}
63
64// TODO: we should replace this with BackendTrieDB
65/// InMemory implementation for the TrieDB trait, with get and put operations.
66#[derive(Default)]
67pub struct InMemoryTrieDB {
68    inner: NodeMap,
69    prefix: Option<Nibbles>,
70}
71
72impl InMemoryTrieDB {
73    pub const fn new(map: NodeMap) -> Self {
74        Self {
75            inner: map,
76            prefix: None,
77        }
78    }
79
80    pub const fn new_with_prefix(map: NodeMap, prefix: Nibbles) -> Self {
81        Self {
82            inner: map,
83            prefix: Some(prefix),
84        }
85    }
86
87    pub fn new_empty() -> Self {
88        Self {
89            inner: Default::default(),
90            prefix: None,
91        }
92    }
93
94    fn apply_prefix(&self, path: Nibbles) -> Nibbles {
95        match &self.prefix {
96            Some(prefix) => prefix.concat(&path),
97            None => path,
98        }
99    }
100
101    // Locks the inner map, bridging the std (`Result`, poisonable) and spin
102    // (infallible, single-threaded guest) `Mutex::lock` APIs.
103    fn lock_inner(&self) -> Result<NodeMapGuard<'_>, TrieError> {
104        #[cfg(feature = "std")]
105        let guard = self.inner.lock().map_err(|_| TrieError::LockError)?;
106        #[cfg(not(feature = "std"))]
107        let guard = self.inner.lock();
108        Ok(guard)
109    }
110
111    // Do not remove or make private as we use this in ethrex-replay
112    pub fn inner(&self) -> NodeMap {
113        Arc::clone(&self.inner)
114    }
115}
116
117// `from_nodes` is host-only by design: nothing in it requires `std`, but the guest
118// never builds a trie from a node dump, so gating it keeps the no_std surface small.
119#[cfg(feature = "std")]
120impl InMemoryTrieDB {
121    // Do not remove or make private as we use this in ethrex-replay
122    pub fn from_nodes(
123        root_hash: H256,
124        state_nodes: &BTreeMap<H256, Node>,
125    ) -> Result<Self, TrieError> {
126        let mut embedded_root = Trie::get_embedded_root(state_nodes, root_hash)?;
127        let mut hashed_nodes = vec![];
128        embedded_root.commit(
129            Nibbles::default(),
130            &mut hashed_nodes,
131            &ethrex_crypto::NativeCrypto,
132        );
133
134        let hashed_nodes = hashed_nodes
135            .into_iter()
136            .map(|(k, v)| (k.into_vec(), v))
137            .collect();
138
139        let in_memory_trie = Arc::new(Mutex::new(hashed_nodes));
140        Ok(Self::new(in_memory_trie))
141    }
142}
143
144impl TrieDB for InMemoryTrieDB {
145    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
146        Ok(self
147            .lock_inner()?
148            .get(self.apply_prefix(key).as_ref())
149            .cloned())
150    }
151
152    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
153        let mut db = self.lock_inner()?;
154
155        for (key, value) in key_values {
156            let prefixed_key = self.apply_prefix(key);
157            db.insert(prefixed_key.into_vec(), value);
158        }
159
160        Ok(())
161    }
162}
ethrex_trie/db.rs

ethrex_trie/
db.rs
