tempo_precompiles/storage/
thread_local.rs

1use alloy::primitives::{Address, LogData, U256};
2use alloy_evm::{Database, EvmInternals};
3use revm::{
4    context::{Block, CfgEnv, JournalTr},
5    state::{AccountInfo, Bytecode},
6};
7use scoped_tls::scoped_thread_local;
8use std::{cell::RefCell, fmt::Debug};
9use tempo_chainspec::hardfork::TempoHardfork;
10
11use crate::{
12    Precompile,
13    error::{Result, TempoPrecompileError},
14    storage::{PrecompileStorageProvider, evm::EvmPrecompileStorageProvider},
15};
16
17scoped_thread_local!(static STORAGE: RefCell<&mut dyn PrecompileStorageProvider>);
18
19/// Thread-local storage accessor that implements `PrecompileStorageProvider` without the trait bound.
20///
21/// This is the only type that exposes access to the thread-local `STORAGE` static.
22///
23/// # Important
24///
25/// Since it provides access to the current thread-local storage context, it MUST be used within
26/// a `StorageCtx::enter` closure.
27///
28/// # Sync with `PrecompileStorageProvider`
29///
30/// This type mirrors `PrecompileStorageProvider` methods but with split mutability:
31/// - Read operations (staticcall) take `&self`
32/// - Write operations take `&mut self`
33#[derive(Debug, Default, Clone, Copy)]
34pub struct StorageCtx;
35
36impl StorageCtx {
37    /// Enter storage context. All storage operations must happen within the closure.
38    ///
39    /// # IMPORTANT
40    ///
41    /// The caller must ensure that:
42    /// 1. Only one `enter` call is active at a time, in the same thread.
43    /// 2. If multiple storage providers are instantiated in parallel threads,
44    ///    they CANNOT point to the same storage addresses.
45    pub fn enter<S, R>(storage: &mut S, f: impl FnOnce() -> R) -> R
46    where
47        S: PrecompileStorageProvider,
48    {
49        // SAFETY: `scoped_tls` ensures the pointer is only accessible within the closure scope.
50        let storage: &mut dyn PrecompileStorageProvider = storage;
51        let storage_static: &mut (dyn PrecompileStorageProvider + 'static) =
52            unsafe { std::mem::transmute(storage) };
53        let cell = RefCell::new(storage_static);
54        STORAGE.set(&cell, f)
55    }
56
57    /// Execute an infallible function with access to the current thread-local storage provider.
58    ///
59    /// # Panics
60    /// Panics if no storage context is set.
61    fn with_storage<F, R>(f: F) -> R
62    where
63        F: FnOnce(&mut dyn PrecompileStorageProvider) -> R,
64    {
65        assert!(
66            STORAGE.is_set(),
67            "No storage context. 'StorageCtx::enter' must be called first"
68        );
69        STORAGE.with(|cell| {
70            // SAFETY: `scoped_tls` ensures the pointer is only accessible within the closure scope.
71            // Holding the guard prevents re-entrant borrows.
72            let mut guard = cell.borrow_mut();
73            f(&mut **guard)
74        })
75    }
76
77    /// Execute a (fallible) function with access to the current thread-local storage provider.
78    fn try_with_storage<F, R>(f: F) -> Result<R>
79    where
80        F: FnOnce(&mut dyn PrecompileStorageProvider) -> Result<R>,
81    {
82        if !STORAGE.is_set() {
83            return Err(TempoPrecompileError::Fatal(
84                "No storage context. 'StorageCtx::enter' must be called first".to_string(),
85            ));
86        }
87        STORAGE.with(|cell| {
88            // SAFETY: `scoped_tls` ensures the pointer is only accessible within the closure scope.
89            // Holding the guard prevents re-entrant borrows.
90            let mut guard = cell.borrow_mut();
91            f(&mut **guard)
92        })
93    }
94
95    // `PrecompileStorageProvider` methods (with modified mutability for read-only methods)
96
97    /// Executes a closure with access to the account info, returning the closure's result.
98    ///
99    /// This is an ergonomic wrapper that flattens the Result, avoiding double `?`.
100    pub fn with_account_info<T>(
101        &self,
102        address: Address,
103        mut f: impl FnMut(&AccountInfo) -> Result<T>,
104    ) -> Result<T> {
105        let mut result: Option<Result<T>> = None;
106        Self::try_with_storage(|s| {
107            s.with_account_info(address, &mut |info| {
108                result = Some(f(info));
109            })
110        })?;
111        result.unwrap()
112    }
113
114    pub fn chain_id(&self) -> u64 {
115        Self::with_storage(|s| s.chain_id())
116    }
117
118    pub fn timestamp(&self) -> U256 {
119        Self::with_storage(|s| s.timestamp())
120    }
121
122    pub fn beneficiary(&self) -> Address {
123        Self::with_storage(|s| s.beneficiary())
124    }
125
126    pub fn set_code(&mut self, address: Address, code: Bytecode) -> Result<()> {
127        Self::try_with_storage(|s| s.set_code(address, code))
128    }
129
130    pub fn sload(&self, address: Address, key: U256) -> Result<U256> {
131        Self::try_with_storage(|s| s.sload(address, key))
132    }
133
134    pub fn tload(&self, address: Address, key: U256) -> Result<U256> {
135        Self::try_with_storage(|s| s.tload(address, key))
136    }
137
138    pub fn sstore(&mut self, address: Address, key: U256, value: U256) -> Result<()> {
139        Self::try_with_storage(|s| s.sstore(address, key, value))
140    }
141
142    pub fn tstore(&mut self, address: Address, key: U256, value: U256) -> Result<()> {
143        Self::try_with_storage(|s| s.tstore(address, key, value))
144    }
145
146    pub fn emit_event(&mut self, address: Address, event: LogData) -> Result<()> {
147        Self::try_with_storage(|s| s.emit_event(address, event))
148    }
149
150    pub fn deduct_gas(&mut self, gas: u64) -> Result<()> {
151        Self::try_with_storage(|s| s.deduct_gas(gas))
152    }
153
154    pub fn refund_gas(&mut self, gas: i64) {
155        Self::with_storage(|s| s.refund_gas(gas))
156    }
157
158    pub fn gas_used(&self) -> u64 {
159        Self::with_storage(|s| s.gas_used())
160    }
161
162    pub fn gas_refunded(&self) -> i64 {
163        Self::with_storage(|s| s.gas_refunded())
164    }
165
166    pub fn spec(&self) -> TempoHardfork {
167        Self::with_storage(|s| s.spec())
168    }
169}
170
171impl<'evm> StorageCtx {
172    /// Generic entry point for EVM-like environments.
173    /// Sets up the storage provider and executes a closure within that context.
174    pub fn enter_evm<J, R>(
175        journal: &'evm mut J,
176        block_env: &'evm dyn Block,
177        cfg: &CfgEnv<TempoHardfork>,
178        f: impl FnOnce() -> R,
179    ) -> R
180    where
181        J: JournalTr<Database: Database> + Debug,
182    {
183        let internals = EvmInternals::new(journal, block_env);
184        let mut provider = EvmPrecompileStorageProvider::new_max_gas(internals, cfg);
185
186        // The core logic of setting up thread-local storage is here.
187        Self::enter(&mut provider, f)
188    }
189
190    /// Entry point for a "canonical" precompile (with unique known address).
191    pub fn enter_precompile<J, P, R>(
192        journal: &'evm mut J,
193        block_env: &'evm dyn Block,
194        cfg: &CfgEnv<TempoHardfork>,
195        f: impl FnOnce(P) -> R,
196    ) -> R
197    where
198        J: JournalTr<Database: Database> + Debug,
199        P: Precompile + Default,
200    {
201        // Delegate all the setup logic to `enter_evm`.
202        // We just need to provide a closure that `enter_evm` expects.
203        Self::enter_evm(journal, block_env, cfg, || f(P::default()))
204    }
205}
206
207#[cfg(any(test, feature = "test-utils"))]
208use crate::storage::hashmap::HashMapStorageProvider;
209
210#[cfg(any(test, feature = "test-utils"))]
211impl StorageCtx {
212    /// Returns a mutable reference to the underlying `HashMapStorageProvider`.
213    ///
214    /// NOTE: takes a non-mutable reference because it's internal. The mutability
215    /// of the storage operation is determined by the public function.
216    #[allow(clippy::mut_from_ref)]
217    fn as_hashmap(&self) -> &mut HashMapStorageProvider {
218        Self::with_storage(|s| {
219            // SAFETY: Test code always uses HashMapStorageProvider.
220            // Reference valid for duration of StorageCtx::enter closure.
221            unsafe {
222                extend_lifetime_mut(
223                    &mut *(s as *mut dyn PrecompileStorageProvider as *mut HashMapStorageProvider),
224                )
225            }
226        })
227    }
228
229    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
230    pub fn get_account_info(&self, address: Address) -> Option<&AccountInfo> {
231        self.as_hashmap().get_account_info(address)
232    }
233
234    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
235    pub fn get_events(&self, address: Address) -> &Vec<LogData> {
236        self.as_hashmap().get_events(address)
237    }
238
239    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
240    pub fn set_nonce(&mut self, address: Address, nonce: u64) {
241        self.as_hashmap().set_nonce(address, nonce)
242    }
243
244    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
245    pub fn set_timestamp(&mut self, timestamp: U256) {
246        self.as_hashmap().set_timestamp(timestamp)
247    }
248
249    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
250    pub fn set_beneficiary(&mut self, beneficiary: Address) {
251        self.as_hashmap().set_beneficiary(beneficiary)
252    }
253
254    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
255    pub fn set_spec(&mut self, spec: TempoHardfork) {
256        self.as_hashmap().set_spec(spec)
257    }
258
259    /// NOTE: assumes storage tests always use the `HashMapStorageProvider`
260    pub fn clear_transient(&mut self) {
261        self.as_hashmap().clear_transient()
262    }
263
264    /// Checks if a contract at the given address has bytecode deployed.
265    pub fn has_bytecode(&self, address: Address) -> bool {
266        if let Some(account_info) = self.get_account_info(address) {
267            !account_info.is_empty_code_hash()
268        } else {
269            false
270        }
271    }
272}
273
274/// Extends the lifetime of a mutable reference: `&'a mut T -> &'b mut T`
275///
276/// SAFETY: the caller must ensure the reference remains valid for the extended lifetime.
277#[cfg(any(test, feature = "test-utils"))]
278unsafe fn extend_lifetime_mut<'b, T: ?Sized>(r: &mut T) -> &'b mut T {
279    unsafe { &mut *(r as *mut T) }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    #[should_panic(expected = "already borrowed")]
288    fn test_reentrant_with_storage_panics() {
289        let mut storage = HashMapStorageProvider::new(1);
290        StorageCtx::enter(&mut storage, || {
291            // first borrow
292            StorageCtx::with_storage(|_| {
293                // re-entrant call should panic
294                StorageCtx::with_storage(|_| ())
295            })
296        });
297    }
298}
tempo_precompiles/storage/thread_local.rs

tempo_precompiles/storage/
thread_local.rs
