arbos/address_set/
mod.rs

1use alloy_primitives::{Address, B256, U256};
2use arb_storage::{
3    STORAGE_READ_GAS as STORAGE_READ_COST, STORAGE_WRITE_ZERO_GAS as STORAGE_WRITE_ZERO_COST,
4    Storage, StorageBackedAddress, StorageBackedUint64, StorageBackend, SystemStateBackend,
5};
6use revm::Database;
7
8mod error;
9pub use error::AddressSetError;
10
11fn write_cost(value: B256) -> u64 {
12    arb_storage::write_cost(value == B256::ZERO)
13}
14
15/// A set of addresses backed by ArbOS storage.
16///
17/// Layout: slot 0 = size, slots 1..size = addresses (as StorageBackedAddress).
18/// Sub-storage at key ]0\] maps address_hash → slot index.
19pub struct AddressSet<'a, D> {
20    backing_storage: Storage<'a, D>,
21    size: StorageBackedUint64,
22    by_address: Storage<'a, D>,
23}
24
25pub fn initialize_address_set<D: Database>(sto: &Storage<'_, D>) -> Result<(), AddressSetError> {
26    Ok(sto.set_by_uint64(0, B256::ZERO)?)
27}
28
29pub fn open_address_set<D>(sto: Storage<'_, D>) -> AddressSet<'_, D> {
30    let size = StorageBackedUint64::new(sto.base_key(), 0);
31    let by_address = sto.open_sub_storage(&[0u8]);
32    AddressSet {
33        backing_storage: sto,
34        size,
35        by_address,
36    }
37}
38
39impl<D> AddressSet<'_, D> {
40    pub fn size<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, AddressSetError> {
41        Ok(self.size.get(backend)?)
42    }
43
44    pub fn is_member<B: SystemStateBackend>(
45        &self,
46        backend: &mut B,
47        addr: Address,
48    ) -> Result<bool, AddressSetError> {
49        let value = self.by_address_get(backend, address_to_hash(addr))?;
50        Ok(value != B256::ZERO)
51    }
52
53    pub fn get_any_member<B: SystemStateBackend>(
54        &self,
55        backend: &mut B,
56    ) -> Result<Option<Address>, AddressSetError> {
57        let size = self.size.get(backend)?;
58        if size == 0 {
59            return Ok(None);
60        }
61        let sba = StorageBackedAddress::new(self.backing_storage.base_key(), 1);
62        Ok(sba.get(backend).map(Some)?)
63    }
64
65    pub fn clear<B: StorageBackend>(&self, backend: &mut B) -> Result<(), AddressSetError> {
66        let size = self.size.get(backend)?;
67        if size == 0 {
68            return Ok(());
69        }
70        for i in 1..=size {
71            let contents = self.backing_get_by_uint64(backend, i)?;
72            self.backing_set_by_uint64(backend, i, B256::ZERO)?;
73            self.by_address_set(backend, contents, B256::ZERO)?;
74        }
75        Ok(self.size.set(backend, 0)?)
76    }
77
78    pub fn all_members<B: SystemStateBackend>(
79        &self,
80        backend: &mut B,
81        max_num: u64,
82    ) -> Result<Vec<Address>, AddressSetError> {
83        let mut size = self.size.get(backend)?;
84        if size > max_num {
85            size = max_num;
86        }
87        let mut ret = Vec::with_capacity(size as usize);
88        for i in 0..size {
89            let sba = StorageBackedAddress::new(self.backing_storage.base_key(), i + 1);
90            ret.push(sba.get(backend)?);
91        }
92        Ok(ret)
93    }
94
95    pub fn clear_list<B: StorageBackend>(&self, backend: &mut B) -> Result<(), AddressSetError> {
96        let size = self.size.get(backend)?;
97        if size == 0 {
98            return Ok(());
99        }
100        for i in 1..=size {
101            self.backing_set_by_uint64(backend, i, B256::ZERO)?;
102        }
103        Ok(self.size.set(backend, 0)?)
104    }
105
106    pub fn rectify_mapping<B: StorageBackend>(
107        &self,
108        backend: &mut B,
109        addr: Address,
110    ) -> Result<(), AddressSetError> {
111        if !self.is_member(backend, addr)? {
112            return Err(AddressSetError::NotMember);
113        }
114
115        let addr_as_hash = address_to_hash(addr);
116        let slot = hash_to_uint64(self.by_address_get(backend, addr_as_hash)?);
117        let at_slot = self.backing_get_by_uint64(backend, slot)?;
118        let size = self.size.get(backend)?;
119
120        if at_slot == addr_as_hash && slot <= size {
121            return Err(AddressSetError::MappingAlreadyConsistent);
122        }
123
124        self.by_address_set(backend, addr_as_hash, B256::ZERO)?;
125        self.add(backend, addr)
126    }
127
128    pub fn add<B: StorageBackend>(
129        &self,
130        backend: &mut B,
131        addr: Address,
132    ) -> Result<(), AddressSetError> {
133        let present = self.is_member(backend, addr)?;
134        if present {
135            return Ok(());
136        }
137
138        let size = self.size.get(backend)?;
139        let slot = uint_to_hash(1 + size);
140        let addr_as_hash = address_to_hash(addr);
141
142        self.by_address_set(backend, addr_as_hash, slot)?;
143
144        let sba = StorageBackedAddress::new(self.backing_storage.base_key(), 1 + size);
145        sba.set(backend, addr)?;
146
147        Ok(self.size.set(backend, size + 1)?)
148    }
149
150    /// Removes `addr`, adding the value-dependent storage gas it consumes to `gas`.
151    pub fn remove<B: StorageBackend>(
152        &self,
153        backend: &mut B,
154        addr: Address,
155        arbos_version: u64,
156        gas: &mut u64,
157    ) -> Result<(), AddressSetError> {
158        let addr_as_hash = address_to_hash(addr);
159        let slot_hash = self.by_address_get(backend, addr_as_hash)?;
160        *gas += STORAGE_READ_COST;
161        let slot = hash_to_uint64(slot_hash);
162
163        if slot == 0 {
164            return Ok(());
165        }
166
167        self.by_address_set(backend, addr_as_hash, B256::ZERO)?;
168        *gas += STORAGE_WRITE_ZERO_COST;
169
170        let size = self.size.get(backend)?;
171        *gas += STORAGE_READ_COST;
172        if slot < size {
173            let at_size = self.backing_get_by_uint64(backend, size)?;
174            *gas += STORAGE_READ_COST;
175            self.backing_set_by_uint64(backend, slot, at_size)?;
176            *gas += write_cost(at_size);
177
178            if arbos_version >= 11 {
179                self.by_address_set(backend, at_size, uint_to_hash(slot))?;
180                *gas += write_cost(uint_to_hash(slot));
181            }
182        }
183
184        self.backing_set_by_uint64(backend, size, B256::ZERO)?;
185        *gas += STORAGE_WRITE_ZERO_COST;
186
187        let new_size = size - 1;
188        *gas += STORAGE_READ_COST;
189        self.size.set(backend, new_size)?;
190        *gas += write_cost(uint_to_hash(new_size));
191        Ok(())
192    }
193
194    fn by_address_get<B: SystemStateBackend>(
195        &self,
196        backend: &mut B,
197        key: B256,
198    ) -> Result<B256, AddressSetError> {
199        let slot = self.by_address.slot_for_key(key);
200        let value = backend
201            .sload_system(self.by_address.account(), slot)
202            .map_err(Into::into)?;
203        Ok(B256::from(value.to_be_bytes::<32>()))
204    }
205
206    fn by_address_set<B: StorageBackend>(
207        &self,
208        backend: &mut B,
209        key: B256,
210        value: B256,
211    ) -> Result<(), AddressSetError> {
212        let slot = self.by_address.slot_for_key(key);
213        backend
214            .sstore(
215                self.by_address.account(),
216                slot,
217                U256::from_be_bytes(value.0),
218            )
219            .map_err(Into::into)?;
220        Ok(())
221    }
222
223    fn backing_get_by_uint64<B: SystemStateBackend>(
224        &self,
225        backend: &mut B,
226        offset: u64,
227    ) -> Result<B256, AddressSetError> {
228        let slot = self.backing_storage.new_slot(offset);
229        let value = backend
230            .sload_system(self.backing_storage.account(), slot)
231            .map_err(Into::into)?;
232        Ok(B256::from(value.to_be_bytes::<32>()))
233    }
234
235    fn backing_set_by_uint64<B: StorageBackend>(
236        &self,
237        backend: &mut B,
238        offset: u64,
239        value: B256,
240    ) -> Result<(), AddressSetError> {
241        let slot = self.backing_storage.new_slot(offset);
242        backend
243            .sstore(
244                self.backing_storage.account(),
245                slot,
246                U256::from_be_bytes(value.0),
247            )
248            .map_err(Into::into)?;
249        Ok(())
250    }
251}
252
253fn address_to_hash(addr: Address) -> B256 {
254    let mut bytes = [0u8; 32];
255    bytes[12..32].copy_from_slice(addr.as_slice());
256    B256::from(bytes)
257}
258
259fn uint_to_hash(val: u64) -> B256 {
260    B256::from(U256::from(val))
261}
262
263fn hash_to_uint64(hash: B256) -> u64 {
264    U256::from_be_bytes(hash.0).to::<u64>()
265}
arbos/address_set/mod.rs

arbos/address_set/
mod.rs
