arb_storage/
queue.rs

1use alloy_primitives::B256;
2use revm::Database;
3
4use crate::{backed_types::StorageBackedUint64, storage::Storage};
5
6/// A FIFO queue backed by ArbOS storage.
7///
8/// Layout: offset 0 = next put position, offset 1 = next get position.
9/// Data stored at offsets 2+.
10pub struct Queue<D> {
11    pub storage: Storage<D>,
12    next_put: StorageBackedUint64<D>,
13    next_get: StorageBackedUint64<D>,
14}
15
16/// Initializes a queue by setting both offsets to 2 (data starts at offset 2).
17pub fn initialize_queue<D: Database>(storage: &Storage<D>) -> Result<(), ()> {
18    storage.set_uint64_by_uint64(0, 2)?;
19    storage.set_uint64_by_uint64(1, 2)?;
20    Ok(())
21}
22
23/// Opens an existing queue from storage.
24pub fn open_queue<D: Database>(storage: Storage<D>) -> Queue<D> {
25    let state = storage.state_ptr();
26    let base_key = storage.base_key();
27    Queue {
28        next_put: StorageBackedUint64::new(state, base_key, 0),
29        next_get: StorageBackedUint64::new(state, base_key, 1),
30        storage,
31    }
32}
33
34impl<D: Database> Queue<D> {
35    pub fn is_empty(&self) -> Result<bool, ()> {
36        let put = self.next_put.get()?;
37        let get = self.next_get.get()?;
38        Ok(put == get)
39    }
40
41    pub fn size(&self) -> Result<u64, ()> {
42        let put = self.next_put.get()?;
43        let get = self.next_get.get()?;
44        Ok(put.saturating_sub(get))
45    }
46
47    pub fn peek(&self) -> Result<Option<B256>, ()> {
48        if self.is_empty()? {
49            return Ok(None);
50        }
51        let get = self.next_get.get()?;
52        let val = self.storage.get_by_uint64(get)?;
53        Ok(Some(val))
54    }
55
56    pub fn get(&self) -> Result<Option<B256>, ()> {
57        if self.is_empty()? {
58            return Ok(None);
59        }
60        let get = self.next_get.get()?;
61        let val = self.storage.get_by_uint64(get)?;
62        self.storage.set_by_uint64(get, B256::ZERO)?;
63        self.next_get.set(get + 1)?;
64        Ok(Some(val))
65    }
66
67    pub fn put(&self, value: B256) -> Result<(), ()> {
68        let put = self.next_put.get()?;
69        self.storage.set_by_uint64(put, value)?;
70        self.next_put.set(put + 1)?;
71        Ok(())
72    }
73
74    /// Removes the last element from the back (most recently put).
75    pub fn shift(&self) -> Result<Option<B256>, ()> {
76        if self.is_empty()? {
77            return Ok(None);
78        }
79        let put = self.next_put.get()?;
80        let idx = put - 1;
81        let val = self.storage.get_by_uint64(idx)?;
82        self.storage.set_by_uint64(idx, B256::ZERO)?;
83        self.next_put.set(idx)?;
84        Ok(Some(val))
85    }
86
87    /// Iterates over all elements in order.
88    pub fn for_each<F>(&self, mut f: F) -> Result<(), ()>
89    where
90        F: FnMut(B256) -> Result<(), ()>,
91    {
92        let get = self.next_get.get()?;
93        let put = self.next_put.get()?;
94        for i in get..put {
95            let val = self.storage.get_by_uint64(i)?;
96            f(val)?;
97        }
98        Ok(())
99    }
100}