arb_storage/
storage.rs

1use alloy_primitives::{Address, B256, U256};
2use revm::Database;
3
4use crate::{
5    slot::{derive_sub_key, storage_key_map, storage_key_map_b256},
6    state_ops::{read_storage_at, write_storage_at, ARBOS_STATE_ADDRESS},
7};
8
9/// Hierarchical storage abstraction over EVM account state.
10///
11/// State lives in a specific account. Storage uses keccak256-based
12/// subspace derivation to create a hierarchical key space.
13pub struct Storage<D> {
14    pub state: *mut revm::database::State<D>,
15    pub base_key: B256,
16    pub account: Address,
17}
18
19impl<D: Database> Storage<D> {
20    /// Creates a new Storage backed by the ArbOS state account.
21    pub fn new(state: *mut revm::database::State<D>, base_key: B256) -> Self {
22        Self {
23            state,
24            base_key,
25            account: ARBOS_STATE_ADDRESS,
26        }
27    }
28
29    /// Creates a new Storage backed by a specific account.
30    pub fn new_with_account(
31        state: *mut revm::database::State<D>,
32        base_key: B256,
33        account: Address,
34    ) -> Self {
35        Self {
36            state,
37            base_key,
38            account,
39        }
40    }
41
42    /// Opens a child subspace by hashing the parent key with the child ID.
43    pub fn open_sub_storage(&self, sub_key: &[u8]) -> Storage<D> {
44        let new_key = derive_sub_key(self.base_key, sub_key);
45        Storage::new_with_account(self.state, new_key, self.account)
46    }
47
48    /// Opens a child subspace using a pre-derived key, avoiding a keccak hash.
49    pub fn open_sub_storage_with_key(&self, key: B256) -> Storage<D> {
50        Storage::new_with_account(self.state, key, self.account)
51    }
52
53    /// Reads a 32-byte value by uint64 offset.
54    pub fn get_by_uint64(&self, offset: u64) -> Result<B256, ()> {
55        let slot = self.compute_slot(offset);
56        unsafe {
57            let state = &mut *self.state;
58            Ok(B256::from(read_storage_at(state, self.account, slot)))
59        }
60    }
61
62    /// Writes a 32-byte value by uint64 offset.
63    pub fn set_by_uint64(&self, offset: u64, value: B256) -> Result<(), ()> {
64        let slot = self.compute_slot(offset);
65        let value_u256 = U256::from_be_bytes(value.0);
66        unsafe {
67            let state = &mut *self.state;
68            write_storage_at(state, self.account, slot, value_u256);
69            Ok(())
70        }
71    }
72
73    /// Reads a U256 by uint64 offset.
74    pub fn get_uint64_by_uint64(&self, offset: u64) -> Result<u64, ()> {
75        let slot = self.compute_slot(offset);
76        unsafe {
77            let state = &mut *self.state;
78            let value = read_storage_at(state, self.account, slot);
79            Ok(value.try_into().unwrap_or(0))
80        }
81    }
82
83    /// Writes a U256 by uint64 offset.
84    pub fn set_uint64_by_uint64(&self, offset: u64, value: u64) -> Result<(), ()> {
85        let slot = self.compute_slot(offset);
86        unsafe {
87            let state = &mut *self.state;
88            write_storage_at(state, self.account, slot, U256::from(value));
89            Ok(())
90        }
91    }
92
93    /// Reads a 32-byte value by B256 key using mapAddress algorithm.
94    pub fn get(&self, key: B256) -> Result<B256, ()> {
95        let slot = self.compute_slot_for_key(key);
96        unsafe {
97            let state = &mut *self.state;
98            Ok(B256::from(read_storage_at(state, self.account, slot)))
99        }
100    }
101
102    /// Writes a 32-byte value by B256 key using mapAddress algorithm.
103    pub fn set(&self, key: B256, value: B256) -> Result<(), ()> {
104        let slot = self.compute_slot_for_key(key);
105        let value_u256 = U256::from_be_bytes(value.0);
106        unsafe {
107            let state = &mut *self.state;
108            write_storage_at(state, self.account, slot, value_u256);
109            Ok(())
110        }
111    }
112
113    fn storage_key(&self) -> &[u8] {
114        if self.base_key == B256::ZERO {
115            &[]
116        } else {
117            self.base_key.as_slice()
118        }
119    }
120
121    fn compute_slot(&self, offset: u64) -> U256 {
122        storage_key_map(self.storage_key(), offset)
123    }
124
125    fn compute_slot_for_key(&self, key: B256) -> U256 {
126        storage_key_map_b256(self.storage_key(), &key.0)
127    }
128
129    /// Creates a StorageSlot handle for a specific offset.
130    pub fn new_slot(&self, offset: u64) -> U256 {
131        self.compute_slot(offset)
132    }
133
134    /// Returns the raw state pointer.
135    pub fn state_ptr(&self) -> *mut revm::database::State<D> {
136        self.state
137    }
138
139    /// Returns the base key for this storage subspace.
140    pub fn base_key(&self) -> B256 {
141        self.base_key
142    }
143}
144
145impl<D> Clone for Storage<D> {
146    fn clone(&self) -> Self {
147        Self {
148            state: self.state,
149            base_key: self.base_key,
150            account: self.account,
151        }
152    }
153}
154
155// Safety: Storage is Send/Sync when D is, since state is only accessed
156// within a single execution context.
157unsafe impl<D: Send> Send for Storage<D> {}
158unsafe impl<D: Sync> Sync for Storage<D> {}
arb_storage/storage.rs

arb_storage/
storage.rs
