Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit 434e5d1

Browse files
committed
refactor: moving DrainFilter into drain_filter.rs
1 parent 5182776 commit 434e5d1

File tree

2 files changed

+148
-138
lines changed

2 files changed

+148
-138
lines changed

library/alloc/src/vec/drain_filter.rs

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
use core::ptr::{self};
2+
use core::slice::{self};
3+
use crate::alloc::{Allocator, Global};
4+
5+
use super::{Vec};
6+
7+
/// An iterator which uses a closure to determine if an element should be removed.
8+
///
9+
/// This struct is created by [`Vec::drain_filter`].
10+
/// See its documentation for more.
11+
///
12+
/// # Example
13+
///
14+
/// ```
15+
/// #![feature(drain_filter)]
16+
///
17+
/// let mut v = vec![0, 1, 2];
18+
/// let iter: std::vec::DrainFilter<_, _> = v.drain_filter(|x| *x % 2 == 0);
19+
/// ```
20+
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
21+
#[derive(Debug)]
22+
pub struct DrainFilter<
23+
'a,
24+
T,
25+
F,
26+
#[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
27+
> where
28+
F: FnMut(&mut T) -> bool,
29+
{
30+
pub(super) vec: &'a mut Vec<T, A>,
31+
/// The index of the item that will be inspected by the next call to `next`.
32+
pub(super) idx: usize,
33+
/// The number of items that have been drained (removed) thus far.
34+
pub(super) del: usize,
35+
/// The original length of `vec` prior to draining.
36+
pub(super) old_len: usize,
37+
/// The filter test predicate.
38+
pub(super) pred: F,
39+
/// A flag that indicates a panic has occurred in the filter test predicate.
40+
/// This is used as a hint in the drop implementation to prevent consumption
41+
/// of the remainder of the `DrainFilter`. Any unprocessed items will be
42+
/// backshifted in the `vec`, but no further items will be dropped or
43+
/// tested by the filter predicate.
44+
pub(super) panic_flag: bool,
45+
}
46+
47+
impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
48+
where
49+
F: FnMut(&mut T) -> bool,
50+
{
51+
/// Returns a reference to the underlying allocator.
52+
#[unstable(feature = "allocator_api", issue = "32838")]
53+
#[inline]
54+
pub fn allocator(&self) -> &A {
55+
self.vec.allocator()
56+
}
57+
}
58+
59+
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
60+
impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
61+
where
62+
F: FnMut(&mut T) -> bool,
63+
{
64+
type Item = T;
65+
66+
fn next(&mut self) -> Option<T> {
67+
unsafe {
68+
while self.idx < self.old_len {
69+
let i = self.idx;
70+
let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
71+
self.panic_flag = true;
72+
let drained = (self.pred)(&mut v[i]);
73+
self.panic_flag = false;
74+
// Update the index *after* the predicate is called. If the index
75+
// is updated prior and the predicate panics, the element at this
76+
// index would be leaked.
77+
self.idx += 1;
78+
if drained {
79+
self.del += 1;
80+
return Some(ptr::read(&v[i]));
81+
} else if self.del > 0 {
82+
let del = self.del;
83+
let src: *const T = &v[i];
84+
let dst: *mut T = &mut v[i - del];
85+
ptr::copy_nonoverlapping(src, dst, 1);
86+
}
87+
}
88+
None
89+
}
90+
}
91+
92+
fn size_hint(&self) -> (usize, Option<usize>) {
93+
(0, Some(self.old_len - self.idx))
94+
}
95+
}
96+
97+
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
98+
impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
99+
where
100+
F: FnMut(&mut T) -> bool,
101+
{
102+
fn drop(&mut self) {
103+
struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
104+
where
105+
F: FnMut(&mut T) -> bool,
106+
{
107+
drain: &'b mut DrainFilter<'a, T, F, A>,
108+
}
109+
110+
impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
111+
where
112+
F: FnMut(&mut T) -> bool,
113+
{
114+
fn drop(&mut self) {
115+
unsafe {
116+
if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
117+
// This is a pretty messed up state, and there isn't really an
118+
// obviously right thing to do. We don't want to keep trying
119+
// to execute `pred`, so we just backshift all the unprocessed
120+
// elements and tell the vec that they still exist. The backshift
121+
// is required to prevent a double-drop of the last successfully
122+
// drained item prior to a panic in the predicate.
123+
let ptr = self.drain.vec.as_mut_ptr();
124+
let src = ptr.add(self.drain.idx);
125+
let dst = src.sub(self.drain.del);
126+
let tail_len = self.drain.old_len - self.drain.idx;
127+
src.copy_to(dst, tail_len);
128+
}
129+
self.drain.vec.set_len(self.drain.old_len - self.drain.del);
130+
}
131+
}
132+
}
133+
134+
let backshift = BackshiftOnDrop { drain: self };
135+
136+
// Attempt to consume any remaining elements if the filter predicate
137+
// has not yet panicked. We'll backshift any remaining elements
138+
// whether we've already panicked or if the consumption here panics.
139+
if !backshift.drain.panic_flag {
140+
backshift.drain.for_each(drop);
141+
}
142+
}
143+
}

library/alloc/src/vec/mod.rs

Lines changed: 5 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,11 @@ use crate::boxed::Box;
7474
use crate::collections::TryReserveError;
7575
use crate::raw_vec::RawVec;
7676

