malloced/
iter.rs

1use crate::sys;
2use core::{
3    marker::PhantomData,
4    mem,
5    ptr::{self, NonNull},
6    slice,
7};
8
9/// An iterator over a
10/// <code>[Malloced](crate::Malloced)<[\[T\]](prim@slice)></code>.
11pub struct SliceIter<T> {
12    pub(super) buf: NonNull<T>,
13    // Marks ownership of an instance of T.
14    pub(super) marker: PhantomData<T>,
15    pub(super) ptr: *mut T,
16    pub(super) end: *mut T,
17}
18
19impl<T> SliceIter<T> {
20    #[inline]
21    fn as_raw_mut_slice(&mut self) -> &mut [T] {
22        unsafe { slice::from_raw_parts_mut(self.ptr, self.len()) }
23    }
24}
25
26impl<T> Drop for SliceIter<T> {
27    #[inline]
28    fn drop(&mut self) {
29        struct DeallocGuard<'a, T: 'a>(&'a mut SliceIter<T>);
30
31        impl<'a, T> Drop for DeallocGuard<'a, T> {
32            #[inline]
33            fn drop(&mut self) {
34                unsafe {
35                    sys::free(self.0.buf.as_ptr() as _);
36                }
37            }
38        }
39
40        // Deallocates the memory slice's on drop. If dropping the elements
41        // panics, the memory will still be deallocated.
42        let guard = DeallocGuard(self);
43
44        // Drop remaining elements.
45        unsafe {
46            ptr::drop_in_place(guard.0.as_raw_mut_slice());
47        }
48    }
49}
50
51impl<T> Iterator for SliceIter<T> {
52    type Item = T;
53
54    #[inline]
55    fn next(&mut self) -> Option<Self::Item> {
56        if self.ptr == self.end {
57            None
58        } else if mem::size_of::<T>() == 0 {
59            // Purposefully don't use `ptr.offset` because for slices with
60            // 0-size elements this would return the same pointer.
61            self.ptr = (self.ptr as *mut i8).wrapping_add(1) as *mut T;
62
63            // Make up a value of this ZST.
64            Some(unsafe { mem::zeroed() })
65        } else {
66            let old = self.ptr;
67            self.ptr = unsafe { self.ptr.offset(1) };
68
69            Some(unsafe { old.read() })
70        }
71    }
72
73    #[inline]
74    fn size_hint(&self) -> (usize, Option<usize>) {
75        let len = self.len();
76        (len, Some(len))
77    }
78
79    #[inline]
80    fn count(self) -> usize {
81        self.len()
82    }
83
84    #[inline]
85    fn last(mut self) -> Option<Self::Item> {
86        self.next_back()
87    }
88}
89
90impl<T> DoubleEndedIterator for SliceIter<T> {
91    #[inline]
92    fn next_back(&mut self) -> Option<Self::Item> {
93        if self.end == self.ptr {
94            None
95        } else if mem::size_of::<T>() == 0 {
96            // Purposefully don't use `ptr.offset` because for slices with
97            // 0-size elements this would return the same pointer.
98            self.ptr = (self.ptr as *mut i8).wrapping_sub(1) as *mut T;
99
100            // Make up a value of this ZST.
101            Some(unsafe { mem::zeroed() })
102        } else {
103            self.end = unsafe { self.end.offset(-1) };
104
105            Some(unsafe { self.end.read() })
106        }
107    }
108}
109
110impl<T> ExactSizeIterator for SliceIter<T> {
111    #[inline]
112    fn len(&self) -> usize {
113        let diff = (self.end as usize).wrapping_sub(self.ptr as usize);
114
115        match diff.checked_div(mem::size_of::<T>()) {
116            Some(len) => len,
117
118            // ZST
119            None => diff,
120        }
121    }
122}
123
124impl<T> core::iter::FusedIterator for SliceIter<T> {}
125
126#[cfg(test)]
127mod tests {
128    use crate::Malloced;
129    use alloc::vec::Vec;
130    use core::fmt::Debug;
131
132    mod collect {
133        use super::*;
134
135        #[track_caller]
136        fn test<T: Copy + Debug + PartialEq>(slice: &[T]) {
137            let result: Vec<T> = Malloced::alloc(slice).unwrap().into_iter().collect();
138            assert_eq!(result, slice);
139        }
140
141        #[test]
142        fn zst() {
143            test(&[()]);
144            test(&[(), ()]);
145        }
146
147        #[test]
148        fn u8() {
149            test(&[1u8]);
150            test(&[1u8, 2u8]);
151        }
152
153        #[test]
154        fn u16() {
155            test(&[1u16]);
156            test(&[1u16, 2u16]);
157        }
158
159        #[test]
160        fn usize() {
161            test(&[1usize]);
162            test(&[1usize, 2usize]);
163        }
164    }
165
166    mod len {
167        use super::*;
168
169        #[track_caller]
170        fn test(slice: &[impl Copy]) {
171            let iter = Malloced::alloc(slice).unwrap().into_iter();
172            assert_eq!(iter.len(), slice.len());
173        }
174
175        #[test]
176        fn zst() {
177            test(&[()]);
178            test(&[(), ()]);
179        }
180
181        #[test]
182        fn u8() {
183            test(&[1u8]);
184            test(&[1u8, 2u8]);
185        }
186
187        #[test]
188        fn u16() {
189            test(&[1u16]);
190            test(&[1u16, 2u16]);
191        }
192
193        #[test]
194        fn usize() {
195            test(&[1usize]);
196            test(&[1usize, 2usize]);
197        }
198    }
199}