From 7f68c73819371165fedde09afe7273752c9948ad Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Sun, 27 Mar 2022 22:02:06 +0400 Subject: [PATCH 1/3] Add vec::Drain{,Filter}::keep_rest These methods allow to cancel draining of unyielded elements. --- alloc/src/vec/drain.rs | 55 +++++++++++++++++++++++++++++++- alloc/src/vec/drain_filter.rs | 41 ++++++++++++++++++++++-- alloc/tests/lib.rs | 1 + alloc/tests/vec.rs | 59 +++++++++++++++++++++++++++++++++++ 4 files changed, 153 insertions(+), 3 deletions(-) diff --git a/alloc/src/vec/drain.rs b/alloc/src/vec/drain.rs index 5cdee0bd4..3e350c2b3 100644 --- a/alloc/src/vec/drain.rs +++ b/alloc/src/vec/drain.rs @@ -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,59 @@ 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`. + #[unstable(feature = "drain_keep_rest", issue = "none")] + 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::() != 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")] diff --git a/alloc/src/vec/drain_filter.rs b/alloc/src/vec/drain_filter.rs index 3c37c92ae..a5c1eab94 100644 --- a/alloc/src/vec/drain_filter.rs +++ b/alloc/src/vec/drain_filter.rs @@ -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,42 @@ where pub fn allocator(&self) -> &A { self.vec.allocator() } + + /// Keep unyielded elements in the source `Vec`. + #[unstable(feature = "drain_keep_rest", issue = "none")] + 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::() != 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")] diff --git a/alloc/tests/lib.rs b/alloc/tests/lib.rs index ffc7944ec..455bd9236 100644 --- a/alloc/tests/lib.rs +++ b/alloc/tests/lib.rs @@ -45,6 +45,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}; diff --git a/alloc/tests/vec.rs b/alloc/tests/vec.rs index cc768c73c..946380162 100644 --- a/alloc/tests/vec.rs +++ b/alloc/tests/vec.rs @@ -793,6 +793,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]; @@ -1478,6 +1508,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 From 52959b4b4602b1db211b0ec38d92eb53252c3e9b Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Sun, 28 Aug 2022 16:38:26 +0400 Subject: [PATCH 2/3] add examples to `vec::Drain{,Filter}::keep_rest` docs --- alloc/src/vec/drain.rs | 18 ++++++++++++++++++ alloc/src/vec/drain_filter.rs | 19 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/alloc/src/vec/drain.rs b/alloc/src/vec/drain.rs index 3e350c2b3..2a3f7ff6d 100644 --- a/alloc/src/vec/drain.rs +++ b/alloc/src/vec/drain.rs @@ -67,6 +67,24 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> { } /// 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 = "none")] pub fn keep_rest(self) { // At this moment layout looks like this: diff --git a/alloc/src/vec/drain_filter.rs b/alloc/src/vec/drain_filter.rs index a5c1eab94..3130c1f1c 100644 --- a/alloc/src/vec/drain_filter.rs +++ b/alloc/src/vec/drain_filter.rs @@ -57,6 +57,25 @@ where } /// 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 = "none")] pub fn keep_rest(self) { // At this moment layout looks like this: From 953ee505bbebe207d6adfd5fdbcdf88172c9cd68 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Sun, 28 Aug 2022 17:02:37 +0400 Subject: [PATCH 3/3] fill-in tracking issue for `feature(drain_keep_rest)` --- alloc/src/vec/drain.rs | 2 +- alloc/src/vec/drain_filter.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/alloc/src/vec/drain.rs b/alloc/src/vec/drain.rs index 2a3f7ff6d..5b73906a1 100644 --- a/alloc/src/vec/drain.rs +++ b/alloc/src/vec/drain.rs @@ -85,7 +85,7 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> { /// // `vec` would be empty. /// assert_eq!(vec, ['b', 'c']); /// ``` - #[unstable(feature = "drain_keep_rest", issue = "none")] + #[unstable(feature = "drain_keep_rest", issue = "101122")] pub fn keep_rest(self) { // At this moment layout looks like this: // diff --git a/alloc/src/vec/drain_filter.rs b/alloc/src/vec/drain_filter.rs index 3130c1f1c..8c03f1692 100644 --- a/alloc/src/vec/drain_filter.rs +++ b/alloc/src/vec/drain_filter.rs @@ -76,7 +76,7 @@ where /// // `vec` would be empty. /// assert_eq!(vec, ['b', 'c']); /// ``` - #[unstable(feature = "drain_keep_rest", issue = "none")] + #[unstable(feature = "drain_keep_rest", issue = "101122")] pub fn keep_rest(self) { // At this moment layout looks like this: //