tiny_vec/
iter.rs

1//! Iterator implementation for TinyVec
2
3use core::iter::FusedIterator;
4use core::marker::PhantomData;
5use core::mem::{self, ManuallyDrop, MaybeUninit};
6use core::ptr;
7
8use alloc::slice;
9
10use crate::raw::RawVec;
11use crate::TinyVec;
12
13enum Kind<T, const N: usize> {
14    Stack([MaybeUninit<T>; N]),
15    Heap(RawVec<T>),
16}
17
18impl<T, const N: usize> Kind<T, N> {
19    const fn ptr(&self) -> *const T {
20        match self {
21            Kind::Stack(s) => s.as_ptr().cast(),
22            Kind::Heap(rv) => rv.ptr.as_ptr()
23        }
24    }
25
26    const fn ptr_mut(&mut self) -> *mut T {
27        match self {
28            Kind::Stack(s) => s.as_mut_ptr().cast(),
29            Kind::Heap(rv) => rv.ptr.as_ptr()
30        }
31    }
32}
33
34/// Iterator over the elements of a [TinyVec]
35///
36/// This struct is returned from the [TinyVec::into_iter] function
37pub struct TinyVecIter<T, const N: usize> {
38    start: usize,
39    len: usize,
40    buf: Kind<T, N>,
41    _marker: PhantomData<TinyVec<T, N>>,
42}
43
44impl<T, const N: usize> TinyVecIter<T, N> {
45    /// Casts the remaining portion of this iterator as a slice of T
46    pub const fn as_slice(&self) -> &[T] {
47        unsafe {
48            let ptr = self.buf.ptr().add(self.start);
49            slice::from_raw_parts(ptr, self.len)
50        }
51    }
52
53    /// Casts the remaining portion of this iterator as a mutable slice of T
54    pub const fn as_mut_slice(&mut self) -> &mut [T] {
55        unsafe {
56            let ptr = self.buf.ptr_mut().add(self.start);
57            slice::from_raw_parts_mut(ptr, self.len)
58        }
59    }
60}
61
62impl<T, const N: usize> Drop for TinyVecIter<T, N> {
63    fn drop(&mut self) {
64        let raw = match self.buf {
65            Kind::Stack(_) => None,
66            Kind::Heap(raw) => Some(raw),
67        };
68
69        if mem::needs_drop::<T>() {
70            for e in self {
71                mem::drop(e);
72            }
73        }
74
75        if let Some(mut raw) = raw {
76            unsafe { raw.destroy(); }
77        }
78    }
79}
80
81impl<T, const N: usize> Iterator for TinyVecIter<T, N> {
82    type Item = T;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        if self.len == 0 {
86            None
87        } else {
88            let e = unsafe {
89                self.buf
90                    .ptr_mut()
91                    .add(self.start)
92                    .read()
93            };
94            self.start += 1;
95            self.len -= 1;
96            Some(e)
97        }
98    }
99
100    fn size_hint(&self) -> (usize, Option<usize>) {
101        (self.len, Some(self.len))
102    }
103}
104
105impl<T, const N: usize> DoubleEndedIterator for TinyVecIter<T, N> {
106    fn next_back(&mut self) -> Option<Self::Item> {
107        if self.len == 0 {
108            None
109        } else {
110            self.len -= 1;
111            let e = unsafe {
112                self.buf
113                    .ptr_mut()
114                    .add(self.start + self.len)
115                    .read()
116            };
117            Some(e)
118        }
119    }
120}
121
122impl<T, const N: usize> ExactSizeIterator for TinyVecIter<T, N> { }
123
124impl<T, const N: usize> FusedIterator for TinyVecIter<T, N> { }
125
126impl<T, const N: usize> IntoIterator for TinyVec<T, N> {
127    type Item = T;
128
129    type IntoIter = TinyVecIter<T, N>;
130
131    fn into_iter(self) -> Self::IntoIter {
132        let vec = ManuallyDrop::new(self);
133
134        let is_stack = vec.lives_on_stack();
135        let len = vec.len();
136
137        let buf = unsafe {
138            let inner = ptr::read( &vec.inner );
139            if is_stack {
140                Kind::Stack(ManuallyDrop::into_inner( inner.stack ))
141            } else {
142                Kind::Heap(inner.raw)
143            }
144        };
145
146        TinyVecIter { start: 0, len, buf, _marker: PhantomData }
147    }
148}