risc0_steel/
state.rs

1// Copyright 2025 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::Account 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                }))
168            }
169            None => {
170                self.account_storage.insert(address, None);
171
172                Ok(None)
173            }
174        }
175    }
176
177    /// Get account code by its hash.
178    #[inline]
179    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
180        let code = self.inner.code_by_hash(code_hash);
181        Ok(Bytecode::new_raw(code.clone()))
182    }
183
184    /// Get storage value of address at index.
185    #[inline]
186    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
187        let storage = self
188            .account_storage
189            .get(&address)
190            .unwrap_or_else(|| panic!("No storage trie with root: {address}"));
191        match storage {
192            Some(storage) => {
193                let val = storage
194                    .get_rlp(keccak256(index.to_be_bytes::<32>()))
195                    .expect("Invalid encoded storage value");
196                Ok(val.unwrap_or_default())
197            }
198            None => Ok(U256::ZERO),
199        }
200    }
201
202    /// Get block hash by block number.
203    #[inline]
204    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
205        Ok(self.inner.block_hash(number))
206    }
207}
208
209impl<H: crate::EvmBlockHeader> crate::EvmDatabase for WrapStateDb<'_, H> {
210    fn logs(
211        &mut self,
212        filter: alloy_rpc_types::Filter,
213    ) -> Result<Vec<Log>, <Self as RevmDatabase>::Error> {
214        assert_eq!(filter.get_block_hash(), Some(self.header.seal()));
215
216        let Some(logs) = self.inner.logs.as_ref() else {
217            // if no logs are stored in the DB, check that the Bloom filter proves non-existence
218            assert!(
219                !event::matches_filter(*self.header.logs_bloom(), &filter),
220                "No logs for matching filter"
221            );
222            return Ok(vec![]);
223        };
224
225        let params = alloy_rpc_types::FilteredParams::new(Some(filter));
226        let filtered = logs
227            .iter()
228            .filter(|log| params.filter_address(&log.address) && params.filter_topics(log.topics()))
229            .cloned()
230            .collect();
231
232        Ok(filtered)
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use alloy_consensus::constants::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
240    #[test]
241    fn default_account() {
242        let account: StateAccount = Default::default();
243        assert_eq!(account.storage_root, EMPTY_ROOT_HASH);
244        assert_eq!(account.code_hash, KECCAK_EMPTY);
245    }
246}
risc0_steel/state.rs

risc0_steel/
state.rs
