From f1fd67d7033187bbfff07ed207f227d103f1c1c7 Mon Sep 17 00:00:00 2001 From: Dave Bryson Date: Fri, 4 Jul 2025 06:55:59 -0500 Subject: [PATCH 1/2] initial support for iter --- stylus-sdk/src/storage/vec.rs | 74 ++++++++++++++++++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) diff --git a/stylus-sdk/src/storage/vec.rs b/stylus-sdk/src/storage/vec.rs index 2b7216b5..d0b98911 100644 --- a/stylus-sdk/src/storage/vec.rs +++ b/stylus-sdk/src/storage/vec.rs @@ -7,7 +7,7 @@ use super::{ use crate::{crypto, host::VM}; use alloy_primitives::U256; use cfg_if::cfg_if; -use core::{cell::OnceCell, marker::PhantomData}; +use core::{cell::OnceCell, marker::PhantomData, ops::Add}; use stylus_core::HostAccess; /// Accessor for a storage-backed vector. @@ -294,10 +294,54 @@ impl<'a, S: SimpleStorageType<'a>> Extend> for StorageVec { } } +pub struct Iter<'a, S> { + idx: usize, + base: U256, + len: usize, + __stylus_host: VM, + marker: PhantomData<&'a S>, +} + +impl<'a, S: SimpleStorageType<'a>> StorageVec { + /// + /// + 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>> Iterator for Iter<'a, S> { + type Item = S::Wraps<'a>; + + fn next(&mut self) -> Option { + if self.idx >= self.len { + return None; + } + 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()) }; + 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 +356,32 @@ 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()); + } } From 1b5be0e148428d90cc9c966a978f29889fc0daf4 Mon Sep 17 00:00:00 2001 From: Dave Bryson Date: Fri, 4 Jul 2025 08:04:34 -0500 Subject: [PATCH 2/2] add support for iter/into_iter for vec --- stylus-sdk/src/storage/vec.rs | 40 +++++++++++++++++++++++++++++++---- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/stylus-sdk/src/storage/vec.rs b/stylus-sdk/src/storage/vec.rs index d0b98911..4097d89a 100644 --- a/stylus-sdk/src/storage/vec.rs +++ b/stylus-sdk/src/storage/vec.rs @@ -7,7 +7,7 @@ use super::{ use crate::{crypto, host::VM}; use alloy_primitives::U256; use cfg_if::cfg_if; -use core::{cell::OnceCell, marker::PhantomData, ops::Add}; +use core::{cell::OnceCell, marker::PhantomData}; use stylus_core::HostAccess; /// Accessor for a storage-backed vector. @@ -294,6 +294,9 @@ 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, @@ -303,8 +306,7 @@ pub struct Iter<'a, S> { } impl<'a, S: SimpleStorageType<'a>> StorageVec { - /// - /// + /// Create an iterator pub fn iter(&self) -> Iter<'a, S> { Iter { idx: 0, @@ -316,6 +318,16 @@ impl<'a, S: SimpleStorageType<'a>> StorageVec { } } +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>; @@ -323,15 +335,17 @@ impl<'a, S: SimpleStorageType<'a>> Iterator for Iter<'a, S> { 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()) } } @@ -384,4 +398,22 @@ mod test { 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]); + } }