diff --git a/src/into_iter.rs b/src/into_iter.rs new file mode 100644 index 0000000..0e726cb --- /dev/null +++ b/src/into_iter.rs @@ -0,0 +1,198 @@ +use core::fmt::Debug; +use core::iter::FusedIterator; +use core::{mem, ptr}; + +use crate::alloc::{Allocator, Global}; +use crate::vector::{RawVector, Vector}; + +/// An iterator that moves out of a vector. +/// +/// This `struct` is created by the `into_iter` method on [`Vector`] +/// (provided by the [`IntoIterator`] trait). +pub struct IntoIter { + // The elements in the `start..end` range are still owned by the iterator, + // the rest of the buffer has been moved out. `raw`'s length is set to zero + // so that it never drops items on its own. + raw: RawVector, + allocator: A, + start: usize, + end: usize, +} + +impl IntoIter { + pub(crate) fn new(mut vector: Vector) -> Self { + unsafe { + let end = vector.len(); + let allocator = ptr::read(&vector.allocator); + let mut raw = mem::take(&mut vector.raw); + raw.header.len = 0; + + // The allocator was moved out of the vector, so it must not be dropped. + mem::forget(vector); + + IntoIter { + raw, + allocator, + start: 0, + end, + } + } + } + + /// Returns the remaining items of this iterator as a slice. + #[inline] + pub fn as_slice(&self) -> &[T] { + unsafe { core::slice::from_raw_parts(self.ptr(self.start), self.len()) } + } + + /// Returns the remaining items of this iterator as a mutable slice. + #[inline] + pub fn as_mut_slice(&mut self) -> &mut [T] { + unsafe { core::slice::from_raw_parts_mut(self.ptr(self.start), self.len()) } + } + + /// Returns a reference to the underlying allocator. + #[inline] + pub fn allocator(&self) -> &A { + &self.allocator + } + + #[inline] + fn len(&self) -> usize { + self.end - self.start + } + + #[inline] + unsafe fn ptr(&self, idx: usize) -> *mut T { + self.raw.data_ptr().add(idx) + } +} + +impl Iterator for IntoIter { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + if self.start == self.end { + return None; + } + + unsafe { + let item = ptr::read(self.ptr(self.start)); + self.start += 1; + + Some(item) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.len(), Some(self.len())) + } + + #[inline] + fn count(self) -> usize { + self.len() + } +} + +impl DoubleEndedIterator for IntoIter { + #[inline] + fn next_back(&mut self) -> Option { + if self.start == self.end { + return None; + } + + unsafe { + self.end -= 1; + + Some(ptr::read(self.ptr(self.end))) + } + } +} + +impl ExactSizeIterator for IntoIter { + #[inline] + fn len(&self) -> usize { + IntoIter::len(self) + } +} + +impl FusedIterator for IntoIter {} + +impl AsRef<[T]> for IntoIter { + fn as_ref(&self) -> &[T] { + self.as_slice() + } +} + +impl Debug for IntoIter { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> { + f.debug_tuple("IntoIter").field(&self.as_slice()).finish() + } +} + +unsafe impl Send for IntoIter {} +unsafe impl Sync for IntoIter {} + +impl Drop for IntoIter { + fn drop(&mut self) { + unsafe { + ptr::drop_in_place(self.as_mut_slice()); + // The raw vector's length is zero, so this only deallocates the buffer. + self.raw.deallocate(&self.allocator); + } + } +} + +#[test] +fn into_iter() { + fn num(val: u32) -> Box { + Box::new(val) + } + + let mut v = Vector::new(); + v.push(num(0)); + v.push(num(1)); + v.push(num(2)); + + let mut iter = v.into_iter(); + assert_eq!(iter.len(), 3); + assert_eq!(iter.as_slice(), &[num(0), num(1), num(2)]); + assert_eq!(iter.next(), Some(num(0))); + assert_eq!(iter.next_back(), Some(num(2))); + assert_eq!(iter.as_slice(), &[num(1)]); + assert_eq!(iter.next(), Some(num(1))); + assert_eq!(iter.next(), None); + assert_eq!(iter.next_back(), None); + + // The remaining items are dropped with the iterator. + let mut v = Vector::new(); + v.push(num(0)); + v.push(num(1)); + let mut iter = v.into_iter(); + assert_eq!(iter.next(), Some(num(0))); + mem::drop(iter); + + // Empty and unallocated vectors. + let v: Vector> = Vector::new(); + assert_eq!(v.into_iter().next(), None); + let v: Vector> = Vector::with_capacity(16); + assert_eq!(v.into_iter().count(), 0); + + // Zero-sized types. + let mut v = Vector::new(); + v.push(()); + v.push(()); + assert_eq!(v.into_iter().collect::>().len(), 2); + + // In a for loop, with a non-global allocator. + let mut v: Vector, &Global> = Vector::new_in(&Global); + v.push(num(1)); + v.push(num(2)); + let mut sum = 0; + for item in v { + sum += *item; + } + assert_eq!(sum, 3); +} diff --git a/src/lib.rs b/src/lib.rs index 6f9c734..22c3239 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,8 +5,10 @@ mod raw; mod shared; mod vector; mod drain; +mod into_iter; mod splice; +pub use into_iter::IntoIter; pub use raw::{AtomicRefCount, BufferSize, DefaultRefCount, RefCount}; pub use shared::{AtomicSharedVector, RefCountedVector, SharedVector}; pub use vector::{Vector, RawVector}; diff --git a/src/vector.rs b/src/vector.rs index 7a47147..f4cf006 100644 --- a/src/vector.rs +++ b/src/vector.rs @@ -6,6 +6,7 @@ use core::ops::RangeBounds; use crate::alloc::{AllocError, Allocator, Global}; use crate::drain::Drain; +use crate::into_iter::IntoIter; use crate::raw::{ self, buffer_layout, AtomicRefCount, BufferSize, Header, HeaderBuffer, RefCount, VecHeader, move_data, }; @@ -144,7 +145,7 @@ impl RawVector { } #[inline] - fn data_ptr(&self) -> *mut T { + pub(crate) fn data_ptr(&self) -> *mut T { self.data.as_ptr() } @@ -1510,6 +1511,14 @@ impl Default for Vector { } } +impl IntoIterator for Vector { + type Item = T; + type IntoIter = IntoIter; + fn into_iter(self) -> IntoIter { + IntoIter::new(self) + } +} + impl<'a, T, A: Allocator> IntoIterator for &'a Vector { type Item = &'a T; type IntoIter = core::slice::Iter<'a, T>;