diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 20a16869cb3f8..d22156636eb1d 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,32 @@ 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))] + #[inline] + 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); + 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..69d2ccef8b9db 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_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 948cf044197c2..33273c134d103 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_trusted_len(iterator); + } vector } }