Tune RepeatWith::try_fold and Take::for_each and Vec::extend_trusted

[rust.git] / library / alloc / src / vec / mod.rs

diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 94ebcb863a421153b4d444108c0907bad220c65a..e147af2ce39c68f905601bfc9c2c66e1b8c1687f 100644 (file)
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -2874,13 +2874,12 @@ fn extend_trusted(&mut self, iterator: impl iter::TrustedLen<Item = T>) {
             );
             self.reserve(additional);
             unsafe {
-                let mut ptr = self.as_mut_ptr().add(self.len());
+                let ptr = self.as_mut_ptr();
                 let mut local_len = SetLenOnDrop::new(&mut self.len);
                 iterator.for_each(move |element| {
-                    ptr::write(ptr, element);
-                    ptr = ptr.add(1);
-                    // Since the loop executes user code which can panic we have to bump the pointer
-                    // after each step.
+                    ptr::write(ptr.add(local_len.current_len()), element);
+                    // Since the loop executes user code which can panic we have to update
+                    // the length every step to correctly drop what we've written.
                     // NB can't overflow since we would have had to alloc the address space
                     local_len.increment_len(1);
                 });