arbos/address_table/
mod.rs

1use alloy_primitives::{Address, B256, U256};
2use alloy_rlp::{Decodable, Encodable};
3use revm::Database;
4
5use arb_storage::{Storage, StorageBackedUint64};
6
7/// A mapping between addresses and compact integer indices.
8///
9/// Allows compressing addresses to small integers for more efficient on-chain encoding.
10/// Slot 0 = number of items, slots 1..N = address hashes.
11/// Sub-storage at key [] maps address_hash → 1-based index.
12pub struct AddressTable<D> {
13    backing_storage: Storage<D>,
14    by_address: Storage<D>,
15    num_items: StorageBackedUint64<D>,
16}
17
18pub fn initialize_address_table<D: Database>(_sto: &Storage<D>) {
19    // no-op
20}
21
22pub fn open_address_table<D: Database>(sto: Storage<D>) -> AddressTable<D> {
23    let num_items = StorageBackedUint64::new(sto.state_ptr(), sto.base_key(), 0);
24    let by_address = sto.open_sub_storage(&[]);
25    AddressTable {
26        backing_storage: sto,
27        by_address,
28        num_items,
29    }
30}
31
32impl<D: Database> AddressTable<D> {
33    pub fn register(&self, addr: Address) -> Result<u64, ()> {
34        let addr_hash = address_to_hash(addr);
35        let rev = self.by_address.get(addr_hash)?;
36
37        if rev != B256::ZERO {
38            return Ok(U256::from_be_bytes(rev.0).to::<u64>() - 1);
39        }
40
41        // Address isn't in the table, add it
42        let current = self.num_items.get()?;
43        let new_num_items = current + 1;
44        self.num_items.set(new_num_items)?;
45
46        self.backing_storage
47            .set_by_uint64(new_num_items, addr_hash)?;
48        self.by_address
49            .set(addr_hash, uint_to_hash(new_num_items))?;
50
51        Ok(new_num_items - 1)
52    }
53
54    pub fn lookup(&self, addr: Address) -> Result<(u64, bool), ()> {
55        let addr_hash = address_to_hash(addr);
56        let res_hash = self.by_address.get(addr_hash)?;
57        let res = U256::from_be_bytes(res_hash.0).to::<u64>();
58
59        if res == 0 {
60            Ok((0, false))
61        } else {
62            Ok((res - 1, true))
63        }
64    }
65
66    pub fn address_exists(&self, addr: Address) -> Result<bool, ()> {
67        let (_, exists) = self.lookup(addr)?;
68        Ok(exists)
69    }
70
71    pub fn size(&self) -> Result<u64, ()> {
72        self.num_items.get()
73    }
74
75    pub fn lookup_index(&self, index: u64) -> Result<Option<Address>, ()> {
76        let items = self.num_items.get()?;
77        if index >= items {
78            return Ok(None);
79        }
80        let value = self.backing_storage.get_by_uint64(index + 1)?;
81        let mut addr_bytes = [0u8; 20];
82        addr_bytes.copy_from_slice(&value.0[12..32]);
83        Ok(Some(Address::from(addr_bytes)))
84    }
85
86    /// Compress an address into an RLP-encoded index or raw address bytes.
87    pub fn compress(&self, addr: Address) -> Result<Vec<u8>, ()> {
88        let (index, exists) = self.lookup(addr)?;
89        if exists {
90            let mut buf = Vec::new();
91            index.encode(&mut buf);
92            Ok(buf)
93        } else {
94            let mut buf = Vec::new();
95            addr.as_slice().encode(&mut buf);
96            Ok(buf)
97        }
98    }
99
100    /// Decompress RLP-encoded data back to an address.
101    /// Returns (address, number_of_bytes_read).
102    pub fn decompress(&self, buf: &[u8]) -> Result<(Address, u64), ()> {
103        // Try decoding as bytes first
104        let mut cursor = buf;
105        let input = <Vec<u8> as Decodable>::decode(&mut cursor).map_err(|_| ())?;
106        let bytes_read = (buf.len() - cursor.len()) as u64;
107
108        if input.len() == 20 {
109            let mut addr_bytes = [0u8; 20];
110            addr_bytes.copy_from_slice(&input);
111            Ok((Address::from(addr_bytes), bytes_read))
112        } else {
113            // Re-decode as u64 index
114            let mut cursor = buf;
115            let index = u64::decode(&mut cursor).map_err(|_| ())?;
116            let bytes_read = (buf.len() - cursor.len()) as u64;
117            let addr = self.lookup_index(index)?.ok_or(())?;
118            Ok((addr, bytes_read))
119        }
120    }
121}
122
123fn address_to_hash(addr: Address) -> B256 {
124    let mut bytes = [0u8; 32];
125    bytes[12..32].copy_from_slice(addr.as_slice());
126    B256::from(bytes)
127}
128
129fn uint_to_hash(val: u64) -> B256 {
130    B256::from(U256::from(val))
131}
arbos/address_table/mod.rs

arbos/address_table/
mod.rs