77+
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
78+
pub use self::drain_filter::DrainFilter;
79+
80+
mod drain_filter;
81+
7782
/// A contiguous growable array type, written `Vec<T>` but pronounced 'vector'.
7883
///
7984
/// # Examples
@@ -3585,141 +3590,3 @@ impl<T, A: Allocator> Drain<'_, T, A> {
35853590
self.tail_start = new_tail_start;
35863591
}
35873592
}
3588-
3589-
/// An iterator which uses a closure to determine if an element should be removed.
3590-
///
3591-
/// This struct is created by [`Vec::drain_filter`].
3592-
/// See its documentation for more.
3593-
///
3594-
/// # Example
3595-
///
3596-
/// ```
3597-
/// #![feature(drain_filter)]
3598-
///
3599-
/// let mut v = vec![0, 1, 2];
3600-
/// let iter: std::vec::DrainFilter<_, _> = v.drain_filter(|x| *x % 2 == 0);
3601-
/// ```
3602-
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
3603-
#[derive(Debug)]
3604-
pub struct DrainFilter<
3605-
'a,
3606-
T,
3607-
F,
3608-
#[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
3609-
> where
3610-
F: FnMut(&mut T) -> bool,
3611-
{
3612-
vec: &'a mut Vec<T, A>,
3613-
/// The index of the item that will be inspected by the next call to `next`.
3614-
idx: usize,
3615-
/// The number of items that have been drained (removed) thus far.
3616-
del: usize,
3617-
/// The original length of `vec` prior to draining.
3618-
old_len: usize,
3619-
/// The filter test predicate.
3620-
pred: F,
3621-
/// A flag that indicates a panic has occurred in the filter test predicate.
3622-
/// This is used as a hint in the drop implementation to prevent consumption
3623-
/// of the remainder of the `DrainFilter`. Any unprocessed items will be
3624-
/// backshifted in the `vec`, but no further items will be dropped or
3625-
/// tested by the filter predicate.
3626-
panic_flag: bool,
3627-
}
3628-
3629-
impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
3630-
where
3631-
F: FnMut(&mut T) -> bool,
3632-
{
3633-
/// Returns a reference to the underlying allocator.
3634-
#[unstable(feature = "allocator_api", issue = "32838")]
3635-
#[inline]
3636-
pub fn allocator(&self) -> &A {
3637-
self.vec.allocator()
3638-
}
3639-
}
3640-
3641-
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
3642-
impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
3643-
where
3644-
F: FnMut(&mut T) -> bool,
3645-
{
3646-
type Item = T;
3647-
3648-
fn next(&mut self) -> Option<T> {
3649-
unsafe {
3650-
while self.idx < self.old_len {
3651-
let i = self.idx;
3652-
let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
3653-
self.panic_flag = true;
3654-
let drained = (self.pred)(&mut v[i]);
3655-
self.panic_flag = false;
3656-
// Update the index *after* the predicate is called. If the index
3657-
// is updated prior and the predicate panics, the element at this
3658-
// index would be leaked.
3659-
self.idx += 1;
3660-
if drained {
3661-
self.del += 1;
3662-
return Some(ptr::read(&v[i]));
3663-
} else if self.del > 0 {
3664-
let del = self.del;
3665-
let src: *const T = &v[i];
3666-
let dst: *mut T = &mut v[i - del];
3667-
ptr::copy_nonoverlapping(src, dst, 1);
3668-
}
3669-
}
3670-
None
3671-
}
3672-
}
3673-
3674-
fn size_hint(&self) -> (usize, Option<usize>) {
3675-
(0, Some(self.old_len - self.idx))
3676-
}
3677-
}
3678-
3679-
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
3680-
impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
3681-
where
3682-
F: FnMut(&mut T) -> bool,
3683-
{
3684-
fn drop(&mut self) {
3685-
struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
3686-
where
3687-
F: FnMut(&mut T) -> bool,
3688-
{
3689-
drain: &'b mut DrainFilter<'a, T, F, A>,
3690-
}
3691-
3692-
impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
3693-
where
3694-
F: FnMut(&mut T) -> bool,
3695-
{
3696-
fn drop(&mut self) {
3697-
unsafe {
3698-
if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
3699-
// This is a pretty messed up state, and there isn't really an
3700-
// obviously right thing to do. We don't want to keep trying
3701-
// to execute `pred`, so we just backshift all the unprocessed
3702-
// elements and tell the vec that they still exist. The backshift
3703-
// is required to prevent a double-drop of the last successfully
3704-
// drained item prior to a panic in the predicate.
3705-
let ptr = self.drain.vec.as_mut_ptr();
3706-
let src = ptr.add(self.drain.idx);
3707-
let dst = src.sub(self.drain.del);
3708-
let tail_len = self.drain.old_len - self.drain.idx;
3709-
src.copy_to(dst, tail_len);
3710-
}
3711-
self.drain.vec.set_len(self.drain.old_len - self.drain.del);
3712-
}
3713-
}
3714-
}
3715-
3716-
let backshift = BackshiftOnDrop { drain: self };
3717-
3718-
// Attempt to consume any remaining elements if the filter predicate
3719-
// has not yet panicked. We'll backshift any remaining elements
3720-
// whether we've already panicked or if the consumption here panics.
3721-
if !backshift.drain.panic_flag {
3722-
backshift.drain.for_each(drop);
3723-
}
3724-
}
3725-
}

0 commit comments

Comments
 (0)