squiid/rust
1
0
Fork 0
Code Issues Pull requests Projects Activity

Rollup merge of #95376 - WaffleLapkin:drain_keep_rest, r=dtolnay

Browse source 
Add `vec::Drain{,Filter}::keep_rest`

This PR adds `keep_rest` methods to `vec::Drain` and `vec::DrainFilter` under `drain_keep_rest` feature gate:
```rust
// mod alloc::vec

impl<T, A: Allocator> Drain<'_, T, A> {
    pub fn keep_rest(self);
}

impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
where
    F: FnMut(&mut T) -> bool,
{
    pub fn keep_rest(self);
}
```

Both these methods cancel draining of elements that were not yet yielded from the iterators. While this needs more testing & documentation, I want at least start the discussion. This may be a potential way out of the "should `DrainFilter` exhaust itself on drop?" argument.
This commit is contained in:
Dylan DPC 2022-08-30 11:26:47 +05:30 committed by GitHub
parent 9f4d5d2a28 7a433e4d00
commit c731157395
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 190 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
library/alloc/src/vec/drain.rs
View file

@ -1,7 +1,7 @@
use crate::alloc::{Allocator, Global};
use core::fmt;
use core::iter::{FusedIterator, TrustedLen};
use core::mem;
use core::mem::{self, ManuallyDrop};
use core::ptr::{self, NonNull};
use core::slice::{self};
@ -65,6 +65,77 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
pub fn allocator(&self) -> &A {
unsafe { self.vec.as_ref().allocator() }
}
/// Keep unyielded elements in the source `Vec`.
///
/// # Examples
///
/// ```
/// #![feature(drain_keep_rest)]
///
/// let mut vec = vec!['a', 'b', 'c'];
/// let mut drain = vec.drain(..);
///
/// assert_eq!(drain.next().unwrap(), 'a');
///
/// // This call keeps 'b' and 'c' in the vec.
/// drain.keep_rest();
///
/// // If we wouldn't call `keep_rest()`,
/// // `vec` would be empty.
/// assert_eq!(vec, ['b', 'c']);
/// ```
#[unstable(feature = "drain_keep_rest", issue = "101122")]
pub fn keep_rest(self) {
// At this moment layout looks like this:
//
// [head] [yielded by next] [unyielded] [yielded by next_back] [tail]
// ^-- start \_________/-- unyielded_len \____/-- self.tail_len
// ^-- unyielded_ptr ^-- tail
//
// Normally `Drop` impl would drop [unyielded] and then move [tail] to the `start`.
// Here we want to
// 1. Move [unyielded] to `start`
// 2. Move [tail] to a new start at `start + len(unyielded)`
// 3. Update length of the original vec to `len(head) + len(unyielded) + len(tail)`
// a. In case of ZST, this is the only thing we want to do
// 4. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
let mut this = ManuallyDrop::new(self);
unsafe {
let source_vec = this.vec.as_mut();
let start = source_vec.len();
let tail = this.tail_start;
let unyielded_len = this.iter.len();
let unyielded_ptr = this.iter.as_slice().as_ptr();
// ZSTs have no identity, so we don't need to move them around.
let needs_move = mem::size_of::<T>() != 0;
if needs_move {
let start_ptr = source_vec.as_mut_ptr().add(start);
// memmove back unyielded elements
if unyielded_ptr != start_ptr {
let src = unyielded_ptr;
let dst = start_ptr;
ptr::copy(src, dst, unyielded_len);
}
// memmove back untouched tail
if tail != (start + unyielded_len) {
let src = source_vec.as_ptr().add(tail);
let dst = start_ptr.add(unyielded_len);
ptr::copy(src, dst, this.tail_len);
}
}
source_vec.set_len(start + unyielded_len + this.tail_len);
}
}
}
#[stable(feature = "vec_drain_as_slice", since = "1.46.0")]

60
library/alloc/src/vec/drain_filter.rs
View file

