phos_data_network_precompiles/storage/
thread_local.rs1use 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#[derive(Debug, Default, Clone, Copy)]
15pub struct StorageCtx;
16
17impl StorageCtx {
18 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 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 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 let mut guard = cell.borrow_mut();
58 f(&mut **guard)
59 })
60 }
61
62 pub fn sload(&self, address: Address, key: U256) -> Result<U256> {
64 Self::try_with_storage(|storage| storage.sload(address, key))
65 }
66
67 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 pub fn is_static(&self) -> bool {
74 Self::with_storage(|storage| storage.is_static())
75 }
76}