rune_alloc/vec_deque/
drain.rs

1use core::fmt;
2use core::iter::FusedIterator;
3use core::marker::PhantomData;
4use core::mem;
5
6use crate::alloc::{Allocator, Global, SizedTypeProperties};
7use crate::ptr::{self, NonNull};
8
9use super::VecDeque;
10
11/// A draining iterator over the elements of a `VecDeque`.
12///
13/// This `struct` is created by the [`drain`] method on [`VecDeque`]. See its
14/// documentation for more.
15///
16/// [`drain`]: VecDeque::drain
17pub struct Drain<'a, T, A = Global>
18where
19    T: 'a,
20    A: Allocator,
21{
22    // We can't just use a &mut VecDeque<T, A>, as that would make Drain invariant over T
23    // and we want it to be covariant instead
24    deque: NonNull<VecDeque<T, A>>,
25    // drain_start is stored in deque.len
26    drain_len: usize,
27    // index into the logical array, not the physical one (always lies in [0..deque.len))
28    idx: usize,
29    // number of elements after the drain range
30    tail_len: usize,
31    remaining: usize,
32    // Needed to make Drain covariant over T
33    _marker: PhantomData<&'a T>,
34}
35
36impl<'a, T, A> Drain<'a, T, A>
37where
38    A: Allocator,
39{
40    pub(super) unsafe fn new(
41        deque: &'a mut VecDeque<T, A>,
42        drain_start: usize,
43        drain_len: usize,
44    ) -> Self {
45        let orig_len = mem::replace(&mut deque.len, drain_start);
46        let tail_len = orig_len - drain_start - drain_len;
47        Drain {
48            deque: NonNull::from(deque),
49            drain_len,
50            idx: drain_start,
51            tail_len,
52            remaining: drain_len,
53            _marker: PhantomData,
54        }
55    }
56
57    // Only returns pointers to the slices, as that's all we need
58    // to drop them. May only be called if `self.remaining != 0`.
59    unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) {
60        unsafe {
61            let deque = self.deque.as_ref();
62
63            // We know that `self.idx + self.remaining <= deque.len <= usize::MAX`, so this won't overflow.
64            let logical_remaining_range = self.idx..self.idx + self.remaining;
65
66            // SAFETY: `logical_remaining_range` represents the
67            // range into the logical buffer of elements that
68            // haven't been drained yet, so they're all initialized,
69            // and `slice::range(start..end, end) == start..end`,
70            // so the preconditions for `slice_ranges` are met.
71            let (a_range, b_range) =
72                deque.slice_ranges(logical_remaining_range.clone(), logical_remaining_range.end);
73            (deque.buffer_range(a_range), deque.buffer_range(b_range))
74        }
75    }
76}
77
78impl<T, A> fmt::Debug for Drain<'_, T, A>
79where
80    T: fmt::Debug,
81    A: Allocator,
82{
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        f.debug_tuple("Drain")
85            .field(&self.drain_len)
86            .field(&self.idx)
87            .field(&self.tail_len)
88            .field(&self.remaining)
89            .finish()
90    }
91}
92
93unsafe impl<T: Sync, A: Allocator + Sync> Sync for Drain<'_, T, A> {}
94unsafe impl<T: Send, A: Allocator + Send> Send for Drain<'_, T, A> {}
95
96impl<T, A> Drop for Drain<'_, T, A>
97where
98    A: Allocator,
99{
100    fn drop(&mut self) {
101        struct DropGuard<'r, 'a, T, A>(&'r mut Drain<'a, T, A>)
102        where
103            A: Allocator;
104
105        impl<T, A> Drop for DropGuard<'_, '_, T, A>
106        where
107            A: Allocator,
108        {
109            fn drop(&mut self) {
110                if self.0.remaining != 0 {
111                    unsafe {
112                        // SAFETY: We just checked that `self.remaining != 0`.
113                        let (front, back) = self.0.as_slices();
114                        ptr::drop_in_place(front);
115                        ptr::drop_in_place(back);
116                    }
117                }
118
119                let source_deque = unsafe { self.0.deque.as_mut() };
120
121                let drain_start = source_deque.len();
122                let drain_len = self.0.drain_len;
123                let drain_end = drain_start + drain_len;
124
125                let orig_len = self.0.tail_len + drain_end;
126
127                if T::IS_ZST {
128                    // no need to copy around any memory if T is a ZST
129                    source_deque.len = orig_len - drain_len;
130                    return;
131                }
132
133                let head_len = drain_start;
134                let tail_len = self.0.tail_len;
135
136                match (head_len, tail_len) {
137                    (0, 0) => {
138                        source_deque.head = 0;
139                        source_deque.len = 0;
140                    }
141                    (0, _) => {
142                        source_deque.head = source_deque.to_physical_idx(drain_len);
143                        source_deque.len = orig_len - drain_len;
144                    }
145                    (_, 0) => {
146                        source_deque.len = orig_len - drain_len;
147                    }
148                    _ => unsafe {
149                        if head_len <= tail_len {
150                            source_deque.wrap_copy(
151                                source_deque.head,
152                                source_deque.to_physical_idx(drain_len),
153                                head_len,
154                            );
155                            source_deque.head = source_deque.to_physical_idx(drain_len);
156                            source_deque.len = orig_len - drain_len;
157                        } else {
158                            source_deque.wrap_copy(
159                                source_deque.to_physical_idx(head_len + drain_len),
160                                source_deque.to_physical_idx(head_len),
161                                tail_len,
162                            );
163                            source_deque.len = orig_len - drain_len;
164                        }
165                    },
166                }
167            }
168        }
169
170        let guard = DropGuard(self);
171
172        if guard.0.remaining != 0 {
173            unsafe {
174                // SAFETY: We just checked that `self.remaining != 0`.
175                let (front, back) = guard.0.as_slices();
176                // since idx is a logical index, we don't need to worry about wrapping.
177                guard.0.idx += front.len();
178                guard.0.remaining -= front.len();
179                ptr::drop_in_place(front);
180                guard.0.remaining = 0;
181                ptr::drop_in_place(back);
182            }
183        }
184
185        // Dropping `guard` handles moving the remaining elements into place.
186    }
187}
188
189impl<T, A> Iterator for Drain<'_, T, A>
190where
191    A: Allocator,
192{
193    type Item = T;
194
195    #[inline]
196    fn next(&mut self) -> Option<T> {
197        if self.remaining == 0 {
198            return None;
199        }
200        let wrapped_idx = unsafe { self.deque.as_ref().to_physical_idx(self.idx) };
201        self.idx += 1;
202        self.remaining -= 1;
203        Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
204    }
205
206    #[inline]
207    fn size_hint(&self) -> (usize, Option<usize>) {
208        let len = self.remaining;
209        (len, Some(len))
210    }
211}
212
213impl<T, A> DoubleEndedIterator for Drain<'_, T, A>
214where
215    A: Allocator,
216{
217    #[inline]
218    fn next_back(&mut self) -> Option<T> {
219        if self.remaining == 0 {
220            return None;
221        }
222        self.remaining -= 1;
223        let wrapped_idx = unsafe {
224            self.deque
225                .as_ref()
226                .to_physical_idx(self.idx + self.remaining)
227        };
228        Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
229    }
230}
231
232impl<T, A> ExactSizeIterator for Drain<'_, T, A> where A: Allocator {}
233
234impl<T, A> FusedIterator for Drain<'_, T, A> where A: Allocator {}
rune_alloc/vec_deque/drain.rs

rune_alloc/vec_deque/
drain.rs