@ -1,6 +1,7 @@
use crate::alloc::{Allocator, Global};
use core::ptr::{self};
use core::slice::{self};
use core::mem::{self, ManuallyDrop};
use core::ptr;
use core::slice;
use super::Vec;
@ -54,6 +55,61 @@ where
pub fn allocator(&self) -> &A {
self.vec.allocator()
}
/// Keep unyielded elements in the source `Vec`.
///
/// # Examples
///
/// ```
/// #![feature(drain_filter)]
/// #![feature(drain_keep_rest)]
///
/// let mut vec = vec!['a', 'b', 'c'];
/// let mut drain = vec.drain_filter(|_| true);
///
/// assert_eq!(drain.next().unwrap(), 'a');
///
/// // This call keeps 'b' and 'c' in the vec.
/// drain.keep_rest();
///
/// // If we wouldn't call `keep_rest()`,
/// // `vec` would be empty.
/// assert_eq!(vec, ['b', 'c']);
/// ```
#[unstable(feature = "drain_keep_rest", issue = "101122")]
pub fn keep_rest(self) {
// At this moment layout looks like this:
//
// _____________________/-- old_len
// / \
// [kept] [yielded] [tail]
// \_______/ ^-- idx
// \-- del
//
// Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
//
// 1. Move [tail] after [kept]
// 2. Update length of the original vec to `old_len - del`
// a. In case of ZST, this is the only thing we want to do
// 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
let mut this = ManuallyDrop::new(self);
unsafe {
// ZSTs have no identity, so we don't need to move them around.
let needs_move = mem::size_of::<T>() != 0;
if needs_move && this.idx < this.old_len && this.del > 0 {
let ptr = this.vec.as_mut_ptr();
let src = ptr.add(this.idx);
let dst = src.sub(this.del);
let tail_len = this.old_len - this.idx;
src.copy_to(dst, tail_len);
}
let new_len = this.old_len - this.del;
this.vec.set_len(new_len);
}
}
}
#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]

1
library/alloc/tests/lib.rs
View file

@ -44,6 +44,7 @@
#![feature(bench_black_box)]
#![feature(strict_provenance)]
#![feature(once_cell)]
#![feature(drain_keep_rest)]
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

59
library/alloc/tests/vec.rs
View file

@ -839,6 +839,36 @@ fn test_drain_leak() {
assert_eq!(v, vec![D(0, false), D(1, false), D(6, false),]);
}
#[test]
fn test_drain_keep_rest() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
let mut drain = v.drain(1..6);
assert_eq!(drain.next(), Some(1));
assert_eq!(drain.next_back(), Some(5));
assert_eq!(drain.next(), Some(2));
drain.keep_rest();
assert_eq!(v, &[0, 3, 4, 6]);
}
#[test]
fn test_drain_keep_rest_all() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
v.drain(1..6).keep_rest();
assert_eq!(v, &[0, 1, 2, 3, 4, 5, 6]);
}
#[test]
fn test_drain_keep_rest_none() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
let mut drain = v.drain(1..6);
drain.by_ref().for_each(drop);
drain.keep_rest();
assert_eq!(v, &[0, 6]);
}
#[test]
fn test_splice() {
let mut v = vec![1, 2, 3, 4, 5];
@ -1533,6 +1563,35 @@ fn drain_filter_unconsumed() {
assert_eq!(vec, [2, 4]);
}
#[test]
fn test_drain_filter_keep_rest() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
let mut drain = v.drain_filter(|&mut x| x % 2 == 0);
assert_eq!(drain.next(), Some(0));
assert_eq!(drain.next(), Some(2));
drain.keep_rest();
assert_eq!(v, &[1, 3, 4, 5, 6]);
}
#[test]
fn test_drain_filter_keep_rest_all() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
v.drain_filter(|_| true).keep_rest();
assert_eq!(v, &[0, 1, 2, 3, 4, 5, 6]);
}
#[test]
fn test_drain_filter_keep_rest_none() {
let mut v = vec![0, 1, 2, 3, 4, 5, 6];
let mut drain = v.drain_filter(|_| true);
drain.by_ref().for_each(drop);
drain.keep_rest();
assert_eq!(v, &[]);
}
#[test]
fn test_reserve_exact() {
// This is all the same as test_reserve