arb_storage/
vector.rs

1use alloy_primitives::B256;
2use arb_storage_errors::StorageError;
3
4use crate::{
5    backed_types::StorageBackedUint64,
6    backend::{StorageBackend, SystemStateBackend},
7    slot::derive_sub_key,
8    storage::Storage,
9};
10
11const LENGTH_OFFSET: u64 = 0;
12
13/// Vector of sub-storages backed by ArbOS storage.
14///
15/// Layout: offset 0 = length; sub-storages live at indices `0..length`.
16#[derive(Clone, Copy, Debug)]
17pub struct SubStorageVector {
18    pub base_key: B256,
19    length: StorageBackedUint64,
20}
21
22pub fn open_sub_storage_vector<D>(storage: Storage<'_, D>) -> SubStorageVector {
23    open_sub_storage_vector_at(storage.base_key())
24}
25
26pub(crate) fn open_sub_storage_vector_at(base_key: B256) -> SubStorageVector {
27    SubStorageVector {
28        base_key,
29        length: StorageBackedUint64::new(base_key, LENGTH_OFFSET),
30    }
31}
32
33impl SubStorageVector {
34    pub fn length<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, StorageError> {
35        self.length.get(backend)
36    }
37
38    /// Returns the base key for the sub-storage at `index`.
39    pub fn at(&self, index: u64) -> B256 {
40        derive_sub_key(self.base_key, &index.to_be_bytes())
41    }
42
43    pub fn push<B: StorageBackend>(&self, backend: &mut B) -> Result<B256, StorageError> {
44        let len = self.length.get(backend)?;
45        self.length.set(backend, len + 1)?;
46        Ok(self.at(len))
47    }
48
49    pub fn pop<B: StorageBackend>(&self, backend: &mut B) -> Result<Option<u64>, StorageError> {
50        let len = self.length.get(backend)?;
51        if len == 0 {
52            return Ok(None);
53        }
54        let new_len = len - 1;
55        self.length.set(backend, new_len)?;
56        Ok(Some(new_len))
57    }
58}