wordvec/
into_iter.rs

1use alloc::alloc::dealloc;
2use core::array;
3use core::iter::{self, FusedIterator};
4use core::mem::{ManuallyDrop, MaybeUninit, needs_drop};
5use core::ptr::{self, NonNull};
6
7use crate::{Allocated, Large, ParsedMarker, WordVec};
8
9impl<T, const N: usize> IntoIterator for WordVec<T, N> {
10    type Item = T;
11    type IntoIter = IntoIter<T, N>;
12
13    fn into_iter(self) -> Self::IntoIter {
14        let mut this = ManuallyDrop::new(self); // the real destructor is called by `IntoIter`.
15        match this.0.parse_marker() {
16            ParsedMarker::Small(len) => {
17                // SAFETY: indicated by marker
18                let small = unsafe { ManuallyDrop::take(&mut this.0.small) };
19                let data = small.data;
20                let valid = data.into_iter().take(len.into());
21                IntoIter(IntoIterInner::Small(valid.into_iter()))
22            }
23            ParsedMarker::Large => {
24                // SAFETY: indicated by marker
25                let alloc = unsafe { this.0.large.0 };
26                IntoIter(IntoIterInner::Large { alloc, start: 0 })
27            }
28        }
29    }
30}
31
32/// Implements [`IntoIterator`] for [`WordVec`].
33pub struct IntoIter<T, const N: usize>(IntoIterInner<T, N>);
34
35enum IntoIterInner<T, const N: usize> {
36    Small(iter::Take<array::IntoIter<MaybeUninit<T>, N>>),
37    Large { alloc: NonNull<Allocated<T>>, start: usize },
38}
39
40impl<T, const N: usize> Iterator for IntoIter<T, N> {
41    type Item = T;
42
43    fn next(&mut self) -> Option<Self::Item> {
44        match &mut self.0 {
45            IntoIterInner::Small(iter) => {
46                // SAFETY: only initialized values are taken
47                iter.next().map(|uninit| unsafe { uninit.assume_init() })
48            }
49            IntoIterInner::Large { alloc, start } => {
50                // SAFETY: the header is always valid before drop.
51                let len = unsafe { alloc.as_mut() }.len;
52
53                let index = *start;
54                if index >= len {
55                    return None;
56                }
57                *start = index + 1;
58
59                let value = unsafe { Allocated::data_start(*alloc) };
60                // SAFETY: index was not consumed before this function was called.
61                let value = unsafe { ptr::read(value.add(index)) };
62                Some(value)
63            }
64        }
65    }
66}
67
68// Small: array::IntoIter implements FusedIterator.
69// Large: we check `index >= len` before incrementing (not doing so may lead to UB caused by
70// overflow).
71impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}
72
73impl<T, const N: usize> Drop for IntoIter<T, N> {
74    fn drop(&mut self) {
75        match &mut self.0 {
76            IntoIterInner::Small(iter) => {
77                // SAFETY: only initialized values are taken
78                iter.by_ref().for_each(|uninit| drop(unsafe { uninit.assume_init() }));
79            }
80            IntoIterInner::Large { alloc, start } => {
81                // SAFETY: the header is always valid before drop.
82                let &mut Allocated { len, cap, .. } = unsafe { alloc.as_mut() };
83
84                if needs_drop::<T>() {
85                    let value = unsafe { Allocated::data_start(*alloc) };
86                    // SAFETY: `start` <= `len`.
87                    let start_ptr = unsafe { value.add(*start) };
88                    // SAFETY: All items in the range start..len are still initialized and need
89                    // dropping.
90                    unsafe {
91                        let to_drop = ptr::slice_from_raw_parts_mut(start_ptr, len - *start);
92                        ptr::drop_in_place(to_drop);
93                    }
94                }
95
96                let layout = Large::<T>::new_layout(cap);
97                // SAFETY: alloc is valid before dropped; layout is provided by the header itself.
98                unsafe {
99                    dealloc(alloc.as_ptr().cast(), layout);
100                }
101            }
102        }
103    }
104}
wordvec/into_iter.rs

wordvec/
into_iter.rs
