Skip to main content

phos_data_network_precompiles/storage/
thread_local.rs

1use std::cell::RefCell;
2
3use alloy_primitives::{Address, U256};
4use scoped_tls::scoped_thread_local;
5
6use crate::{
7    error::{DataNetworkPrecompileError, Result},
8    storage::PrecompileStorageProvider,
9};
10
11scoped_thread_local!(static STORAGE: RefCell<&mut dyn PrecompileStorageProvider>);
12
13/// Thread-local storage accessor that exposes the active precompile storage provider.
14#[derive(Debug, Default, Clone, Copy)]
15pub struct StorageCtx;
16
17impl StorageCtx {
18    /// Enters a storage context for the duration of the supplied closure.
19    pub fn enter<S, R>(storage: &mut S, f: impl FnOnce() -> R) -> R
20    where
21        S: PrecompileStorageProvider,
22    {
23        let storage: &mut dyn PrecompileStorageProvider = storage;
24        let storage_static: &mut (dyn PrecompileStorageProvider + 'static) =
25            unsafe { std::mem::transmute(storage) };
26        let cell = RefCell::new(storage_static);
27        STORAGE.set(&cell, f)
28    }
29
30    /// Executes an infallible function with the active storage provider.
31    fn with_storage<F, R>(f: F) -> R
32    where
33        F: FnOnce(&mut dyn PrecompileStorageProvider) -> R,
34    {
35        assert!(
36            STORAGE.is_set(),
37            "No storage context. 'StorageCtx::enter' must be called first"
38        );
39        STORAGE.with(|cell| {
40            let mut guard = cell.borrow_mut();
41            f(&mut **guard)
42        })
43    }
44
45    /// Executes a fallible function with the active storage provider.
46    fn try_with_storage<F, R>(f: F) -> Result<R>
47    where
48        F: FnOnce(&mut dyn PrecompileStorageProvider) -> Result<R>,
49    {
50        if !STORAGE.is_set() {
51            return Err(DataNetworkPrecompileError::Fatal(
52                "No storage context. 'StorageCtx::enter' must be called first".to_string(),
53            ));
54        }
55        STORAGE.with(|cell| {
56            // Holding the guard prevents re-entrant mutable borrows.
57            let mut guard = cell.borrow_mut();
58            f(&mut **guard)
59        })
60    }
61
62    /// Performs an SLOAD operation.
63    pub fn sload(&self, address: Address, key: U256) -> Result<U256> {
64        Self::try_with_storage(|storage| storage.sload(address, key))
65    }
66
67    /// Performs an SSTORE operation.
68    pub fn sstore(&mut self, address: Address, key: U256, value: U256) -> Result<()> {
69        Self::try_with_storage(|storage| storage.sstore(address, key, value))
70    }
71
72    /// Returns whether the current call context is static.
73    pub fn is_static(&self) -> bool {
74        Self::with_storage(|storage| storage.is_static())
75    }
76}