arb_storage/
extra_types.rs

1use alloy_primitives::{B256, U256};
2use revm::Database;
3
4use crate::{
5    slot::storage_key_map,
6    state_ops::{read_arbos_storage, write_arbos_storage},
7};
8
9fn compute_slot(base_key: B256, offset: u64) -> U256 {
10    if base_key == B256::ZERO {
11        storage_key_map(&[], offset)
12    } else {
13        storage_key_map(base_key.as_slice(), offset)
14    }
15}
16
17fn read_slot<D: Database>(state: *mut revm::database::State<D>, slot: U256) -> Result<U256, ()> {
18    unsafe {
19        let state = &mut *state;
20        Ok(read_arbos_storage(state, slot))
21    }
22}
23
24fn write_slot<D: Database>(
25    state: *mut revm::database::State<D>,
26    slot: U256,
27    value: U256,
28) -> Result<(), ()> {
29    unsafe {
30        let state = &mut *state;
31        write_arbos_storage(state, slot, value);
32        Ok(())
33    }
34}
35
36/// Basis points stored as signed i64. 10000 bips = 100%.
37pub struct StorageBackedBips<D> {
38    state: *mut revm::database::State<D>,
39    slot: U256,
40}
41
42impl<D: Database> StorageBackedBips<D> {
43    pub fn new(state: *mut revm::database::State<D>, base_key: B256, offset: u64) -> Self {
44        Self {
45            state,
46            slot: compute_slot(base_key, offset),
47        }
48    }
49
50    pub fn get(&self) -> Result<i64, ()> {
51        let value = read_slot(self.state, self.slot)?;
52        let value_u64: u64 = value.try_into().unwrap_or(0);
53        Ok(value_u64 as i64)
54    }
55
56    pub fn set(&self, value: i64) -> Result<(), ()> {
57        write_slot(self.state, self.slot, U256::from(value as u64))
58    }
59}
60
61impl<D> Clone for StorageBackedBips<D> {
62    fn clone(&self) -> Self {
63        Self {
64            state: self.state,
65            slot: self.slot,
66        }
67    }
68}
69
70unsafe impl<D: Send> Send for StorageBackedBips<D> {}
71unsafe impl<D: Sync> Sync for StorageBackedBips<D> {}
72
73/// Unsigned basis points stored as u64. 10000 ubips = 100%.
74pub struct StorageBackedUBips<D> {
75    state: *mut revm::database::State<D>,
76    slot: U256,
77}
78
79impl<D: Database> StorageBackedUBips<D> {
80    pub fn new(state: *mut revm::database::State<D>, base_key: B256, offset: u64) -> Self {
81        Self {
82            state,
83            slot: compute_slot(base_key, offset),
84        }
85    }
86
87    pub fn get(&self) -> Result<u64, ()> {
88        let value = read_slot(self.state, self.slot)?;
89        Ok(value.try_into().unwrap_or(0))
90    }
91
92    pub fn set(&self, value: u64) -> Result<(), ()> {
93        write_slot(self.state, self.slot, U256::from(value))
94    }
95}
96
97impl<D> Clone for StorageBackedUBips<D> {
98    fn clone(&self) -> Self {
99        Self {
100            state: self.state,
101            slot: self.slot,
102        }
103    }
104}
105
106unsafe impl<D: Send> Send for StorageBackedUBips<D> {}
107unsafe impl<D: Sync> Sync for StorageBackedUBips<D> {}
108
109/// Storage-backed 16-bit unsigned integer.
110pub struct StorageBackedUint16<D> {
111    state: *mut revm::database::State<D>,
112    slot: U256,
113}
114
115impl<D: Database> StorageBackedUint16<D> {
116    pub fn new(state: *mut revm::database::State<D>, base_key: B256, offset: u64) -> Self {
117        Self {
118            state,
119            slot: compute_slot(base_key, offset),
120        }
121    }
122
123    pub fn get(&self) -> Result<u16, ()> {
124        let value = read_slot(self.state, self.slot)?;
125        Ok(value.try_into().unwrap_or(0))
126    }
127
128    pub fn set(&self, value: u16) -> Result<(), ()> {
129        write_slot(self.state, self.slot, U256::from(value))
130    }
131}
132
133impl<D> Clone for StorageBackedUint16<D> {
134    fn clone(&self) -> Self {
135        Self {
136            state: self.state,
137            slot: self.slot,
138        }
139    }
140}
141
142unsafe impl<D: Send> Send for StorageBackedUint16<D> {}
143unsafe impl<D: Sync> Sync for StorageBackedUint16<D> {}
144
145/// Storage-backed 24-bit unsigned integer.
146pub struct StorageBackedUint24<D> {
147    state: *mut revm::database::State<D>,
148    slot: U256,
149}
150
151impl<D: Database> StorageBackedUint24<D> {
152    pub fn new(state: *mut revm::database::State<D>, base_key: B256, offset: u64) -> Self {
153        Self {
154            state,
155            slot: compute_slot(base_key, offset),
156        }
157    }
158
159    pub fn get(&self) -> Result<u32, ()> {
160        let value = read_slot(self.state, self.slot)?;
161        let raw: u32 = value.try_into().unwrap_or(0);
162        Ok(raw & 0xFF_FFFF)
163    }
164
165    pub fn set(&self, value: u32) -> Result<(), ()> {
166        write_slot(self.state, self.slot, U256::from(value & 0xFF_FFFF))
167    }
168}
169
170impl<D> Clone for StorageBackedUint24<D> {
171    fn clone(&self) -> Self {
172        Self {
173            state: self.state,
174            slot: self.slot,
175        }
176    }
177}
178
179unsafe impl<D: Send> Send for StorageBackedUint24<D> {}
180unsafe impl<D: Sync> Sync for StorageBackedUint24<D> {}
181
182/// Storage-backed 32-bit unsigned integer.
183pub struct StorageBackedUint32<D> {
184    state: *mut revm::database::State<D>,
185    slot: U256,
186}
187
188impl<D: Database> StorageBackedUint32<D> {
189    pub fn new(state: *mut revm::database::State<D>, base_key: B256, offset: u64) -> Self {
190        Self {
191            state,
192            slot: compute_slot(base_key, offset),
193        }
194    }
195
196    pub fn get(&self) -> Result<u32, ()> {
197        let value = read_slot(self.state, self.slot)?;
198        Ok(value.try_into().unwrap_or(0))
199    }
200
201    pub fn set(&self, value: u32) -> Result<(), ()> {
202        write_slot(self.state, self.slot, U256::from(value))
203    }
204
205    pub fn clear(&self) -> Result<(), ()> {
206        self.set(0)
207    }
208}
209
210impl<D> Clone for StorageBackedUint32<D> {
211    fn clone(&self) -> Self {
212        Self {
213            state: self.state,
214            slot: self.slot,
215        }
216    }
217}
218
219unsafe impl<D: Send> Send for StorageBackedUint32<D> {}
220unsafe impl<D: Sync> Sync for StorageBackedUint32<D> {}