tempo_precompiles/storage/
hashmap.rs

1use alloy::primitives::{Address, LogData, U256};
2use revm::{
3    context::journaled_state::JournalCheckpoint,
4    state::{AccountInfo, Bytecode},
5};
6use std::collections::HashMap;
7use tempo_chainspec::hardfork::TempoHardfork;
8
9use crate::{error::TempoPrecompileError, storage::PrecompileStorageProvider};
10
11/// In-memory [`PrecompileStorageProvider`] for unit tests.
12///
13/// Stores all state in `HashMap`s, avoiding the need for a real EVM context.
14pub struct HashMapStorageProvider {
15    internals: HashMap<(Address, U256), U256>,
16    transient: HashMap<(Address, U256), U256>,
17    accounts: HashMap<Address, AccountInfo>,
18    fail_on_sload: Option<(Address, U256)>,
19    chain_id: u64,
20    timestamp: U256,
21    beneficiary: Address,
22    block_number: u64,
23    spec: TempoHardfork,
24    amsterdam_eip8037_enabled: bool,
25    is_static: bool,
26    counter_sload: u64,
27    counter_sstore: u64,
28    snapshots: Vec<Snapshot>,
29
30    /// Emitted events keyed by contract address.
31    pub events: HashMap<Address, Vec<LogData>>,
32}
33
34/// Snapshot of mutable state for checkpoint/revert support.
35///
36/// PERF: naive cloning strategy due to its limited usage.
37struct Snapshot {
38    internals: HashMap<(Address, U256), U256>,
39    events: HashMap<Address, Vec<LogData>>,
40}
41
42impl HashMapStorageProvider {
43    /// Creates a new provider with the given chain ID and default hardfork.
44    pub fn new(chain_id: u64) -> Self {
45        Self::new_with_spec(chain_id, TempoHardfork::default())
46    }
47
48    /// Creates a new provider with the given chain ID and hardfork spec.
49    pub fn new_with_spec(chain_id: u64, spec: TempoHardfork) -> Self {
50        Self {
51            internals: HashMap::new(),
52            transient: HashMap::new(),
53            accounts: HashMap::new(),
54            fail_on_sload: None,
55            events: HashMap::new(),
56            snapshots: Vec::new(),
57            chain_id,
58            #[expect(clippy::disallowed_methods)]
59            timestamp: U256::from(
60                std::time::SystemTime::now()
61                    .duration_since(std::time::UNIX_EPOCH)
62                    .unwrap()
63                    .as_secs(),
64            ),
65            beneficiary: Address::ZERO,
66            block_number: 0,
67            spec,
68            amsterdam_eip8037_enabled: false,
69            is_static: false,
70            counter_sload: 0,
71            counter_sstore: 0,
72        }
73    }
74
75    /// Returns self with the hardfork spec overridden (builder pattern).
76    pub fn with_spec(mut self, spec: TempoHardfork) -> Self {
77        self.spec = spec;
78        self
79    }
80
81    /// Returns self with `amsterdam_eip8037_enabled` overridden (builder pattern).
82    pub fn with_amsterdam_eip8037_enabled(mut self, enabled: bool) -> Self {
83        self.amsterdam_eip8037_enabled = enabled;
84        self
85    }
86}
87
88impl PrecompileStorageProvider for HashMapStorageProvider {
89    fn chain_id(&self) -> u64 {
90        self.chain_id
91    }
92
93    fn timestamp(&self) -> U256 {
94        self.timestamp
95    }
96
97    fn beneficiary(&self) -> Address {
98        self.beneficiary
99    }
100
101    fn block_number(&self) -> u64 {
102        self.block_number
103    }
104
105    fn set_code(&mut self, address: Address, code: Bytecode) -> Result<(), TempoPrecompileError> {
106        let account = self.accounts.entry(address).or_default();
107        account.code_hash = code.hash_slow();
108        account.code = Some(code);
109        Ok(())
110    }
111
112    fn with_account_info(
113        &mut self,
114        address: Address,
115        f: &mut dyn FnMut(&AccountInfo),
116    ) -> Result<(), TempoPrecompileError> {
117        let account = self.accounts.entry(address).or_default();
118        f(&*account);
119        Ok(())
120    }
121
122    fn sstore(
123        &mut self,
124        address: Address,
125        key: U256,
126        value: U256,
127    ) -> Result<(), TempoPrecompileError> {
128        self.counter_sstore += 1;
129        self.internals.insert((address, key), value);
130        Ok(())
131    }
132
133    fn tstore(
134        &mut self,
135        address: Address,
136        key: U256,
137        value: U256,
138    ) -> Result<(), TempoPrecompileError> {
139        self.transient.insert((address, key), value);
140        Ok(())
141    }
142
143    fn emit_event(&mut self, address: Address, event: LogData) -> Result<(), TempoPrecompileError> {
144        self.events.entry(address).or_default().push(event);
145        Ok(())
146    }
147
148    fn sload(&mut self, address: Address, key: U256) -> Result<U256, TempoPrecompileError> {
149        if self.fail_on_sload == Some((address, key)) {
150            return Err(TempoPrecompileError::Fatal("injected sload failure".into()));
151        }
152
153        self.counter_sload += 1;
154        Ok(self
155            .internals
156            .get(&(address, key))
157            .copied()
158            .unwrap_or(U256::ZERO))
159    }
160
161    fn tload(&mut self, address: Address, key: U256) -> Result<U256, TempoPrecompileError> {
162        Ok(self
163            .transient
164            .get(&(address, key))
165            .copied()
166            .unwrap_or(U256::ZERO))
167    }
168
169    fn deduct_gas(&mut self, _gas: u64) -> Result<(), TempoPrecompileError> {
170        Ok(())
171    }
172
173    fn refund_gas(&mut self, _gas: i64) {
174        // No-op
175    }
176
177    fn gas_limit(&self) -> u64 {
178        0
179    }
180
181    fn gas_used(&self) -> u64 {
182        0
183    }
184
185    fn state_gas_used(&self) -> u64 {
186        0
187    }
188
189    fn gas_refunded(&self) -> i64 {
190        0
191    }
192
193    fn reservoir(&self) -> u64 {
194        0
195    }
196
197    fn spec(&self) -> TempoHardfork {
198        self.spec
199    }
200
201    fn amsterdam_eip8037_enabled(&self) -> bool {
202        self.amsterdam_eip8037_enabled
203    }
204
205    fn is_static(&self) -> bool {
206        self.is_static
207    }
208
209    fn checkpoint(&mut self) -> JournalCheckpoint {
210        let idx = self.snapshots.len();
211        self.snapshots.push(Snapshot {
212            internals: self.internals.clone(),
213            events: self.events.clone(),
214        });
215        JournalCheckpoint {
216            log_i: 0,
217            journal_i: idx,
218            selfdestructed_i: 0,
219        }
220    }
221
222    fn checkpoint_commit(&mut self, checkpoint: JournalCheckpoint) {
223        assert_eq!(
224            checkpoint.journal_i,
225            self.snapshots.len() - 1,
226            "out-of-order checkpoint commit (expected top of stack)"
227        );
228        self.snapshots.pop();
229    }
230
231    fn checkpoint_revert(&mut self, checkpoint: JournalCheckpoint) {
232        assert_eq!(
233            checkpoint.journal_i,
234            self.snapshots.len() - 1,
235            "out-of-order checkpoint revert (expected top of stack)"
236        );
237        if let Some(snapshot) = self.snapshots.drain(checkpoint.journal_i..).next() {
238            self.internals = snapshot.internals;
239            self.events = snapshot.events;
240        }
241    }
242}
243
244#[cfg(any(test, feature = "test-utils"))]
245impl HashMapStorageProvider {
246    pub fn fail_next_sload_at(&mut self, address: Address, slot: U256) {
247        self.fail_on_sload = Some((address, slot));
248    }
249
250    /// Returns the account info for the given address, if it exists.
251    pub fn get_account_info(&self, address: Address) -> Option<&AccountInfo> {
252        self.accounts.get(&address)
253    }
254
255    /// Returns all emitted events for the given address.
256    pub fn get_events(&self, address: Address) -> &Vec<LogData> {
257        static EMPTY: Vec<LogData> = Vec::new();
258        self.events.get(&address).unwrap_or(&EMPTY)
259    }
260
261    /// Sets the nonce for the given address.
262    pub fn set_nonce(&mut self, address: Address, nonce: u64) {
263        let account = self.accounts.entry(address).or_default();
264        account.nonce = nonce;
265    }
266
267    /// Overrides the block timestamp.
268    pub fn set_timestamp(&mut self, timestamp: U256) {
269        self.timestamp = timestamp;
270    }
271
272    /// Overrides the block beneficiary (coinbase).
273    pub fn set_beneficiary(&mut self, beneficiary: Address) {
274        self.beneficiary = beneficiary;
275    }
276
277    /// Overrides the block number.
278    pub fn set_block_number(&mut self, block_number: u64) {
279        self.block_number = block_number;
280    }
281
282    /// Overrides the active hardfork spec.
283    pub fn set_spec(&mut self, spec: TempoHardfork) {
284        self.spec = spec;
285    }
286
287    /// Clears all transient storage (simulates a new block).
288    pub fn clear_transient(&mut self) {
289        self.transient.clear();
290    }
291
292    /// Clears all emitted events for the given address.
293    pub fn clear_events(&mut self, address: Address) {
294        let _ = self
295            .events
296            .entry(address)
297            .and_modify(|v| v.clear())
298            .or_default();
299    }
300
301    /// Returns the amount of counted SLOADs.
302    pub fn counter_sload(&self) -> u64 {
303        self.counter_sload
304    }
305
306    /// Returns the amount of counted SSTOREs.
307    pub fn counter_sstore(&self) -> u64 {
308        self.counter_sstore
309    }
310
311    /// Resets the SLOAD and SSTORE counters.
312    pub fn reset_counters(&mut self) {
313        self.counter_sload = 0;
314        self.counter_sstore = 0;
315    }
316
317    /// Returns all storage entries as `(address, slot, value)`.
318    pub fn into_storage(self) -> impl Iterator<Item = (Address, U256, U256)> {
319        self.internals
320            .into_iter()
321            .map(|((addr, slot), value)| (addr, slot, value))
322    }
323}
tempo_precompiles/storage/hashmap.rs

tempo_precompiles/storage/
hashmap.rs
