tiny_vec/
drain.rs

1//! [drain](TinyVec::drain) implementation
2//!
3use crate::TinyVec;
4use core::iter::FusedIterator;
5use core::marker::PhantomData;
6use core::mem::{self, ManuallyDrop};
7use core::ptr::{self, NonNull};
8use core::slice;
9
10/// A draining iterator for [TinyVec]
11///
12/// This struct is created by the [TinyVec::drain] method
13pub struct Drain<'a, T: 'a, const N: usize> {
14    pub (super) remaining_start: usize,
15    pub (super) remaining_len: usize,
16    pub (super) tail_start: usize,
17    pub (super) tail_len: usize,
18    pub (super) vec: NonNull<TinyVec<T, N>>,
19    pub (super) _marker: PhantomData<&'a mut TinyVec<T, N>>,
20}
21
22impl<'a, T: 'a, const N: usize> Drain<'a, T, N> {
23
24    /// Returns the remaining elements as a slice
25    #[inline]
26    pub fn as_slice(&self) -> &[T] {
27        unsafe {
28            let ptr = self.vec.as_ref().as_ptr().add(self.remaining_start);
29            slice::from_raw_parts(ptr, self.remaining_len)
30        }
31    }
32
33    /// Consumes this [Drain], preserving any un-yielded elements
34    ///
35    /// # Example
36    /// ```
37    /// use tiny_vec::TinyVec;
38    ///
39    /// let mut tv = TinyVec::from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
40    ///
41    /// let mut drain = tv.drain(3..7);
42    ///
43    /// assert_eq!(drain.next(), Some(3));
44    /// assert_eq!(drain.next(), Some(4));
45    ///
46    /// drain.keep_rest();
47    ///
48    /// assert_eq!(tv.as_slice(), &[0, 1, 2, 5, 6, 7, 8, 9]);
49    /// ```
50    pub fn keep_rest(self) {
51        let mut slf = ManuallyDrop::new(self);
52
53        /* [ HEAD ] [ yieled  | remaining | yielded_back ] [ TAIL ] */
54        unsafe {
55            let vec = slf.vec.as_mut();
56            let start = vec.len();
57
58            if mem::size_of::<T>() != 0 {
59                let buf = vec.as_mut_ptr();
60
61                let start_ptr = buf.add(start);
62                let remaining_ptr = buf.add(slf.remaining_start);
63
64                if remaining_ptr != start_ptr {
65                    /* First: We move the remaining chunk to the start */
66                    ptr::copy(remaining_ptr, start_ptr, slf.remaining_len);
67                }
68
69                if slf.tail_start != (start + slf.remaining_len) {
70                    let src = buf.add(slf.tail_start);
71                    let dst = start_ptr.add(slf.remaining_len);
72                    /* Now we move the tail */
73                    ptr::copy(src, dst, slf.tail_len);
74                }
75            }
76            vec.set_len(start + slf.remaining_len + slf.tail_len);
77        }
78    }
79}
80
81impl<'a, T: 'a, const N: usize> Iterator for Drain<'a, T, N> {
82    type Item = T;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        if self.remaining_len == 0 {
86            None
87        } else {
88            unsafe {
89                let e = self.vec.as_ref().as_ptr().add(self.remaining_start).read();
90                self.remaining_start += 1;
91                self.remaining_len -= 1;
92                Some(e)
93            }
94        }
95    }
96
97    fn size_hint(&self) -> (usize, Option<usize>) {
98        (self.remaining_len, Some(self.remaining_len))
99    }
100}
101
102impl<'a, T: 'a, const N: usize> DoubleEndedIterator for Drain<'a, T, N> {
103    fn next_back(&mut self) -> Option<Self::Item> {
104        if self.remaining_len == 0 {
105            None
106        } else {
107            unsafe {
108                let e = self.vec.as_ref()
109                                .as_ptr()
110                                .add(self.remaining_start + self.remaining_len - 1)
111                                .read();
112                self.remaining_len -= 1;
113                Some(e)
114            }
115        }
116    }
117}
118
119impl<'a, T: 'a, const N: usize> ExactSizeIterator for Drain<'a, T, N> { }
120
121impl<'a, T: 'a, const N: usize> FusedIterator for Drain<'a, T, N> { }
122
123impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
124    fn drop(&mut self) {
125        /* FIXME: The std implementation does some complex
126         * thing with a DropGuard struct. */
127        for e in &mut *self { drop(e) }
128        unsafe {
129            let vec = self.vec.as_mut();
130
131            let len = vec.len();
132            let ptr = vec.as_mut_ptr();
133
134            let src = ptr.add(self.tail_start);
135            let dst = ptr.add(len);
136
137            ptr::copy(src, dst, self.tail_len);
138
139            vec.set_len(len + self.tail_len);
140        }
141    }
142}
tiny_vec/drain.rs

tiny_vec/
drain.rs
