extract code path shared between FromIterator and Extend

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<T, A: Allocator> Vec<T, A> {
         }
     }
 
+    /// 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<I: TrustedLen<Item = T>>(&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`.

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<T, I> {
@@ -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

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
     }
 }