arb_storage/
bytes_storage.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
10/// Variable-length byte storage.
11///
12/// Layout: offset 0 holds the byte length; offsets 1..N hold 32-byte chunks
13/// with the trailing partial chunk right-aligned in its slot.
14#[derive(Clone, Copy, Debug)]
15pub struct StorageBackedBytes {
16    pub base_key: B256,
17    pub account: Address,
18}
19
20impl StorageBackedBytes {
21    pub fn new(base_key: B256) -> Self {
22        Self {
23            base_key,
24            account: ARBOS_STATE_ADDRESS,
25        }
26    }
27
28    pub fn new_with_account(base_key: B256, account: Address) -> Self {
29        Self { base_key, account }
30    }
31
32    fn slot(&self, offset: u64) -> U256 {
33        let key: &[u8] = if self.base_key == B256::ZERO {
34            &[]
35        } else {
36            self.base_key.as_slice()
37        };
38        storage_key_map(key, offset)
39    }
40
41    fn load_u64<B: SystemStateBackend>(
42        &self,
43        backend: &mut B,
44        offset: u64,
45    ) -> Result<u64, StorageError> {
46        let value = backend
47            .sload_system(self.account, self.slot(offset))
48            .map_err(Into::into)?;
49        Ok(value.try_into().unwrap_or(0))
50    }
51
52    fn store_u64<B: StorageBackend>(
53        &self,
54        backend: &mut B,
55        offset: u64,
56        value: u64,
57    ) -> Result<(), StorageError> {
58        backend
59            .sstore(self.account, self.slot(offset), U256::from(value))
60            .map_err(Into::into)
61    }
62
63    fn load_word<B: SystemStateBackend>(
64        &self,
65        backend: &mut B,
66        offset: u64,
67    ) -> Result<[u8; 32], StorageError> {
68        let value = backend
69            .sload_system(self.account, self.slot(offset))
70            .map_err(Into::into)?;
71        Ok(value.to_be_bytes::<32>())
72    }
73
74    fn store_word<B: StorageBackend>(
75        &self,
76        backend: &mut B,
77        offset: u64,
78        word: [u8; 32],
79    ) -> Result<(), StorageError> {
80        backend
81            .sstore(self.account, self.slot(offset), U256::from_be_bytes(word))
82            .map_err(Into::into)
83    }
84
85    pub fn get<B: SystemStateBackend>(&self, backend: &mut B) -> Result<Vec<u8>, StorageError> {
86        let mut bytes_left = self.load_u64(backend, 0)? as usize;
87        if bytes_left == 0 {
88            return Ok(Vec::new());
89        }
90        let mut ret = Vec::with_capacity(bytes_left);
91        let mut offset = 1u64;
92        while bytes_left >= 32 {
93            let word = self.load_word(backend, offset)?;
94            ret.extend_from_slice(&word);
95            bytes_left -= 32;
96            offset += 1;
97        }
98        if bytes_left > 0 {
99            let word = self.load_word(backend, offset)?;
100            ret.extend_from_slice(&word[32 - bytes_left..]);
101        }
102        Ok(ret)
103    }
104
105    pub fn set<B: StorageBackend>(&self, backend: &mut B, b: &[u8]) -> Result<(), StorageError> {
106        self.clear(backend)?;
107        self.store_u64(backend, 0, b.len() as u64)?;
108        let mut remaining = b;
109        let mut offset = 1u64;
110        while remaining.len() >= 32 {
111            let mut word = [0u8; 32];
112            word.copy_from_slice(&remaining[..32]);
113            self.store_word(backend, offset, word)?;
114            remaining = &remaining[32..];
115            offset += 1;
116        }
117        if !remaining.is_empty() {
118            let mut word = [0u8; 32];
119            word[32 - remaining.len()..].copy_from_slice(remaining);
120            self.store_word(backend, offset, word)?;
121        }
122        Ok(())
123    }
124
125    pub fn clear<B: StorageBackend>(&self, backend: &mut B) -> Result<(), StorageError> {
126        let bytes_left = self.load_u64(backend, 0)?;
127        let mut offset = 1u64;
128        let mut remaining = bytes_left;
129        while remaining > 0 {
130            self.store_word(backend, offset, [0u8; 32])?;
131            offset += 1;
132            remaining = remaining.saturating_sub(32);
133        }
134        self.store_u64(backend, 0, 0)
135    }
136
137    pub fn size<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, StorageError> {
138        self.load_u64(backend, 0)
139    }
140}
arb_storage/bytes_storage.rs

arb_storage/
bytes_storage.rs
