1use 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
32pub 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
46pub trait IntoBoxedIter<'a, T> {
48 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#[derive(Copy, Clone, Debug, PartialOrd, Eq, PartialEq, Default)]
65pub enum IterDirection {
66 #[default]
68 Forward,
69 Reverse,
71}
72
73#[impl_tools::autoimpl(for<T: trait> &T, &mut T, Box<T>)]
75pub trait IterableStore: KeyValueInspect {
76 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 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
124pub trait IterableTable<M>
126where
127 M: Mappable,
128{
129 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 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 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 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
228pub trait IteratorOverTable {
230 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 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 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 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 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 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 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 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
339pub 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
397pub 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 (_, _) => iterator(tree, prefix, start, direction)
417 .map(|(key, _)| key)
418 .into_boxed(),
419 }
420}