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    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError>;
34    // TODO: replace putbatch with this function.
35    fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
36        self.put_batch(
37            key_values
38                .iter()
39                .map(|node| (node.0.clone(), node.1.encode_to_vec()))
40                .collect(),
41        )
42    }
43    fn put(&self, key: Nibbles, value: Vec<u8>) -> Result<(), TrieError> {
44        self.put_batch(vec![(key, value)])
45    }
46    /// Commits any pending changes to the underlying storage
47    /// For read-only or in-memory implementations, this is a no-op
48    fn commit(&self) -> Result<(), TrieError> {
49        Ok(())
50    }
51
52    fn flatkeyvalue_computed(&self, _key: Nibbles) -> bool {
53        false
54    }
55}
56
57// TODO: we should replace this with BackendTrieDB
58/// InMemory implementation for the TrieDB trait, with get and put operations.
59#[derive(Default)]
60pub struct InMemoryTrieDB {
61    inner: NodeMap,
62    prefix: Option<Nibbles>,
63}
64
65impl InMemoryTrieDB {
66    pub const fn new(map: NodeMap) -> Self {
67        Self {
68            inner: map,
69            prefix: None,
70        }
71    }
72
73    pub const fn new_with_prefix(map: NodeMap, prefix: Nibbles) -> Self {
74        Self {
75            inner: map,
76            prefix: Some(prefix),
77        }
78    }
79
80    pub fn new_empty() -> Self {
81        Self {
82            inner: Default::default(),
83            prefix: None,
84        }
85    }
86
87    fn apply_prefix(&self, path: Nibbles) -> Nibbles {
88        match &self.prefix {
89            Some(prefix) => prefix.concat(&path),
90            None => path,
91        }
92    }
93
94    // Locks the inner map, bridging the std (`Result`, poisonable) and spin
95    // (infallible, single-threaded guest) `Mutex::lock` APIs.
96    fn lock_inner(&self) -> Result<NodeMapGuard<'_>, TrieError> {
97        #[cfg(feature = "std")]
98        let guard = self.inner.lock().map_err(|_| TrieError::LockError)?;
99        #[cfg(not(feature = "std"))]
100        let guard = self.inner.lock();
101        Ok(guard)
102    }
103
104    // Do not remove or make private as we use this in ethrex-replay
105    pub fn inner(&self) -> NodeMap {
106        Arc::clone(&self.inner)
107    }
108}
109
110// `from_nodes` is host-only by design: nothing in it requires `std`, but the guest
111// never builds a trie from a node dump, so gating it keeps the no_std surface small.
112#[cfg(feature = "std")]
113impl InMemoryTrieDB {
114    // Do not remove or make private as we use this in ethrex-replay
115    pub fn from_nodes(
116        root_hash: H256,
117        state_nodes: &BTreeMap<H256, Node>,
118    ) -> Result<Self, TrieError> {
119        let mut embedded_root = Trie::get_embedded_root(state_nodes, root_hash)?;
120        let mut hashed_nodes = vec![];
121        embedded_root.commit(
122            Nibbles::default(),
123            &mut hashed_nodes,
124            &ethrex_crypto::NativeCrypto,
125        );
126
127        let hashed_nodes = hashed_nodes
128            .into_iter()
129            .map(|(k, v)| (k.into_vec(), v))
130            .collect();
131
132        let in_memory_trie = Arc::new(Mutex::new(hashed_nodes));
133        Ok(Self::new(in_memory_trie))
134    }
135}
136
137impl TrieDB for InMemoryTrieDB {
138    fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
139        Ok(self
140            .lock_inner()?
141            .get(self.apply_prefix(key).as_ref())
142            .cloned())
143    }
144
145    fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError> {
146        let mut db = self.lock_inner()?;
147
148        for (key, value) in key_values {
149            let prefixed_key = self.apply_prefix(key);
150            db.insert(prefixed_key.into_vec(), value);
151        }
152
153        Ok(())
154    }
155}
ethrex_trie/db.rs

ethrex_trie/
db.rs
