diff --git a/corelib/src/starknet/storage.cairo b/corelib/src/starknet/storage.cairo index 34ed12c175f..04e3424d143 100644 --- a/corelib/src/starknet/storage.cairo +++ b/corelib/src/starknet/storage.cairo @@ -139,7 +139,6 @@ //! address (the `sn_keccak` hash of the variable name) combined with the mapping keys or vector //! indices. //! See their respective module documentation for more details. - use core::hash::HashStateTrait; #[allow(unused_imports)] use core::pedersen::HashState; @@ -161,7 +160,10 @@ mod sub_pointers; pub use sub_pointers::{SubPointers, SubPointersForward, SubPointersMut, SubPointersMutForward}; mod vec; -use vec::{MutableVecIndexView, VecIndexView}; +use vec::{ + MutableVecIndexView, MutableVecIntoIterRange, PathableMutableVecIntoIterRange, + PathableVecIntoIterRange, VecIndexView, VecIntoIterRange, +}; pub use vec::{MutableVecTrait, Vec, VecTrait}; /// A pointer to an address in storage, can be used to read and write values, if the generic type @@ -539,3 +541,13 @@ trait MutableTrait { impl MutableImpl of MutableTrait> { type InnerType = T; } + +/// Trait for turning collection of values into an iterator over a specific range. +pub trait IntoIterRange { + type IntoIter; + impl Iterator: Iterator; + /// Creates an iterator over a range from a collection. + fn into_iter_range(self: T, range: core::ops::Range) -> Self::IntoIter; + /// Creates an iterator over the full range of a collection. + fn into_iter_full_range(self: T) -> Self::IntoIter; +} diff --git a/corelib/src/starknet/storage/vec.cairo b/corelib/src/starknet/storage/vec.cairo index 5e03fa885c4..2a4a3bc8b0d 100644 --- a/corelib/src/starknet/storage/vec.cairo +++ b/corelib/src/starknet/storage/vec.cairo @@ -77,11 +77,11 @@ //! arr //! } //! ``` - -use core::Option; +use core::ops::Range; use super::{ - Mutable, StorageAsPath, StorageAsPointer, StoragePath, StoragePathTrait, StoragePathUpdateTrait, - StoragePointer0Offset, StoragePointerReadAccess, StoragePointerWriteAccess, + IntoIterRange, Mutable, StorageAsPath, StorageAsPointer, StoragePath, StoragePathTrait, + StoragePathUpdateTrait, StoragePointer0Offset, StoragePointerReadAccess, + StoragePointerWriteAccess, }; /// Represents a dynamic array in contract storage. @@ -396,3 +396,105 @@ pub impl MutableVecIndexView< (*self).at(index) } } + +/// An iterator struct over a `Vec` in storage. +#[derive(Drop)] +pub struct VecIter> { + vec: T, + current_index: crate::ops::RangeIterator, +} + +impl VecIterator, +Drop, +Copy> of Iterator> { + type Item = StoragePath; + fn next(ref self: VecIter) -> Option { + self.vec.get(self.current_index.next()?) + } +} + +// Implement `IntoIterRange` for `StoragePath>` +pub impl VecIntoIterRange< + T, impl VecTraitImpl: VecTrait>>, +> of IntoIterRange>> { + type IntoIter = VecIter>, VecTraitImpl>; + #[inline] + fn into_iter_range(self: StoragePath>, range: Range) -> Self::IntoIter { + VecIter { current_index: range.into_iter(), vec: self } + } + #[inline] + fn into_iter_full_range(self: StoragePath>) -> Self::IntoIter { + VecIter { current_index: (0..core::num::traits::Bounded::MAX).into_iter(), vec: self } + } +} + +/// Implement `IntoIterRange` for any type that implements StorageAsPath into a storage path +/// that implements VecTrait. +pub impl PathableVecIntoIterRange< + T, + +Destruct, + impl PathImpl: StorageAsPath, + impl VecTraitImpl: VecTrait>, +> of IntoIterRange { + type IntoIter = VecIter, VecTraitImpl>; + #[inline] + fn into_iter_range(self: T, range: Range) -> Self::IntoIter { + VecIter { current_index: range.into_iter(), vec: self.as_path() } + } + #[inline] + fn into_iter_full_range(self: T) -> Self::IntoIter { + let vec = self.as_path(); + VecIter { current_index: (0..core::num::traits::Bounded::MAX).into_iter(), vec } + } +} + +/// An iterator struct over a `Mutable` in storage. +#[derive(Drop)] +struct MutableVecIter> { + vec: T, + current_index: crate::ops::RangeIterator, +} + +impl MutableVecIterator< + T, +Drop, +Copy, impl MutVecTraitImpl: MutableVecTrait, +> of Iterator> { + type Item = StoragePath>; + fn next(ref self: MutableVecIter) -> Option { + self.vec.get(self.current_index.next()?) + } +} + +// Implement `IntoIterRange` for `StoragePath>>` +pub impl MutableVecIntoIterRange< + T, impl MutVecTraitImpl: MutableVecTrait>>>, +> of IntoIterRange>>> { + type IntoIter = MutableVecIter>>, MutVecTraitImpl>; + #[inline] + fn into_iter_range(self: StoragePath>>, range: Range) -> Self::IntoIter { + MutableVecIter { current_index: range.into_iter(), vec: self } + } + #[inline] + fn into_iter_full_range(self: StoragePath>>) -> Self::IntoIter { + MutableVecIter { + current_index: (0..core::num::traits::Bounded::MAX).into_iter(), vec: self, + } + } +} + +/// Implement `IntoIterRange` for any type that implements StorageAsPath into a storage path +/// that implements MutableVecTrait. +pub impl PathableMutableVecIntoIterRange< + T, + +Destruct, + impl PathImpl: StorageAsPath, + impl MutVecTraitImpl: MutableVecTrait>, +> of IntoIterRange { + type IntoIter = MutableVecIter, MutVecTraitImpl>; + #[inline] + fn into_iter_range(self: T, range: Range) -> Self::IntoIter { + MutableVecIter { current_index: range.into_iter(), vec: self.as_path() } + } + #[inline] + fn into_iter_full_range(self: T) -> Self::IntoIter { + let vec = self.as_path(); + MutableVecIter { current_index: (0..core::num::traits::Bounded::MAX).into_iter(), vec } + } +} diff --git a/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo b/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo index e0b2427e084..e5f492a7b49 100644 --- a/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo +++ b/crates/cairo-lang-starknet/cairo_level_tests/collections_test.cairo @@ -1,8 +1,8 @@ use starknet::storage::{ - MutableVecTrait, StoragePathEntry, StoragePointerReadAccess, StoragePointerWriteAccess, + IntoIterRange, MutableVecTrait, StoragePathEntry, StoragePointerReadAccess, + StoragePointerWriteAccess, }; - #[starknet::contract] mod contract_with_map { use starknet::storage::Map; @@ -32,6 +32,51 @@ fn test_simple_member_write_to_map() { assert_eq!(vec_entry.read(), 1); } +#[test] +fn test_vec_iter() { + let mut mut_state = contract_with_vec::contract_state_for_testing(); + for i in 0..9_usize { + mut_state.simple.append().write(i); + }; + + let state = @contract_with_vec::contract_state_for_testing(); + let mut i = 0; + for entry in state.simple.into_iter_full_range() { + assert_eq!(entry.read(), i); + i += 1; + }; + assert_eq!(i, 9); + + let mut i = 2; + for entry in state.simple.into_iter_range(2..5) { + assert_eq!(entry.read(), i); + i += 1; + }; + assert_eq!(i, 5); +} + +#[test] +fn test_mut_vec_iter() { + let mut mut_state = contract_with_vec::contract_state_for_testing(); + for i in 0..9_usize { + mut_state.simple.append().write(i); + }; + + let mut i = 0; + for entry in mut_state.simple.into_iter_full_range() { + assert_eq!(entry.read(), i); + i += 1; + }; + assert_eq!(i, 9); + + let mut i = 2; + for entry in mut_state.simple.into_iter_range(2..5) { + assert_eq!(entry.read(), i); + i += 1; + }; + assert_eq!(i, 5); +} + #[test] fn test_simple_member_write_to_vec() { let mut map_contract_state = contract_with_map::contract_state_for_testing();