feat(soavec): add SoAVec::iter and SoAVec::iter_mut

Author: omfj

soavec/src/iter.rs

@@ -0,0 +1,130 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+use std::{marker::PhantomData, ptr::NonNull};
+
+use crate::{
+    SoAVec,
+    soable::{SoATuple, SoAble},
+};
+
+impl<'a, T: SoAble> IntoIterator for &'a SoAVec<T> {
+    type Item = T::Ref<'a>;
+    type IntoIter = SoAIter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+impl<'a, T: SoAble> IntoIterator for &'a mut SoAVec<T> {
+    type Item = T::Mut<'a>;
+    type IntoIter = SoAIterMut<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter_mut()
+    }
+}
+
+/// An iterator over the elements of a `SoAVec`.
+///
+/// This struct is created by the [`iter`] method on [`SoAVec`].
+///
+/// [`iter`]: SoAVec::iter
+pub struct SoAIter<'a, T: SoAble> {
+    ptr: NonNull<u8>,
+    capacity: u32,
+    index: u32,
+    end: u32,
+    _marker: PhantomData<&'a T>,
+}
+
+impl<'a, T: SoAble> SoAIter<'a, T> {
+    pub(crate) fn new(ptr: NonNull<u8>, capacity: u32, end: u32) -> SoAIter<'a, T> {
+        SoAIter {
+            ptr,
+            capacity,
+            index: 0,
+            end,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<'a, T: SoAble> Iterator for SoAIter<'a, T> {
+    type Item = T::Ref<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.end {
+            return None;
+        }
+        let ptrs = unsafe { T::TupleRepr::get_pointers(self.ptr, self.index, self.capacity) };
+        self.index += 1;
+        Some(T::as_ref(PhantomData, ptrs))
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+}
+
+impl<'a, T: SoAble> ExactSizeIterator for SoAIter<'a, T> {
+    #[inline]
+    fn len(&self) -> usize {
+        (self.end - self.index) as usize
+    }
+}
+
+/// A mutable iterator over the elements of a `SoAVec`.
+///
+/// This struct is created by the [`iter_mut`] method on [`SoAVec`].
+///
+/// [`iter_mut`]: SoAVec::iter_mut
+pub struct SoAIterMut<'a, T: SoAble> {
+    ptr: NonNull<u8>,
+    capacity: u32,
+    index: u32,
+    end: u32,
+    _marker: PhantomData<&'a mut T>,
+}
+
+impl<'a, T: SoAble> SoAIterMut<'a, T> {
+    pub(crate) fn new(ptr: NonNull<u8>, capacity: u32, end: u32) -> SoAIterMut<'a, T> {
+        SoAIterMut {
+            ptr,
+            capacity,
+            index: 0,
+            end,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<'a, T: SoAble> Iterator for SoAIterMut<'a, T> {
+    type Item = T::Mut<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.end {
+            return None;
+        }
+        let ptrs = unsafe { T::TupleRepr::get_pointers(self.ptr, self.index, self.capacity) };
+        self.index += 1;
+        Some(T::as_mut(PhantomData, ptrs))
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+}
+
+impl<'a, T: SoAble> ExactSizeIterator for SoAIterMut<'a, T> {
+    #[inline]
+    fn len(&self) -> usize {
+        (self.end - self.index) as usize
+    }
+}

soavec/src/lib.rs

@@ -58,6 +58,7 @@
 //! [`as_slice`]: SoAVec::as_slice
 //! [`as_mut_slice`]: SoAVec::as_mut_slice
 
+mod iter;
 mod macros;
 mod raw_vec;
 mod raw_vec_inner;
@@ -66,6 +67,7 @@ mod soable;
 use core::marker::PhantomData;
 use std::ptr::NonNull;
 
+use iter::{SoAIter, SoAIterMut};
 use raw_vec::RawSoAVec;
 use raw_vec_inner::AllocError;
 pub use soable::{SoATuple, SoAble};
@@ -1100,6 +1102,52 @@ impl<T: SoAble> SoAVec<T> {
             Ok(value)
         }
     }
+
+    /// Returns an iterator over the soavec.
+    ///
+    /// The iterator yields all items from start to end.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use soavec::soavec;
+    ///
+    /// let vec = soavec![(1, 2), (3, 4), (5, 6)].unwrap();
+    /// let mut iterator = vec.iter();
+    ///
+    /// assert_eq!(iterator.next(), Some((&1, &2)));
+    /// assert_eq!(iterator.next(), Some((&3, &4)));
+    /// assert_eq!(iterator.next(), Some((&5, &6)));
+    /// assert_eq!(iterator.next(), None);
+    /// ```
+    #[inline]
+    pub fn iter(&self) -> SoAIter<'_, T> {
+        SoAIter::new(self.buf.as_ptr(), self.capacity(), self.len())
+    }
+
+    /// Returns an iterator that allows modifying each value.
+    ///
+    /// The iterator yields all items from start to end.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use soavec::soavec;
+    ///
+    /// let mut vec = soavec![(1, 2), (3, 4), (5, 6)].unwrap();
+    /// for (a, b) in vec.iter_mut() {
+    ///     *a += 10;
+    /// }
+    /// assert_eq!(vec.get(0), Some((&11, &2)));
+    /// assert_eq!(vec.get(1), Some((&13, &4)));
+    /// assert_eq!(vec.get(2), Some((&15, &6)));
+    /// ```
+    #[inline]
+    pub fn iter_mut(&mut self) -> SoAIterMut<'_, T> {
+        let len = self.len();
+        let capacity = self.capacity();
+        SoAIterMut::new(self.buf.as_mut_ptr(), capacity, len)
+    }
 }
 
 impl<T: SoAble> Drop for SoAVec<T> {
@@ -2084,4 +2132,37 @@ mod tests {
         assert_eq!(**third.a, 168);
         assert_eq!(third.b, &[6]);
     }
+
+    #[test]
+    fn test_iter() {
+        let mut vec = SoAVec::<(u32, u32)>::new();
+
+        vec.push((1, 10)).unwrap();
+        vec.push((2, 20)).unwrap();
+        vec.push((3, 30)).unwrap();
+
+        let mut iter = vec.iter();
+        assert_eq!(iter.next(), Some((&1, &10)));
+        assert_eq!(iter.next(), Some((&2, &20)));
+        assert_eq!(iter.next(), Some((&3, &30)));
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_iter_mut() {
+        let mut vec = SoAVec::<(u32, u32)>::new();
+
+        vec.push((1, 10)).unwrap();
+        vec.push((2, 20)).unwrap();
+        vec.push((3, 30)).unwrap();
+
+        for (a, b) in vec.iter_mut() {
+            *a += 100;
+            *b += 5;
+        }
+
+        assert_eq!(vec.get(0), Some((&101, &15)));
+        assert_eq!(vec.get(1), Some((&102, &25)));
+        assert_eq!(vec.get(2), Some((&103, &35)));
+    }
 }