Skip to main content

risc0_steel/
state.rs

1// Copyright 2026 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub use alloy_consensus::TrieAccount as StateAccount;
16
17use crate::{event, mpt::MerkleTrie};
18use alloy_primitives::{
19    keccak256,
20    map::{AddressHashMap, B256HashMap, HashMap},
21    Address, Bytes, Log, Sealed, B256, U256,
22};
23use revm::{
24    state::{AccountInfo, Bytecode},
25    Database as RevmDatabase,
26};
27use std::{
28    convert::Infallible,
29    fmt::{self, Debug},
30    rc::Rc,
31};
32
33/// A simple MPT-based read-only EVM database implementation.
34///
35/// It is backed by a single [MerkleTrie] for the accounts and one [MerkleTrie] each for the
36/// accounts' stores. It panics when querying data not contained in the tries. This design allows
37/// storage keys to be queried by storage trie root hash, but not by account, which is required to
38/// implement [DatabaseRef]. Thus, in order to use [StateDb] in revm, a wrapper must be
39/// used, which caches the appropriate storage trie root for each basic account query, thus
40/// requiring mutability.
41///
42/// [DatabaseRef]: revm::DatabaseRef
43#[derive(Debug)]
44pub struct StateDb {
45    /// State MPT.
46    state_trie: MerkleTrie,
47    /// Storage MPTs to their root hash.
48    /// [Rc] is used for MPT deduplication.
49    storage_tries: B256HashMap<Rc<MerkleTrie>>,
50    /// Contracts by their hash.
51    contracts: B256HashMap<Bytes>,
52    /// Block hashes by their number.
53    block_hashes: HashMap<u64, B256>,
54    /// All the logs for this block.
55    logs: Option<Vec<Log>>,
56}
57
58impl StateDb {
59    /// Creates a new state database from the given tries.
60    pub fn new(
61        state_trie: MerkleTrie,
62        storage_tries: impl IntoIterator<Item = MerkleTrie>,
63        contracts: impl IntoIterator<Item = Bytes>,
64        block_hashes: HashMap<u64, B256>,
65        logs: Option<Vec<Log>>,
66    ) -> Self {
67        let contracts = contracts
68            .into_iter()
69            .map(|code| (keccak256(&code), code))
70            .collect();
71        let storage_tries = storage_tries
72            .into_iter()
73            .map(|trie| (trie.hash_slow(), Rc::new(trie)))
74            .collect();
75
76        Self {
77            state_trie,
78            contracts,
79            storage_tries,
80            block_hashes,
81            logs,
82        }
83    }
84
85    #[inline]
86    pub(crate) fn account(&self, address: Address) -> Option<StateAccount> {
87        self.state_trie
88            .get_rlp(keccak256(address))
89            .expect("Invalid encoded state trie value")
90    }
91
92    #[inline]
93    pub(crate) fn code_by_hash(&self, hash: B256) -> &Bytes {
94        self.contracts
95            .get(&hash)
96            .unwrap_or_else(|| panic!("No code with hash: {hash}"))
97    }
98
99    #[inline]
100    pub(crate) fn block_hash(&self, number: u64) -> B256 {
101        let hash = self
102            .block_hashes
103            .get(&number)
104            .unwrap_or_else(|| panic!("No block with number: {number}"));
105        *hash
106    }
107
108    #[inline]
109    pub(crate) fn storage_trie(&self, root: &B256) -> Option<&Rc<MerkleTrie>> {
110        self.storage_tries.get(root)
111    }
112}
113
114/// A simple wrapper for [StateDb] to implement the [Database] trait.
115///
116/// In addition to translating the actual [Database] queries into MPT calls, it also maps account
117/// addresses to their respective storage trie when the account is first accessed. This works
118/// because [Database::basic] must always be called before any [Database::storage] calls for that
119/// account.
120pub struct WrapStateDb<'a, H> {
121    inner: &'a StateDb,
122    header: &'a Sealed<H>,
123    account_storage: AddressHashMap<Option<Rc<MerkleTrie>>>,
124}
125
126impl<'a, H> WrapStateDb<'a, H> {
127    /// Creates a new [Database] from the given [StateDb].
128    pub fn new(inner: &'a StateDb, header: &'a Sealed<H>) -> Self {
129        Self {
130            inner,
131            header,
132            account_storage: Default::default(),
133        }
134    }
135}
136
137impl<H> Debug for WrapStateDb<'_, H> {
138    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139        f.debug_struct("WrapStateDb")
140            .field("inner", &self.inner)
141            .field("block_hash", &self.header.seal())
142            .finish_non_exhaustive()
143    }
144}
145
146impl<H> RevmDatabase for WrapStateDb<'_, H> {
147    /// The [StateDb] does not return any errors.
148    type Error = Infallible;
149
150    /// Get basic account information.
151    #[inline]
152    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
153        let account = self.inner.account(address);
154        match account {
155            Some(account) => {
156                // link storage trie to the account, if it exists
157                if let Some(storage_trie) = self.inner.storage_trie(&account.storage_root) {
158                    self.account_storage
159                        .insert(address, Some(storage_trie.clone()));
160                }
161
162                Ok(Some(AccountInfo {
163                    balance: account.balance,
164                    nonce: account.nonce,
165                    code_hash: account.code_hash,
166                    code: None, // we don't need the code here, `code_by_hash` will be used instead
167                    account_id: None,
168                }))
169            }
170            None => {
171                self.account_storage.insert(address, None);
172
173                Ok(None)
174            }
175        }
176    }
177
178    /// Get account code by its hash.
179    #[inline]
180    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
181        let code = self.inner.code_by_hash(code_hash);
182        Ok(Bytecode::new_raw(code.clone()))
183    }
184
185    /// Get storage value of address at index.
186    #[inline]
187    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
188        let storage = self
189            .account_storage
190            .get(&address)
191            .unwrap_or_else(|| panic!("No storage trie with root: {address}"));
192        match storage {
193            Some(storage) => {
194                let val = storage
195                    .get_rlp(keccak256(index.to_be_bytes::<32>()))
196                    .expect("Invalid encoded storage value");
197                Ok(val.unwrap_or_default())
198            }
199            None => Ok(U256::ZERO),
200        }
201    }
202
203    /// Get block hash by block number.
204    #[inline]
205    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
206        Ok(self.inner.block_hash(number))
207    }
208}
209
210impl<H: crate::EvmBlockHeader> crate::EvmDatabase for WrapStateDb<'_, H> {
211    fn logs(
212        &mut self,
213        filter: alloy_rpc_types::Filter,
214    ) -> Result<Vec<Log>, <Self as RevmDatabase>::Error> {
215        assert_eq!(filter.get_block_hash(), Some(self.header.seal()));
216
217        let Some(logs) = self.inner.logs.as_ref() else {
218            // if no logs are stored in the DB, check that the Bloom filter proves non-existence
219            assert!(
220                !event::matches_filter(*self.header.logs_bloom(), &filter),
221                "No logs for matching filter"
222            );
223            return Ok(vec![]);
224        };
225
226        let params = alloy_rpc_types::FilteredParams::new(Some(filter));
227        let filtered = logs
228            .iter()
229            .filter(|log| params.filter_address(&log.address) && params.filter_topics(log.topics()))
230            .cloned()
231            .collect();
232
233        Ok(filtered)
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use alloy_consensus::constants::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
241    #[test]
242    fn default_account() {
243        let account: StateAccount = Default::default();
244        assert_eq!(account.storage_root, EMPTY_ROOT_HASH);
245        assert_eq!(account.code_hash, KECCAK_EMPTY);
246    }
247}