From c4a1d5d9ae34df4f52672c7d057a7b040ed19840 Mon Sep 17 00:00:00 2001 From: The8472 Date: Fri, 16 Apr 2021 21:08:34 +0200 Subject: [PATCH 1/3] extract code path shared between FromIterator and Extend Previously FromIterator specialization called with_capacity() and then delegated to Extend which called reserve(). The reserve() is only needed when extending. This should reduce the amount of generated IR. --- library/alloc/src/vec/mod.rs | 27 ++++++++++++++++++- library/alloc/src/vec/spec_extend.rs | 16 +++-------- .../alloc/src/vec/spec_from_iter_nested.rs | 7 +++-- 3 files changed, 35 insertions(+), 15 deletions(-) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 20a16869cb3f8..88cf2305c3014 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -62,7 +62,7 @@ use core::hash::{Hash, Hasher}; use core::intrinsics::{arith_offset, assume}; use core::iter; #[cfg(not(no_global_oom_handling))] -use core::iter::FromIterator; +use core::iter::{FromIterator, TrustedLen}; use core::marker::PhantomData; use core::mem::{self, ManuallyDrop, MaybeUninit}; use core::ops::{self, Index, IndexMut, Range, RangeBounds}; @@ -2627,6 +2627,31 @@ impl Vec { } } + /// Appends the iterator's items to the vec without allocating. + /// + /// # Safety + /// + /// The caller must ensure that `self` has sufficient spare capacity + /// to hold the items returned by the iterator. + /// + /// The bound `I: TrustedLen` ensures that the caller can safely know + /// how much needs to be allocated. + #[cfg(not(no_global_oom_handling))] + unsafe fn extend_prealloc_trustedlen>(&mut self, iterator: I) { + unsafe { + let mut ptr = self.as_mut_ptr().add(self.len()); + let mut local_len = SetLenOnDrop::new(&mut self.len); + iterator.for_each(move |element| { + ptr::write(ptr, element); + ptr = ptr.offset(1); + // Since the loop executes user code which can panic we have to bump the length + // after each step. + // NB can't overflow since we would have had to alloc the address space + local_len.increment_len(1); + }); + } + } + /// Creates a splicing iterator that replaces the specified range in the vector /// with the given `replace_with` iterator and yields the removed items. /// `replace_with` does not need to be the same length as `range`. diff --git a/library/alloc/src/vec/spec_extend.rs b/library/alloc/src/vec/spec_extend.rs index c3b4534096de5..7b6bc1137a5a9 100644 --- a/library/alloc/src/vec/spec_extend.rs +++ b/library/alloc/src/vec/spec_extend.rs @@ -1,9 +1,8 @@ use crate::alloc::Allocator; use core::iter::TrustedLen; -use core::ptr::{self}; use core::slice::{self}; -use super::{IntoIter, SetLenOnDrop, Vec}; +use super::{IntoIter, Vec}; // Specialization trait used for Vec::extend pub(super) trait SpecExtend { @@ -34,17 +33,10 @@ where (low, high) ); self.reserve(additional); + // Safety: We rely on the TrustedLen contract to know how much capacity needs to be + // reserved. And we reserved at least that amount above. unsafe { - let mut ptr = self.as_mut_ptr().add(self.len()); - let mut local_len = SetLenOnDrop::new(&mut self.len); - iterator.for_each(move |element| { - ptr::write(ptr, element); - ptr = ptr.offset(1); - // Since the loop executes user code which can panic we have to bump the pointer - // after each step. - // NB can't overflow since we would have had to alloc the address space - local_len.increment_len(1); - }); + self.extend_prealloc_trustedlen(iterator); } } else { // Per TrustedLen contract a `None` upper bound means that the iterator length diff --git a/library/alloc/src/vec/spec_from_iter_nested.rs b/library/alloc/src/vec/spec_from_iter_nested.rs index 948cf044197c2..6019ea89ee961 100644 --- a/library/alloc/src/vec/spec_from_iter_nested.rs +++ b/library/alloc/src/vec/spec_from_iter_nested.rs @@ -52,8 +52,11 @@ where // (via `with_capacity`) we do the same here. _ => panic!("capacity overflow"), }; - // reuse extend specialization for TrustedLen - vector.spec_extend(iterator); + // Safety: The TrustedLen contract together with the `with_capacity` + // above guarantee that no further allocations should be needed to collect the iterator + unsafe { + vector.extend_prealloc_trustedlen(iterator); + } vector } } From dab430b1f28cc52e9cb540cec46728fb05aef54d Mon Sep 17 00:00:00 2001 From: the8472 Date: Sat, 17 Apr 2021 01:08:30 +0200 Subject: [PATCH 2/3] from review: improve method name Co-authored-by: bjorn3 --- library/alloc/src/vec/mod.rs | 2 +- library/alloc/src/vec/spec_extend.rs | 2 +- library/alloc/src/vec/spec_from_iter_nested.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 88cf2305c3014..a9b28746735f1 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2637,7 +2637,7 @@ impl Vec { /// The bound `I: TrustedLen` ensures that the caller can safely know /// how much needs to be allocated. #[cfg(not(no_global_oom_handling))] - unsafe fn extend_prealloc_trustedlen>(&mut self, iterator: I) { + unsafe fn extend_prealloc_trusted_len>(&mut self, iterator: I) { unsafe { let mut ptr = self.as_mut_ptr().add(self.len()); let mut local_len = SetLenOnDrop::new(&mut self.len); diff --git a/library/alloc/src/vec/spec_extend.rs b/library/alloc/src/vec/spec_extend.rs index 7b6bc1137a5a9..69d2ccef8b9db 100644 --- a/library/alloc/src/vec/spec_extend.rs +++ b/library/alloc/src/vec/spec_extend.rs @@ -36,7 +36,7 @@ where // Safety: We rely on the TrustedLen contract to know how much capacity needs to be // reserved. And we reserved at least that amount above. unsafe { - self.extend_prealloc_trustedlen(iterator); + self.extend_prealloc_trusted_len(iterator); } } else { // Per TrustedLen contract a `None` upper bound means that the iterator length diff --git a/library/alloc/src/vec/spec_from_iter_nested.rs b/library/alloc/src/vec/spec_from_iter_nested.rs index 6019ea89ee961..33273c134d103 100644 --- a/library/alloc/src/vec/spec_from_iter_nested.rs +++ b/library/alloc/src/vec/spec_from_iter_nested.rs @@ -55,7 +55,7 @@ where // Safety: The TrustedLen contract together with the `with_capacity` // above guarantee that no further allocations should be needed to collect the iterator unsafe { - vector.extend_prealloc_trustedlen(iterator); + vector.extend_prealloc_trusted_len(iterator); } vector } From 7f904fb41ed68f43c3338805c3e87991f38d678b Mon Sep 17 00:00:00 2001 From: The8472 Date: Tue, 19 Oct 2021 19:39:28 +0200 Subject: [PATCH 3/3] inline extracted function --- library/alloc/src/vec/mod.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index a9b28746735f1..d22156636eb1d 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2637,6 +2637,7 @@ impl Vec { /// The bound `I: TrustedLen` ensures that the caller can safely know /// how much needs to be allocated. #[cfg(not(no_global_oom_handling))] + #[inline] unsafe fn extend_prealloc_trusted_len>(&mut self, iterator: I) { unsafe { let mut ptr = self.as_mut_ptr().add(self.len());