arb_storage/
backed_types.rs

1use alloy_primitives::{Address, B256, U256};
2use arb_storage_errors::StorageError;
3
4use crate::{
5    backend::{StorageBackend, SystemStateBackend},
6    slot::storage_key_map,
7    state_ops::ARBOS_STATE_ADDRESS,
8};
9
10fn compute_slot(base_key: B256, offset: u64) -> U256 {
11    if base_key == B256::ZERO {
12        storage_key_map(&[], offset)
13    } else {
14        storage_key_map(base_key.as_slice(), offset)
15    }
16}
17
18fn decode_address(slot: U256, value: U256) -> Result<Address, StorageError> {
19    let bytes = value.to_be_bytes::<32>();
20    let addr_bytes: [u8; 20] =
21        bytes[12..32]
22            .try_into()
23            .map_err(|_| StorageError::InvalidLayout {
24                slot,
25                reason: "address slot must be 20 bytes right-aligned",
26            })?;
27    Ok(Address::from(addr_bytes))
28}
29
30/// Storage-backed 64-bit unsigned integer.
31///
32/// Holds only the storage slot; the caller passes a [`StorageBackend`] at
33/// access time.
34#[derive(Clone, Copy, Debug)]
35pub struct StorageBackedUint64 {
36    pub slot: U256,
37}
38
39impl StorageBackedUint64 {
40    pub fn new(base_key: B256, offset: u64) -> Self {
41        Self {
42            slot: compute_slot(base_key, offset),
43        }
44    }
45
46    pub fn get<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, StorageError> {
47        let value = backend
48            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
49            .map_err(Into::into)?;
50        Ok(value.try_into().unwrap_or(0))
51    }
52
53    pub fn set<B: StorageBackend>(&self, backend: &mut B, value: u64) -> Result<(), StorageError> {
54        backend
55            .sstore(ARBOS_STATE_ADDRESS, self.slot, U256::from(value))
56            .map_err(Into::into)
57    }
58}
59/// Storage-backed 256-bit unsigned integer.
60///
61/// Holds only the storage slot; the caller passes a [`StorageBackend`] at
62/// access time.
63#[derive(Clone, Copy, Debug)]
64pub struct StorageBackedBigUint {
65    pub slot: U256,
66}
67
68impl StorageBackedBigUint {
69    pub fn new(base_key: B256, offset: u64) -> Self {
70        Self {
71            slot: compute_slot(base_key, offset),
72        }
73    }
74
75    pub fn get<B: SystemStateBackend>(&self, backend: &mut B) -> Result<U256, StorageError> {
76        backend
77            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
78            .map_err(Into::into)
79    }
80
81    pub fn set<B: StorageBackend>(&self, backend: &mut B, value: U256) -> Result<(), StorageError> {
82        backend
83            .sstore(ARBOS_STATE_ADDRESS, self.slot, value)
84            .map_err(Into::into)
85    }
86}
87
88/// Storage-backed Ethereum address (20 bytes, right-aligned in 32-byte slot).
89#[derive(Clone, Copy, Debug)]
90pub struct StorageBackedAddress {
91    pub slot: U256,
92}
93
94impl StorageBackedAddress {
95    pub fn new(base_key: B256, offset: u64) -> Self {
96        Self {
97            slot: compute_slot(base_key, offset),
98        }
99    }
100
101    pub fn get<B: SystemStateBackend>(&self, backend: &mut B) -> Result<Address, StorageError> {
102        let value = backend
103            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
104            .map_err(Into::into)?;
105        decode_address(self.slot, value)
106    }
107
108    pub fn set<B: StorageBackend>(
109        &self,
110        backend: &mut B,
111        value: Address,
112    ) -> Result<(), StorageError> {
113        let mut value_bytes = [0u8; 32];
114        value_bytes[12..32].copy_from_slice(value.as_slice());
115        backend
116            .sstore(
117                ARBOS_STATE_ADDRESS,
118                self.slot,
119                U256::from_be_bytes(value_bytes),
120            )
121            .map_err(Into::into)
122    }
123}
124
125/// Storage-backed signed 64-bit integer, bit-reinterpreting `i64` as `u64`.
126#[derive(Clone, Copy, Debug)]
127pub struct StorageBackedInt64 {
128    pub slot: U256,
129}
130
131impl StorageBackedInt64 {
132    pub fn new(base_key: B256, offset: u64) -> Self {
133        Self {
134            slot: compute_slot(base_key, offset),
135        }
136    }
137
138    pub fn get<B: SystemStateBackend>(&self, backend: &mut B) -> Result<i64, StorageError> {
139        let value = backend
140            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
141            .map_err(Into::into)?;
142        let value_u64: u64 = value.try_into().unwrap_or(0);
143        Ok(value_u64 as i64)
144    }
145
146    pub fn set<B: StorageBackend>(&self, backend: &mut B, value: i64) -> Result<(), StorageError> {
147        backend
148            .sstore(ARBOS_STATE_ADDRESS, self.slot, U256::from(value as u64))
149            .map_err(Into::into)
150    }
151}
152
153/// Storage-backed signed 256-bit integer using two's complement.
154#[derive(Clone, Copy, Debug)]
155pub struct StorageBackedBigInt {
156    pub slot: U256,
157}
158
159impl StorageBackedBigInt {
160    pub fn new(base_key: B256, offset: u64) -> Self {
161        Self {
162            slot: compute_slot(base_key, offset),
163        }
164    }
165
166    pub fn get_raw<B: SystemStateBackend>(&self, backend: &mut B) -> Result<U256, StorageError> {
167        backend
168            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
169            .map_err(Into::into)
170    }
171
172    pub fn is_negative<B: SystemStateBackend>(
173        &self,
174        backend: &mut B,
175    ) -> Result<bool, StorageError> {
176        Ok(self.get_raw(backend)?.bit(255))
177    }
178
179    /// Returns `(magnitude, is_negative)` decoded from two's complement.
180    pub fn get_signed<B: SystemStateBackend>(
181        &self,
182        backend: &mut B,
183    ) -> Result<(U256, bool), StorageError> {
184        let raw = self.get_raw(backend)?;
185        if raw.bit(255) {
186            let magnitude = (!raw).wrapping_add(U256::from(1));
187            Ok((magnitude, true))
188        } else {
189            Ok((raw, false))
190        }
191    }
192
193    pub fn set<B: StorageBackend>(&self, backend: &mut B, value: U256) -> Result<(), StorageError> {
194        backend
195            .sstore(ARBOS_STATE_ADDRESS, self.slot, value)
196            .map_err(Into::into)
197    }
198
199    pub fn set_negative<B: StorageBackend>(
200        &self,
201        backend: &mut B,
202        magnitude: U256,
203    ) -> Result<(), StorageError> {
204        let neg_value = (!magnitude).wrapping_add(U256::from(1));
205        self.set(backend, neg_value)
206    }
207}
208
209/// Sentinel value for nil addresses: `1 << 255`.
210fn nil_address_representation() -> U256 {
211    U256::from(1u64) << 255
212}
213
214/// Storage-backed optional address, using `1 << 255` to represent `None`.
215#[derive(Clone, Copy, Debug)]
216pub struct StorageBackedAddressOrNil {
217    pub slot: U256,
218}
219
220impl StorageBackedAddressOrNil {
221    pub fn new(base_key: B256, offset: u64) -> Self {
222        Self {
223            slot: compute_slot(base_key, offset),
224        }
225    }
226
227    pub fn get<B: SystemStateBackend>(
228        &self,
229        backend: &mut B,
230    ) -> Result<Option<Address>, StorageError> {
231        let value = backend
232            .sload_system(ARBOS_STATE_ADDRESS, self.slot)
233            .map_err(Into::into)?;
234        if value == nil_address_representation() {
235            return Ok(None);
236        }
237        decode_address(self.slot, value).map(Some)
238    }
239
240    pub fn set<B: StorageBackend>(
241        &self,
242        backend: &mut B,
243        value: Option<Address>,
244    ) -> Result<(), StorageError> {
245        let value_u256 = match value {
246            None => nil_address_representation(),
247            Some(addr) => {
248                let mut bytes = [0u8; 32];
249                bytes[12..32].copy_from_slice(addr.as_slice());
250                U256::from_be_bytes(bytes)
251            }
252        };
253        backend
254            .sstore(ARBOS_STATE_ADDRESS, self.slot, value_u256)
255            .map_err(Into::into)
256    }
257}