flex_alloc/vec/
drain.rs

1use core::fmt;
2use core::fmt::Debug;
3use core::iter::FusedIterator;
4use core::ops::Range;
5use core::ptr;
6use core::slice;
7
8use crate::capacity::Index;
9
10use super::buffer::VecBuffer;
11use super::index_panic;
12
13/// A struct used for draining items from a Vec as an iterator.
14pub struct Drain<'d, B: VecBuffer> {
15    pub(super) range: Range<usize>,
16    pub(super) remain: Range<usize>,
17    pub(super) tail_length: usize,
18    pub(super) buf: &'d mut B,
19}
20
21impl<'d, B: VecBuffer> Drain<'d, B> {
22    pub(super) fn new(buf: &'d mut B, range: Range<usize>) -> Self {
23        let len = buf.length().to_usize();
24        if range.end < range.start || range.end > len {
25            index_panic();
26        }
27        if len > 0 {
28            // SAFETY: buffer capacity is established as > 0
29            unsafe { buf.set_length(B::Index::from_usize(range.start)) };
30        }
31        let tail_length = len - range.end;
32        Self {
33            range: range.clone(),
34            remain: range,
35            tail_length,
36            buf,
37        }
38    }
39
40    /// Access the remaining items as a slice reference.
41    pub fn as_slice(&self) -> &[B::Item] {
42        unsafe {
43            slice::from_raw_parts(
44                self.buf.data_ptr().add(self.remain.start),
45                self.remain.len(),
46            )
47        }
48    }
49
50    /// Access the remaining items as a mutable slice reference.
51    pub fn as_mut_slice(&mut self) -> &mut [B::Item] {
52        unsafe {
53            slice::from_raw_parts_mut(
54                self.buf.data_ptr_mut().add(self.remain.start),
55                self.remain.len(),
56            )
57        }
58    }
59
60    /// Check if there are remaining items in the iterator.
61    pub const fn is_empty(&self) -> bool {
62        self.len() == 0
63    }
64
65    /// Get the number of remaining items in the iterator.
66    pub const fn len(&self) -> usize {
67        self.remain.end - self.remain.start
68    }
69
70    pub(super) fn clear_remain(&mut self) {
71        let remain_len = self.len();
72        if remain_len > 0 {
73            unsafe {
74                ptr::drop_in_place(self.as_mut_slice().as_mut_ptr());
75            }
76            self.remain.start = self.remain.end;
77        }
78    }
79
80    /// Abort the drain operation, leaving the remaining items contained in the `Vec` instance.
81    pub fn keep_rest(mut self) {
82        let len = self.len();
83        let shift = self.remain.start - self.range.start;
84        if len > 0 && shift > 0 {
85            unsafe {
86                let head = self.buf.data_ptr_mut().add(self.range.start);
87                ptr::copy(head.add(shift), head, len);
88            }
89        }
90        self.range.start += len;
91        self.remain = Range {
92            start: self.range.start,
93            end: self.range.start,
94        };
95    }
96}
97
98impl<'d, B: VecBuffer> AsRef<[B::Item]> for Drain<'d, B> {
99    fn as_ref(&self) -> &[B::Item] {
100        self.as_slice()
101    }
102}
103
104impl<'d, B: VecBuffer> fmt::Debug for Drain<'d, B>
105where
106    B::Item: Debug,
107{
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        f.debug_tuple("Drain").field(&self.as_slice()).finish()
110    }
111}
112
113impl<'d, B: VecBuffer> Iterator for Drain<'d, B> {
114    type Item = B::Item;
115
116    fn next(&mut self) -> Option<Self::Item> {
117        let index = self.remain.start;
118        if index != self.remain.end {
119            self.remain.start = index + 1;
120            unsafe {
121                let read = self.buf.data_ptr().add(index);
122                Some(ptr::read(read))
123            }
124        } else {
125            None
126        }
127    }
128
129    #[inline]
130    fn count(self) -> usize
131    where
132        Self: Sized,
133    {
134        self.len()
135    }
136
137    #[inline]
138    fn size_hint(&self) -> (usize, Option<usize>) {
139        let len = self.len();
140        (len, Some(len))
141    }
142}
143
144impl<'d, B: VecBuffer> DoubleEndedIterator for Drain<'d, B> {
145    fn next_back(&mut self) -> Option<Self::Item> {
146        let mut index = self.remain.end;
147        if index != self.remain.start {
148            index -= 1;
149            self.remain.end = index;
150            unsafe {
151                let read = self.buf.data_ptr().add(index);
152                Some(ptr::read(read))
153            }
154        } else {
155            None
156        }
157    }
158}
159
160impl<'d, B: VecBuffer> ExactSizeIterator for Drain<'d, B> {}
161
162impl<'d, B: VecBuffer> FusedIterator for Drain<'d, B> {}
163
164impl<'d, B: VecBuffer> Drop for Drain<'d, B> {
165    fn drop(&mut self) {
166        self.clear_remain();
167        if self.tail_length > 0 {
168            unsafe {
169                let head = self.buf.data_ptr_mut().add(self.range.start);
170                ptr::copy(head.add(self.range.len()), head, self.tail_length);
171            }
172        }
173        let len = self.range.start + self.tail_length;
174        if len > 0 {
175            // SAFETY: capacity is established as > 0
176            unsafe {
177                self.buf
178                    .set_length(B::Index::from_usize(self.range.start + self.tail_length))
179            }
180        }
181    }
182}
flex_alloc/vec/drain.rs

flex_alloc/vec/
drain.rs
