arbos/address_table/
mod.rs

1use alloy_primitives::{Address, B256, U256};
2use alloy_rlp::{Decodable, Encodable, Header};
3use arb_storage::{Storage, StorageBackedUint64, StorageBackend, SystemStateBackend};
4use revm::Database;
5
6mod error;
7pub use error::AddressTableError;
8
9/// A mapping between addresses and compact integer indices.
10///
11/// Allows compressing addresses to small integers for more efficient on-chain encoding.
12/// Slot 0 = number of items, slots 1..N = address hashes.
13/// Sub-storage at key [] maps address_hash → 1-based index.
14pub struct AddressTable<'a, D> {
15    backing_storage: Storage<'a, D>,
16    by_address: Storage<'a, D>,
17    num_items: StorageBackedUint64,
18}
19
20pub fn initialize_address_table<D: Database>(_sto: &Storage<'_, D>) {
21    // no-op
22}
23
24pub fn open_address_table<D>(sto: Storage<'_, D>) -> AddressTable<'_, D> {
25    let num_items = StorageBackedUint64::new(sto.base_key(), 0);
26    let by_address = sto.open_sub_storage(&[]);
27    AddressTable {
28        backing_storage: sto,
29        by_address,
30        num_items,
31    }
32}
33
34impl<D> AddressTable<'_, D> {
35    /// Registers `addr` if not already present and returns its 0-based index
36    /// together with a flag indicating whether the address was already
37    /// registered.
38    pub fn register<B: StorageBackend>(
39        &self,
40        backend: &mut B,
41        addr: Address,
42    ) -> Result<(u64, bool), AddressTableError> {
43        let addr_hash = address_to_hash(addr);
44        let rev = self.by_address_get(backend, addr_hash)?;
45
46        if rev != B256::ZERO {
47            return Ok((U256::from_be_bytes(rev.0).to::<u64>() - 1, true));
48        }
49
50        let current = self.num_items.get(backend)?;
51        let new_num_items = current + 1;
52        self.num_items.set(backend, new_num_items)?;
53
54        self.backing_set_by_uint64(backend, new_num_items, addr_hash)?;
55        self.by_address_set(backend, addr_hash, uint_to_hash(new_num_items))?;
56
57        Ok((new_num_items - 1, false))
58    }
59
60    pub fn lookup<B: SystemStateBackend>(
61        &self,
62        backend: &mut B,
63        addr: Address,
64    ) -> Result<(u64, bool), AddressTableError> {
65        let addr_hash = address_to_hash(addr);
66        let res_hash = self.by_address_get(backend, addr_hash)?;
67        let res = U256::from_be_bytes(res_hash.0).to::<u64>();
68
69        if res == 0 {
70            Ok((0, false))
71        } else {
72            Ok((res - 1, true))
73        }
74    }
75
76    pub fn address_exists<B: SystemStateBackend>(
77        &self,
78        backend: &mut B,
79        addr: Address,
80    ) -> Result<bool, AddressTableError> {
81        let (_, exists) = self.lookup(backend, addr)?;
82        Ok(exists)
83    }
84
85    pub fn size<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, AddressTableError> {
86        Ok(self.num_items.get(backend)?)
87    }
88
89    pub fn lookup_index<B: SystemStateBackend>(
90        &self,
91        backend: &mut B,
92        index: u64,
93    ) -> Result<Option<Address>, AddressTableError> {
94        let items = self.num_items.get(backend)?;
95        if index >= items {
96            return Ok(None);
97        }
98        let value = self.backing_get_by_uint64(backend, index + 1)?;
99        let mut addr_bytes = [0u8; 20];
100        addr_bytes.copy_from_slice(&value.0[12..32]);
101        Ok(Some(Address::from(addr_bytes)))
102    }
103
104    /// Compress an address into an RLP-encoded index or raw address bytes.
105    pub fn compress<B: SystemStateBackend>(
106        &self,
107        backend: &mut B,
108        addr: Address,
109    ) -> Result<Vec<u8>, AddressTableError> {
110        let (index, exists) = self.lookup(backend, addr)?;
111        if exists {
112            let mut buf = Vec::new();
113            index.encode(&mut buf);
114            Ok(buf)
115        } else {
116            let mut buf = Vec::new();
117            addr.as_slice().encode(&mut buf);
118            Ok(buf)
119        }
120    }
121
122    /// Decompress RLP-encoded data back to an address. Returns the
123    /// resolved address, the number of bytes consumed, and whether the
124    /// encoding was a raw 20-byte address (vs. a table index).
125    pub fn decompress<B: SystemStateBackend>(
126        &self,
127        backend: &mut B,
128        buf: &[u8],
129    ) -> Result<(Address, u64, bool), AddressTableError> {
130        if buf.is_empty() {
131            return Err(AddressTableError::InvalidEncoding);
132        }
133        let mut header_cursor = buf;
134        let header =
135            Header::decode(&mut header_cursor).map_err(|_| AddressTableError::InvalidEncoding)?;
136        if header.list {
137            return Err(AddressTableError::InvalidEncoding);
138        }
139
140        if header.payload_length == 20 {
141            if header_cursor.len() < 20 {
142                return Err(AddressTableError::InvalidEncoding);
143            }
144            let mut addr_bytes = [0u8; 20];
145            addr_bytes.copy_from_slice(&header_cursor[..20]);
146            let consumed = buf.len() - header_cursor.len() + 20;
147            Ok((Address::from(addr_bytes), consumed as u64, true))
148        } else {
149            let mut cursor = buf;
150            let index = u64::decode(&mut cursor).map_err(|_| AddressTableError::InvalidEncoding)?;
151            let bytes_read = (buf.len() - cursor.len()) as u64;
152            let addr = self
153                .lookup_index(backend, index)?
154                .ok_or(AddressTableError::IndexOutOfRange(index))?;
155            Ok((addr, bytes_read, false))
156        }
157    }
158
159    fn by_address_get<B: SystemStateBackend>(
160        &self,
161        backend: &mut B,
162        key: B256,
163    ) -> Result<B256, AddressTableError> {
164        let slot = self.by_address.slot_for_key(key);
165        let value = backend
166            .sload_system(self.by_address.account(), slot)
167            .map_err(Into::into)?;
168        Ok(B256::from(value.to_be_bytes::<32>()))
169    }
170
171    fn by_address_set<B: StorageBackend>(
172        &self,
173        backend: &mut B,
174        key: B256,
175        value: B256,
176    ) -> Result<(), AddressTableError> {
177        let slot = self.by_address.slot_for_key(key);
178        backend
179            .sstore(
180                self.by_address.account(),
181                slot,
182                U256::from_be_bytes(value.0),
183            )
184            .map_err(Into::into)?;
185        Ok(())
186    }
187
188    fn backing_get_by_uint64<B: SystemStateBackend>(
189        &self,
190        backend: &mut B,
191        offset: u64,
192    ) -> Result<B256, AddressTableError> {
193        let slot = self.backing_storage.new_slot(offset);
194        let value = backend
195            .sload_system(self.backing_storage.account(), slot)
196            .map_err(Into::into)?;
197        Ok(B256::from(value.to_be_bytes::<32>()))
198    }
199
200    fn backing_set_by_uint64<B: StorageBackend>(
201        &self,
202        backend: &mut B,
203        offset: u64,
204        value: B256,
205    ) -> Result<(), AddressTableError> {
206        let slot = self.backing_storage.new_slot(offset);
207        backend
208            .sstore(
209                self.backing_storage.account(),
210                slot,
211                U256::from_be_bytes(value.0),
212            )
213            .map_err(Into::into)?;
214        Ok(())
215    }
216}
217
218fn address_to_hash(addr: Address) -> B256 {
219    let mut bytes = [0u8; 32];
220    bytes[12..32].copy_from_slice(addr.as_slice());
221    B256::from(bytes)
222}
223
224fn uint_to_hash(val: u64) -> B256 {
225    B256::from(U256::from(val))
226}
arbos/address_table/mod.rs

arbos/address_table/
mod.rs
