Skip to main content

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