fuel_core_storage/
iter.rs

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

fuel_core_storage/
iter.rs
