arb_storage/
queue.rs

1use alloy_primitives::{B256, U256};
2use arb_storage_errors::StorageError;
3
4use crate::{
5    backed_types::StorageBackedUint64,
6    backend::{StorageBackend, SystemStateBackend},
7    slot::storage_key_map,
8    state_ops::ARBOS_STATE_ADDRESS,
9    storage::Storage,
10};
11
12/// FIFO queue backed by ArbOS storage.
13///
14/// Layout: offset 0 = next put position, offset 1 = next get position;
15/// data lives at offsets 2+.
16#[derive(Clone, Copy, Debug)]
17pub struct Queue {
18    pub base_key: B256,
19    next_put: StorageBackedUint64,
20    next_get: StorageBackedUint64,
21}
22
23fn compute_slot(base_key: B256, offset: u64) -> U256 {
24    if base_key == B256::ZERO {
25        storage_key_map(&[], offset)
26    } else {
27        storage_key_map(base_key.as_slice(), offset)
28    }
29}
30
31pub fn initialize_queue<D: revm::Database>(storage: &Storage<'_, D>) -> Result<(), StorageError> {
32    storage.set_uint64_by_uint64(0, 2)?;
33    storage.set_uint64_by_uint64(1, 2)?;
34    Ok(())
35}
36
37pub fn open_queue<D>(storage: Storage<'_, D>) -> Queue {
38    open_queue_at(storage.base_key())
39}
40
41pub(crate) fn open_queue_at(base_key: B256) -> Queue {
42    Queue {
43        base_key,
44        next_put: StorageBackedUint64::new(base_key, 0),
45        next_get: StorageBackedUint64::new(base_key, 1),
46    }
47}
48
49impl Queue {
50    fn load_slot<B: SystemStateBackend>(
51        &self,
52        backend: &mut B,
53        offset: u64,
54    ) -> Result<B256, StorageError> {
55        let slot = compute_slot(self.base_key, offset);
56        let value = backend
57            .sload_system(ARBOS_STATE_ADDRESS, slot)
58            .map_err(Into::into)?;
59        Ok(B256::from(value.to_be_bytes::<32>()))
60    }
61
62    fn store_slot<B: StorageBackend>(
63        &self,
64        backend: &mut B,
65        offset: u64,
66        value: B256,
67    ) -> Result<(), StorageError> {
68        let slot = compute_slot(self.base_key, offset);
69        backend
70            .sstore(ARBOS_STATE_ADDRESS, slot, U256::from_be_bytes(value.0))
71            .map_err(Into::into)
72    }
73
74    pub fn is_empty<B: SystemStateBackend>(&self, backend: &mut B) -> Result<bool, StorageError> {
75        let put = self.next_put.get(backend)?;
76        let get = self.next_get.get(backend)?;
77        Ok(put == get)
78    }
79
80    pub fn size<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, StorageError> {
81        let put = self.next_put.get(backend)?;
82        let get = self.next_get.get(backend)?;
83        Ok(put.saturating_sub(get))
84    }
85
86    pub fn peek<B: SystemStateBackend>(
87        &self,
88        backend: &mut B,
89    ) -> Result<Option<B256>, StorageError> {
90        if self.is_empty(backend)? {
91            return Ok(None);
92        }
93        let get = self.next_get.get(backend)?;
94        let val = self.load_slot(backend, get)?;
95        Ok(Some(val))
96    }
97
98    pub fn get<B: StorageBackend>(&self, backend: &mut B) -> Result<Option<B256>, StorageError> {
99        if self.is_empty(backend)? {
100            return Ok(None);
101        }
102        let get = self.next_get.get(backend)?;
103        let val = self.load_slot(backend, get)?;
104        self.store_slot(backend, get, B256::ZERO)?;
105        self.next_get.set(backend, get + 1)?;
106        Ok(Some(val))
107    }
108
109    pub fn put<B: StorageBackend>(&self, backend: &mut B, value: B256) -> Result<(), StorageError> {
110        let put = self.next_put.get(backend)?;
111        self.store_slot(backend, put, value)?;
112        self.next_put.set(backend, put + 1)?;
113        Ok(())
114    }
115
116    /// Removes the last element from the back (most recently put).
117    pub fn shift<B: StorageBackend>(&self, backend: &mut B) -> Result<Option<B256>, StorageError> {
118        if self.is_empty(backend)? {
119            return Ok(None);
120        }
121        let put = self.next_put.get(backend)?;
122        let idx = put - 1;
123        let val = self.load_slot(backend, idx)?;
124        self.store_slot(backend, idx, B256::ZERO)?;
125        self.next_put.set(backend, idx)?;
126        Ok(Some(val))
127    }
128
129    pub fn for_each<F, E, B>(&self, backend: &mut B, mut f: F) -> Result<(), E>
130    where
131        F: FnMut(B256) -> Result<(), E>,
132        E: From<StorageError>,
133        B: SystemStateBackend,
134    {
135        let get = self.next_get.get(backend)?;
136        let put = self.next_put.get(backend)?;
137        for i in get..put {
138            let val = self.load_slot(backend, i)?;
139            f(val)?;
140        }
141        Ok(())
142    }
143}
arb_storage/queue.rs

arb_storage/
queue.rs
