From ac2c21a62338f916cac8b58de808013e3eb7a9f0 Mon Sep 17 00:00:00 2001 From: Paolo Barbolini Date: Sat, 11 Jun 2022 22:53:46 +0200 Subject: [PATCH 1/3] Add VecDeque::extend TrustedLen benchmark --- library/alloc/benches/vec_deque.rs | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/library/alloc/benches/vec_deque.rs b/library/alloc/benches/vec_deque.rs index 6660380e4beb0..7c78561ebf109 100644 --- a/library/alloc/benches/vec_deque.rs +++ b/library/alloc/benches/vec_deque.rs @@ -91,3 +91,35 @@ fn bench_extend_vec(b: &mut Bencher) { ring.extend(black_box(input)); }); } + +#[bench] +fn bench_extend_trustedlen(b: &mut Bencher) { + let mut ring: VecDeque = VecDeque::with_capacity(1000); + + b.iter(|| { + ring.clear(); + ring.extend(black_box(0..512)); + }); +} + +#[bench] +fn bench_extend_chained_trustedlen(b: &mut Bencher) { + let mut ring: VecDeque = VecDeque::with_capacity(1000); + + b.iter(|| { + ring.clear(); + ring.extend(black_box((0..256).chain(768..1024))); + }); +} + +#[bench] +fn bench_extend_chained_bytes(b: &mut Bencher) { + let mut ring: VecDeque = VecDeque::with_capacity(1000); + let input1: &[u16] = &[128; 256]; + let input2: &[u16] = &[255; 256]; + + b.iter(|| { + ring.clear(); + ring.extend(black_box(input1.iter().chain(input2.iter()))); + }); +} From bc3fae4dc1891de19368622d7b77ba3006b8f4ea Mon Sep 17 00:00:00 2001 From: Paolo Barbolini Date: Sun, 12 Jun 2022 00:54:20 +0200 Subject: [PATCH 2/3] Add VecDeque::extend from TrustedLen specialization --- .../alloc/src/collections/vec_deque/mod.rs | 19 ++++ .../src/collections/vec_deque/spec_extend.rs | 59 +++++++++++ .../alloc/src/collections/vec_deque/tests.rs | 97 +++++++++++++++++++ 3 files changed, 175 insertions(+) diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index e28a94386c7cd..f92e5759b6f9e 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -453,6 +453,25 @@ impl VecDeque { } } + /// Writes all values from `iter` to `dst`. + /// + /// # Safety + /// + /// Assumes no wrapping around happens. + /// Assumes capacity is sufficient. + #[inline] + unsafe fn write_iter( + &mut self, + dst: usize, + iter: impl Iterator, + written: &mut usize, + ) { + iter.enumerate().for_each(|(i, element)| unsafe { + self.buffer_write(dst + i, element); + *written += 1; + }); + } + /// Frobs the head and tail sections around to handle the fact that we /// just reallocated. Unsafe because it trusts old_capacity. #[inline] diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs index 172f2e9068fe6..cea2b7543ce44 100644 --- a/library/alloc/src/collections/vec_deque/spec_extend.rs +++ b/library/alloc/src/collections/vec_deque/spec_extend.rs @@ -1,5 +1,6 @@ use crate::alloc::Allocator; use crate::vec; +use core::iter::TrustedLen; use core::slice; use super::VecDeque; @@ -34,6 +35,64 @@ where } } +impl SpecExtend for VecDeque +where + I: TrustedLen, +{ + default fn spec_extend(&mut self, mut iter: I) { + // This is the case for a TrustedLen iterator. + let (low, high) = iter.size_hint(); + if let Some(additional) = high { + debug_assert_eq!( + low, + additional, + "TrustedLen iterator's size hint is not exact: {:?}", + (low, high) + ); + self.reserve(additional); + + struct WrapAddOnDrop<'a, T, A: Allocator> { + vec_deque: &'a mut VecDeque, + written: usize, + } + + impl<'a, T, A: Allocator> Drop for WrapAddOnDrop<'a, T, A> { + fn drop(&mut self) { + self.vec_deque.head = + self.vec_deque.wrap_add(self.vec_deque.head, self.written); + } + } + + let mut wrapper = WrapAddOnDrop { vec_deque: self, written: 0 }; + + let head_room = wrapper.vec_deque.cap() - wrapper.vec_deque.head; + unsafe { + wrapper.vec_deque.write_iter( + wrapper.vec_deque.head, + ByRefSized(&mut iter).take(head_room), + &mut wrapper.written, + ); + + if additional > head_room { + wrapper.vec_deque.write_iter(0, iter, &mut wrapper.written); + } + } + + debug_assert_eq!( + additional, wrapper.written, + "The number of items written to VecDeque doesn't match the TrustedLen size hint" + ); + } else { + // Per TrustedLen contract a `None` upper bound means that the iterator length + // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway. + // Since the other branch already panics eagerly (via `reserve()`) we do the same here. + // This avoids additional codegen for a fallback code path which would eventually + // panic anyway. + panic!("capacity overflow"); + } + } +} + impl SpecExtend> for VecDeque { fn spec_extend(&mut self, mut iterator: vec::IntoIter) { let slice = iterator.as_slice(); diff --git a/library/alloc/src/collections/vec_deque/tests.rs b/library/alloc/src/collections/vec_deque/tests.rs index f2869a54713d2..1f2daef213c3b 100644 --- a/library/alloc/src/collections/vec_deque/tests.rs +++ b/library/alloc/src/collections/vec_deque/tests.rs @@ -1,3 +1,5 @@ +use core::iter::TrustedLen; + use super::*; #[bench] @@ -791,6 +793,101 @@ fn test_from_vec() { assert_eq!(vd.len(), vec.len()); } +#[test] +fn test_extend_basic() { + test_extend_impl(false); +} + +#[test] +fn test_extend_trusted_len() { + test_extend_impl(true); +} + +fn test_extend_impl(trusted_len: bool) { + struct VecDequeTester { + test: VecDeque, + expected: VecDeque, + trusted_len: bool, + } + + impl VecDequeTester { + fn new(trusted_len: bool) -> Self { + Self { test: VecDeque::new(), expected: VecDeque::new(), trusted_len } + } + + fn test_extend(&mut self, iter: I) + where + I: Iterator + TrustedLen + Clone, + { + struct BasicIterator(I); + impl Iterator for BasicIterator + where + I: Iterator, + { + type Item = usize; + + fn next(&mut self) -> Option { + self.0.next() + } + } + + if self.trusted_len { + self.test.extend(iter.clone()); + } else { + self.test.extend(BasicIterator(iter.clone())); + } + + for item in iter { + self.expected.push_back(item) + } + + assert_eq!(self.test, self.expected); + let (a1, b1) = self.test.as_slices(); + let (a2, b2) = self.expected.as_slices(); + assert_eq!(a1, a2); + assert_eq!(b1, b2); + } + + fn drain + Clone>(&mut self, range: R) { + self.test.drain(range.clone()); + self.expected.drain(range); + + assert_eq!(self.test, self.expected); + } + + fn clear(&mut self) { + self.test.clear(); + self.expected.clear(); + } + + fn remaining_capacity(&self) -> usize { + self.test.capacity() - self.test.len() + } + } + + let mut tester = VecDequeTester::new(trusted_len); + + // Initial capacity + tester.test_extend(0..tester.remaining_capacity() - 1); + + // Grow + tester.test_extend(1024..2048); + + // Wrap around + tester.drain(..128); + + tester.test_extend(0..tester.remaining_capacity() - 1); + + // Continue + tester.drain(256..); + tester.test_extend(4096..8196); + + tester.clear(); + + // Start again + tester.test_extend(0..32); +} + #[test] #[should_panic = "capacity overflow"] fn test_from_vec_zst_overflow() { From ce3b6f505ed8c0a285461f4f9b729b62e7703354 Mon Sep 17 00:00:00 2001 From: Paolo Barbolini Date: Sun, 12 Jun 2022 00:55:09 +0200 Subject: [PATCH 3/3] Expose iter::ByRefSized as unstable feature and use it --- library/alloc/src/collections/vec_deque/spec_extend.rs | 2 +- library/alloc/src/lib.rs | 1 + library/core/src/iter/adapters/by_ref_sized.rs | 6 +++++- library/core/src/iter/adapters/mod.rs | 3 ++- library/core/src/iter/mod.rs | 4 +++- 5 files changed, 12 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/collections/vec_deque/spec_extend.rs b/library/alloc/src/collections/vec_deque/spec_extend.rs index cea2b7543ce44..97ff8b7652403 100644 --- a/library/alloc/src/collections/vec_deque/spec_extend.rs +++ b/library/alloc/src/collections/vec_deque/spec_extend.rs @@ -1,6 +1,6 @@ use crate::alloc::Allocator; use crate::vec; -use core::iter::TrustedLen; +use core::iter::{ByRefSized, TrustedLen}; use core::slice; use super::VecDeque; diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index baa1106a0dd7b..c08caa7b93ed2 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -143,6 +143,7 @@ #![feature(unchecked_math)] #![feature(unicode_internals)] #![feature(unsize)] +#![feature(std_internals)] // // Language features: #![feature(allocator_internals)] diff --git a/library/core/src/iter/adapters/by_ref_sized.rs b/library/core/src/iter/adapters/by_ref_sized.rs index bf2e3e182e2a3..cc1e8e8a27fa6 100644 --- a/library/core/src/iter/adapters/by_ref_sized.rs +++ b/library/core/src/iter/adapters/by_ref_sized.rs @@ -4,8 +4,11 @@ use crate::ops::Try; /// /// Ideally this will no longer be required, eventually, but as can be seen in /// the benchmarks (as of Feb 2022 at least) `by_ref` can have performance cost. -pub(crate) struct ByRefSized<'a, I>(pub &'a mut I); +#[unstable(feature = "std_internals", issue = "none")] +#[derive(Debug)] +pub struct ByRefSized<'a, I>(pub &'a mut I); +#[unstable(feature = "std_internals", issue = "none")] impl Iterator for ByRefSized<'_, I> { type Item = I::Item; @@ -47,6 +50,7 @@ impl Iterator for ByRefSized<'_, I> { } } +#[unstable(feature = "std_internals", issue = "none")] impl DoubleEndedIterator for ByRefSized<'_, I> { #[inline] fn next_back(&mut self) -> Option { diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index 4500b44b7e9af..916a26e242466 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/library/core/src/iter/adapters/mod.rs @@ -32,7 +32,8 @@ pub use self::{ scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip, }; -pub(crate) use self::by_ref_sized::ByRefSized; +#[unstable(feature = "std_internals", issue = "none")] +pub use self::by_ref_sized::ByRefSized; #[stable(feature = "iter_cloned", since = "1.1.0")] pub use self::cloned::Cloned; diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index 15362d2330acc..d5c6aed5b6c8a 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -398,6 +398,8 @@ pub use self::traits::{ #[stable(feature = "iter_zip", since = "1.59.0")] pub use self::adapters::zip; +#[unstable(feature = "std_internals", issue = "none")] +pub use self::adapters::ByRefSized; #[stable(feature = "iter_cloned", since = "1.1.0")] pub use self::adapters::Cloned; #[stable(feature = "iter_copied", since = "1.36.0")] @@ -422,7 +424,7 @@ pub use self::adapters::{ #[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] pub use self::adapters::{Intersperse, IntersperseWith}; -pub(crate) use self::adapters::{try_process, ByRefSized}; +pub(crate) use self::adapters::try_process; mod adapters; mod range;