fuel_core_storage/
iter.rs

1//! The module defines primitives that allow iterating of the storage.
2
3use crate::{
4    blueprint::{
5        BlueprintCodec,
6        BlueprintInspect,
7    },
8    codec::{
9        Decode,
10        Encode,
11        Encoder,
12    },
13    kv_store::{
14        KVItem,
15        KeyItem,
16        KeyValueInspect,
17    },
18    structured_storage::TableWithBlueprint,
19    transactional::ReferenceBytesKey,
20};
21use fuel_vm_private::fuel_storage::Mappable;
22
23#[cfg(feature = "alloc")]
24use alloc::{
25    boxed::Box,
26    collections::BTreeMap,
27    vec::Vec,
28};
29
30pub mod changes_iterator;
31
32// TODO: BoxedIter to be used until RPITIT lands in stable rust.
33/// A boxed variant of the iterator that can be used as a return type of the traits.
34pub struct BoxedIter<'a, T> {
35    iter: Box<dyn Iterator<Item = T> + 'a + Send>,
36}
37
38impl<T> Iterator for BoxedIter<'_, T> {
39    type Item = T;
40
41    fn next(&mut self) -> Option<Self::Item> {
42        self.iter.next()
43    }
44}
45
46/// The traits simplifies conversion into `BoxedIter`.
47pub trait IntoBoxedIter<'a, T> {
48    /// Converts `Self` iterator into `BoxedIter`.
49    fn into_boxed(self) -> BoxedIter<'a, T>;
50}
51
52impl<'a, T, I> IntoBoxedIter<'a, T> for I
53where
54    I: Iterator<Item = T> + 'a + Send,
55{
56    fn into_boxed(self) -> BoxedIter<'a, T> {
57        BoxedIter {
58            iter: Box::new(self),
59        }
60    }
61}
62
63/// A enum for iterating across the database
64#[derive(Copy, Clone, Debug, PartialOrd, Eq, PartialEq)]
65pub enum IterDirection {
66    /// Iterate forward
67    Forward,
68    /// Iterate backward
69    Reverse,
70}
71
72impl Default for IterDirection {
73    fn default() -> Self {
74        Self::Forward
75    }
76}
77
78/// A trait for iterating over the storage of [`KeyValueInspect`].
79#[impl_tools::autoimpl(for<T: trait> &T, &mut T, Box<T>)]
80pub trait IterableStore: KeyValueInspect {
81    /// Returns an iterator over the values in the storage.
82    fn iter_store(
83        &self,
84        column: Self::Column,
85        prefix: Option<&[u8]>,
86        start: Option<&[u8]>,
87        direction: IterDirection,
88    ) -> BoxedIter<'_, KVItem>;
89
90    /// Returns an iterator over keys in the storage.
91    fn iter_store_keys(
92        &self,
93        column: Self::Column,
94        prefix: Option<&[u8]>,
95        start: Option<&[u8]>,
96        direction: IterDirection,
97    ) -> BoxedIter<'_, KeyItem>;
98}
99
100#[cfg(feature = "std")]
101impl<T> IterableStore for std::sync::Arc<T>
102where
103    T: IterableStore,
104{
105    fn iter_store(
106        &self,
107        column: Self::Column,
108        prefix: Option<&[u8]>,
109        start: Option<&[u8]>,
110        direction: IterDirection,
111    ) -> BoxedIter<'_, KVItem> {
112        use core::ops::Deref;
113        self.deref().iter_store(column, prefix, start, direction)
114    }
115
116    fn iter_store_keys(
117        &self,
118        column: Self::Column,
119        prefix: Option<&[u8]>,
120        start: Option<&[u8]>,
121        direction: IterDirection,
122    ) -> BoxedIter<'_, KeyItem> {
123        use core::ops::Deref;
124        self.deref()
125            .iter_store_keys(column, prefix, start, direction)
126    }
127}
128
129/// A trait for iterating over the `Mappable` table.
130pub trait IterableTable<M>
131where
132    M: Mappable,
133{
134    /// Returns an iterator over the all keys in the table with a prefix after a specific start key.
135    fn iter_table_keys<P>(
136        &self,
137        prefix: Option<P>,
138        start: Option<&M::Key>,
139        direction: Option<IterDirection>,
140    ) -> BoxedIter<'_, super::Result<M::OwnedKey>>
141    where
142        P: AsRef<[u8]>;
143
144    /// Returns an iterator over the all entries in the table with a prefix after a specific start key.
145    fn iter_table<P>(
146        &self,
147        prefix: Option<P>,
148        start: Option<&M::Key>,
149        direction: Option<IterDirection>,
150    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
151    where
152        P: AsRef<[u8]>;
153}
154
155impl<Column, M, S> IterableTable<M> for S
156where
157    M: TableWithBlueprint<Column = Column>,
158    M::Blueprint: BlueprintInspect<M, S>,
159    S: IterableStore<Column = Column>,
160{
161    fn iter_table_keys<P>(
162        &self,
163        prefix: Option<P>,
164        start: Option<&M::Key>,
165        direction: Option<IterDirection>,
166    ) -> BoxedIter<'_, crate::Result<M::OwnedKey>>
167    where
168        P: AsRef<[u8]>,
169    {
170        #[allow(clippy::redundant_closure)]
171        // false positive: https://github.com/rust-lang/rust-clippy/issues/14215
172        let encoder = start
173            .map(|start| <M::Blueprint as BlueprintCodec<M>>::KeyCodec::encode(start));
174
175        let start = encoder.as_ref().map(|encoder| encoder.as_bytes());
176
177        IterableStore::iter_store_keys(
178            self,
179            M::column(),
180            prefix.as_ref().map(|p| p.as_ref()),
181            start.as_ref().map(|cow| cow.as_ref()),
182            direction.unwrap_or_default(),
183        )
184        .map(|res| {
185            res.and_then(|key| {
186                let key =
187                    <M::Blueprint as BlueprintCodec<M>>::KeyCodec::decode(key.as_slice())
188                        .map_err(|e| crate::Error::Codec(anyhow::anyhow!(e)))?;
189                Ok(key)
190            })
191        })
192        .into_boxed()
193    }
194
195    fn iter_table<P>(
196        &self,
197        prefix: Option<P>,
198        start: Option<&M::Key>,
199        direction: Option<IterDirection>,
200    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
201    where
202        P: AsRef<[u8]>,
203    {
204        #[allow(clippy::redundant_closure)]
205        // false positive: https://github.com/rust-lang/rust-clippy/issues/14215
206        let encoder = start
207            .map(|start| <M::Blueprint as BlueprintCodec<M>>::KeyCodec::encode(start));
208
209        let start = encoder.as_ref().map(|encoder| encoder.as_bytes());
210
211        IterableStore::iter_store(
212            self,
213            M::column(),
214            prefix.as_ref().map(|p| p.as_ref()),
215            start.as_ref().map(|cow| cow.as_ref()),
216            direction.unwrap_or_default(),
217        )
218        .map(|val| {
219            val.and_then(|(key, value)| {
220                let key =
221                    <M::Blueprint as BlueprintCodec<M>>::KeyCodec::decode(key.as_slice())
222                        .map_err(|e| crate::Error::Codec(anyhow::anyhow!(e)))?;
223                let value =
224                    <M::Blueprint as BlueprintCodec<M>>::ValueCodec::decode(&value)
225                        .map_err(|e| crate::Error::Codec(anyhow::anyhow!(e)))?;
226                Ok((key, value))
227            })
228        })
229        .into_boxed()
230    }
231}
232
233/// A helper trait to provide a user-friendly API over table iteration.
234pub trait IteratorOverTable {
235    /// Returns an iterator over the all keys in the table.
236    fn iter_all_keys<M>(
237        &self,
238        direction: Option<IterDirection>,
239    ) -> BoxedIter<'_, super::Result<M::OwnedKey>>
240    where
241        M: Mappable,
242        Self: IterableTable<M>,
243    {
244        self.iter_all_filtered_keys::<M, [u8; 0]>(None, None, direction)
245    }
246
247    /// Returns an iterator over the all keys in the table with the specified prefix.
248    fn iter_all_by_prefix_keys<M, P>(
249        &self,
250        prefix: Option<P>,
251    ) -> BoxedIter<'_, super::Result<M::OwnedKey>>
252    where
253        M: Mappable,
254        P: AsRef<[u8]>,
255        Self: IterableTable<M>,
256    {
257        self.iter_all_filtered_keys::<M, P>(prefix, None, None)
258    }
259
260    /// Returns an iterator over the all keys in the table after a specific start key.
261    fn iter_all_by_start_keys<M>(
262        &self,
263        start: Option<&M::Key>,
264        direction: Option<IterDirection>,
265    ) -> BoxedIter<'_, super::Result<M::OwnedKey>>
266    where
267        M: Mappable,
268        Self: IterableTable<M>,
269    {
270        self.iter_all_filtered_keys::<M, [u8; 0]>(None, start, direction)
271    }
272
273    /// Returns an iterator over the all keys in the table with a prefix after a specific start key.
274    fn iter_all_filtered_keys<M, P>(
275        &self,
276        prefix: Option<P>,
277        start: Option<&M::Key>,
278        direction: Option<IterDirection>,
279    ) -> BoxedIter<'_, super::Result<M::OwnedKey>>
280    where
281        M: Mappable,
282        P: AsRef<[u8]>,
283        Self: IterableTable<M>,
284    {
285        self.iter_table_keys(prefix, start, direction)
286    }
287
288    /// Returns an iterator over the all entries in the table.
289    fn iter_all<M>(
290        &self,
291        direction: Option<IterDirection>,
292    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
293    where
294        M: Mappable,
295        Self: IterableTable<M>,
296    {
297        self.iter_all_filtered::<M, [u8; 0]>(None, None, direction)
298    }
299
300    /// Returns an iterator over the all entries in the table with the specified prefix.
301    fn iter_all_by_prefix<M, P>(
302        &self,
303        prefix: Option<P>,
304    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
305    where
306        M: Mappable,
307        P: AsRef<[u8]>,
308        Self: IterableTable<M>,
309    {
310        self.iter_all_filtered::<M, P>(prefix, None, None)
311    }
312
313    /// Returns an iterator over the all entries in the table after a specific start key.
314    fn iter_all_by_start<M>(
315        &self,
316        start: Option<&M::Key>,
317        direction: Option<IterDirection>,
318    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
319    where
320        M: Mappable,
321        Self: IterableTable<M>,
322    {
323        self.iter_all_filtered::<M, [u8; 0]>(None, start, direction)
324    }
325
326    /// Returns an iterator over the all entries in the table with a prefix after a specific start key.
327    fn iter_all_filtered<M, P>(
328        &self,
329        prefix: Option<P>,
330        start: Option<&M::Key>,
331        direction: Option<IterDirection>,
332    ) -> BoxedIter<'_, super::Result<(M::OwnedKey, M::OwnedValue)>>
333    where
334        M: Mappable,
335        P: AsRef<[u8]>,
336        Self: IterableTable<M>,
337    {
338        self.iter_table(prefix, start, direction)
339    }
340}
341
342impl<S> IteratorOverTable for S {}
343
344/// Returns an iterator over the values in the `BTreeMap`.
345pub fn iterator<'a, V>(
346    tree: &'a BTreeMap<ReferenceBytesKey, V>,
347    prefix: Option<&[u8]>,
348    start: Option<&[u8]>,
349    direction: IterDirection,
350) -> impl Iterator<Item = (&'a ReferenceBytesKey, &'a V)> + 'a + use<'a, V>
351where
352    V: Send + Sync,
353{
354    match (prefix, start) {
355        (None, None) => {
356            if direction == IterDirection::Forward {
357                tree.iter().into_boxed()
358            } else {
359                tree.iter().rev().into_boxed()
360            }
361        }
362        (Some(prefix), None) => {
363            let prefix = prefix.to_vec();
364            if direction == IterDirection::Forward {
365                tree.range(prefix.clone()..)
366                    .take_while(move |(key, _)| key.starts_with(prefix.as_slice()))
367                    .into_boxed()
368            } else {
369                let mut vec: Vec<_> = tree
370                    .range(prefix.clone()..)
371                    .into_boxed()
372                    .take_while(|(key, _)| key.starts_with(prefix.as_slice()))
373                    .collect();
374
375                vec.reverse();
376                vec.into_iter().into_boxed()
377            }
378        }
379        (None, Some(start)) => {
380            if direction == IterDirection::Forward {
381                tree.range(start.to_vec()..).into_boxed()
382            } else {
383                tree.range(..=start.to_vec()).rev().into_boxed()
384            }
385        }
386        (Some(prefix), Some(start)) => {
387            let prefix = prefix.to_vec();
388            if direction == IterDirection::Forward {
389                tree.range(start.to_vec()..)
390                    .take_while(move |(key, _)| key.starts_with(prefix.as_slice()))
391                    .into_boxed()
392            } else {
393                tree.range(..=start.to_vec())
394                    .rev()
395                    .take_while(move |(key, _)| key.starts_with(prefix.as_slice()))
396                    .into_boxed()
397            }
398        }
399    }
400}
401
402/// Returns an iterator over the keys in the `BTreeMap`.
403pub fn keys_iterator<'a, V>(
404    tree: &'a BTreeMap<ReferenceBytesKey, V>,
405    prefix: Option<&[u8]>,
406    start: Option<&[u8]>,
407    direction: IterDirection,
408) -> impl Iterator<Item = &'a ReferenceBytesKey> + 'a + use<'a, V>
409where
410    V: Send + Sync,
411{
412    match (prefix, start) {
413        (None, None) => {
414            if direction == IterDirection::Forward {
415                tree.keys().into_boxed()
416            } else {
417                tree.keys().rev().into_boxed()
418            }
419        }
420        // all the other cases require using a range, so we can't use the keys() method
421        (_, _) => iterator(tree, prefix, start, direction)
422            .map(|(key, _)| key)
423            .into_boxed(),
424    }
425}