arb_storage/
storage.rs

1use alloy_primitives::{keccak256, Address, B256, U256};
2use revm::Database;
3
4use crate::{
5    slot::{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 base_slice: &[u8] = if self.base_key == B256::ZERO {
45            &[]
46        } else {
47            self.base_key.as_slice()
48        };
49        let mut combined = Vec::with_capacity(base_slice.len() + sub_key.len());
50        combined.extend_from_slice(base_slice);
51        combined.extend_from_slice(sub_key);
52        let new_key = keccak256(&combined);
53        Storage::new_with_account(self.state, new_key, self.account)
54    }
55
56    /// Reads a 32-byte value by uint64 offset.
57    pub fn get_by_uint64(&self, offset: u64) -> Result<B256, ()> {
58        let slot = self.compute_slot(offset);
59        unsafe {
60            let state = &mut *self.state;
61            Ok(B256::from(read_storage_at(state, self.account, slot)))
62        }
63    }
64
65    /// Writes a 32-byte value by uint64 offset.
66    pub fn set_by_uint64(&self, offset: u64, value: B256) -> Result<(), ()> {
67        let slot = self.compute_slot(offset);
68        let value_u256 = U256::from_be_bytes(value.0);
69        unsafe {
70            let state = &mut *self.state;
71            write_storage_at(state, self.account, slot, value_u256);
72            Ok(())
73        }
74    }
75
76    /// Reads a U256 by uint64 offset.
77    pub fn get_uint64_by_uint64(&self, offset: u64) -> Result<u64, ()> {
78        let slot = self.compute_slot(offset);
79        unsafe {
80            let state = &mut *self.state;
81            let value = read_storage_at(state, self.account, slot);
82            Ok(value.try_into().unwrap_or(0))
83        }
84    }
85
86    /// Writes a U256 by uint64 offset.
87    pub fn set_uint64_by_uint64(&self, offset: u64, value: u64) -> Result<(), ()> {
88        let slot = self.compute_slot(offset);
89        unsafe {
90            let state = &mut *self.state;
91            write_storage_at(state, self.account, slot, U256::from(value));
92            Ok(())
93        }
94    }
95
96    /// Reads a 32-byte value by B256 key using mapAddress algorithm.
97    pub fn get(&self, key: B256) -> Result<B256, ()> {
98        let slot = self.compute_slot_for_key(key);
99        unsafe {
100            let state = &mut *self.state;
101            Ok(B256::from(read_storage_at(state, self.account, slot)))
102        }
103    }
104
105    /// Writes a 32-byte value by B256 key using mapAddress algorithm.
106    pub fn set(&self, key: B256, value: B256) -> Result<(), ()> {
107        let slot = self.compute_slot_for_key(key);
108        let value_u256 = U256::from_be_bytes(value.0);
109        unsafe {
110            let state = &mut *self.state;
111            write_storage_at(state, self.account, slot, value_u256);
112            Ok(())
113        }
114    }
115
116    fn storage_key(&self) -> &[u8] {
117        if self.base_key == B256::ZERO {
118            &[]
119        } else {
120            self.base_key.as_slice()
121        }
122    }
123
124    fn compute_slot(&self, offset: u64) -> U256 {
125        storage_key_map(self.storage_key(), offset)
126    }
127
128    fn compute_slot_for_key(&self, key: B256) -> U256 {
129        storage_key_map_b256(self.storage_key(), &key.0)
130    }
131
132    /// Creates a StorageSlot handle for a specific offset.
133    pub fn new_slot(&self, offset: u64) -> U256 {
134        self.compute_slot(offset)
135    }
136
137    /// Returns the raw state pointer.
138    pub fn state_ptr(&self) -> *mut revm::database::State<D> {
139        self.state
140    }
141
142    /// Returns the base key for this storage subspace.
143    pub fn base_key(&self) -> B256 {
144        self.base_key
145    }
146}
147
148impl<D> Clone for Storage<D> {
149    fn clone(&self) -> Self {
150        Self {
151            state: self.state,
152            base_key: self.base_key,
153            account: self.account,
154        }
155    }
156}
157
158// Safety: Storage is Send/Sync when D is, since state is only accessed
159// within a single execution context.
160unsafe impl<D: Send> Send for Storage<D> {}
161unsafe impl<D: Sync> Sync for Storage<D> {}
arb_storage/storage.rs

arb_storage/
storage.rs
