diff --git a/stylus-sdk/src/storage/vec.rs b/stylus-sdk/src/storage/vec.rs index 2b7216b5..4097d89a 100644 --- a/stylus-sdk/src/storage/vec.rs +++ b/stylus-sdk/src/storage/vec.rs @@ -294,10 +294,68 @@ impl<'a, S: SimpleStorageType<'a>> Extend> for StorageVec { } } +/// An iterator over the values of a `StorageVec`. +/// +/// Create the iterator by calling: [`StorageVec.iter()`] +pub struct Iter<'a, S> { + idx: usize, + base: U256, + len: usize, + __stylus_host: VM, + marker: PhantomData<&'a S>, +} + +impl<'a, S: SimpleStorageType<'a>> StorageVec { + /// Create an iterator + pub fn iter(&self) -> Iter<'a, S> { + Iter { + idx: 0, + base: *self.base(), + len: self.len(), + __stylus_host: self.__stylus_host.clone(), + marker: PhantomData, + } + } +} + +impl<'a, S: SimpleStorageType<'a>> IntoIterator for &'a StorageVec { + type Item = S::Wraps<'a>; + type IntoIter = Iter<'a, S>; + + /// Convert to an iterator + fn into_iter(self) -> Iter<'a, S> { + self.iter() + } +} + +impl<'a, S: SimpleStorageType<'a>> Iterator for Iter<'a, S> { + type Item = S::Wraps<'a>; + + fn next(&mut self) -> Option { + if self.idx >= self.len { + return None; + } + // calculate where to find the next value + let width = S::SLOT_BYTES; + let words = S::REQUIRED_SLOTS.max(1); + let density = 32 / S::SLOT_BYTES; + let slot = self.base + U256::from(words * self.idx / density); + let offset = 32 - (width * (1 + self.idx % density)) as u8; + + let store = unsafe { S::new(slot, offset, self.__stylus_host.clone()) }; + // increment and track the state of the index + self.idx += 1; + + Some(store.load()) + } +} + #[cfg(test)] mod test { use stylus_test::vm::TestVM; + use crate::storage::StorageU256; + #[test] fn test_storage_vec() { use super::super::StorageBool; @@ -312,4 +370,50 @@ mod test { assert!(!vec.get(1).unwrap()); assert!(vec.get(2).unwrap()); } + + #[test] + fn test_storage_vec_iter() { + use super::*; + use alloy_primitives::U256; + + let host = TestVM::new(); + let mut vec: StorageVec = StorageVec::from(&host); + vec.push(U256::from(1)); + vec.push(U256::from(2)); + vec.push(U256::from(3)); + + let mut iter = vec.iter(); + assert_eq!(U256::from(1), iter.next().unwrap()); + assert_eq!(U256::from(2), iter.next().unwrap()); + assert_eq!(U256::from(3), iter.next().unwrap()); + assert_eq!(None, iter.next()); + } + + #[test] + fn test_empty_storage_vec_iter() { + use super::*; + + let host = TestVM::new(); + let vec: StorageVec = StorageVec::from(&host); + let mut iter = vec.iter(); + assert_eq!(None, iter.next()); + } + + #[test] + fn test_storage_vec_into_iter() { + use super::*; + use alloy_primitives::U256; + + let host = TestVM::new(); + let mut vec: StorageVec = StorageVec::from(&host); + vec.push(U256::from(1)); + vec.push(U256::from(2)); + vec.push(U256::from(3)); + + let v = vec.into_iter().collect::>(); + assert_eq!(3, v.len()); + assert_eq!(U256::from(1), v[0]); + assert_eq!(U256::from(2), v[1]); + assert_eq!(U256::from(3), v[2]); + } }