1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! `Vec::drain` and its iterator.
use core::iter::FusedIterator;
use core::marker::PhantomData;
use core::ops::{Bound, RangeBounds};
use core::{fmt, mem};
use allocator_api2::alloc::Allocator;
use super::Vec;
impl<'a, T, A: Allocator + Clone> Vec<'a, T, A> {
/// Drain a range of elements.
///
/// # Panics
///
/// Panics if the start of the range is greater than the end, or if the
/// end is greater than `len`.
pub fn drain<R: RangeBounds<usize>>(&mut self, range: R) -> Drain<'_, 'a, T, A> {
let len = self.buf.len();
let start = match range.start_bound() {
Bound::Included(&i) => i,
Bound::Excluded(&i) => i + 1,
Bound::Unbounded => 0,
};
let end = match range.end_bound() {
Bound::Included(&i) => i + 1,
Bound::Excluded(&i) => i,
Bound::Unbounded => len,
};
assert!(start <= end, "drain: start > end");
assert!(end <= len, "drain: end > len");
// Strategy: move the [start, end) elements into a heap-backed
// staging vec, leaving a "hole" in `self.buf`. The Drain iterator
// yields from the staging vec; on drop, it closes the hole by
// sliding the tail (already-initialized) elements left and
// truncating `self.buf` to `len - drained`.
//
// We use safe `swap_remove` semantics implemented via element-wise
// shifting: pull out the drained values, then compact.
//
// For the staging vec we use a `std`-allocated Vec to keep the
// drain iterator simple and avoid arena chunk churn.
// 1. Move out the drained slice in order.
let drained_count = end - start;
let mut staged: allocator_api2::vec::Vec<T> = allocator_api2::vec::Vec::with_capacity(drained_count);
// Take each element by index; we rebuild the vector after.
// We do this via mem::replace on the slot using a "ghost" value:
// since `T` isn't `Default`, we can't fill — so instead, we
// truncate-then-rebuild. We pop the whole tail (after `start`)
// into a temp, then push the kept tail back.
let tail_count = len - start;
let mut tail_staging: allocator_api2::vec::Vec<T> = allocator_api2::vec::Vec::with_capacity(tail_count);
// Pop produces reverse order; collect then reverse.
for _ in 0..tail_count {
tail_staging.push(self.buf.pop().expect("tail length matches len-start"));
}
tail_staging.reverse();
// `tail_staging` now contains [start, len) in original order.
// First `drained_count` go into `staged`; remainder go back.
let mut iter = tail_staging.into_iter();
for _ in 0..drained_count {
staged.push(iter.next().expect("drained_count <= tail"));
}
// Remaining items go back into `self.buf`. We may need capacity,
// but `self.buf` still has its existing capacity; pop didn't free
// it.
for item in iter {
self.buf
.push_within_cap(item)
.ok()
.expect("returning original tail to existing capacity");
}
Drain {
staged: staged.into_iter(),
_marker: PhantomData,
}
}
}
/// Draining iterator returned from [`Vec::drain`].
pub struct Drain<'d, 'a, T, A: Allocator + Clone> {
staged: allocator_api2::vec::IntoIter<T>,
_marker: PhantomData<&'d mut Vec<'a, T, A>>,
}
impl<T: fmt::Debug, A: Allocator + Clone> fmt::Debug for Drain<'_, '_, T, A> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Drain").field("remaining", &self.staged.len()).finish()
}
}
// `Drain` is `!Send`/`!Sync` for the same reason as `Vec`: it borrows a
// vector tied to one arena thread.
impl<T, A: Allocator + Clone> Iterator for Drain<'_, '_, T, A> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
self.staged.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.staged.size_hint()
}
}
impl<T, A: Allocator + Clone> DoubleEndedIterator for Drain<'_, '_, T, A> {
#[inline]
fn next_back(&mut self) -> Option<T> {
self.staged.next_back()
}
}
impl<T, A: Allocator + Clone> ExactSizeIterator for Drain<'_, '_, T, A> {}
impl<T, A: Allocator + Clone> FusedIterator for Drain<'_, '_, T, A> {}
impl<T, A: Allocator + Clone> Drop for Drain<'_, '_, T, A> {
#[inline]
#[cfg_attr(test, mutants::skip)] // body has no observable side effects
fn drop(&mut self) {
// Remaining staged items are dropped by `allocator_api2::vec::IntoIter`'s
// own `Drop`. Hole has already been closed at construction time.
let _ = &mut self.staged;
let _ = mem::size_of::<T>();
}
}