[beta] backports

library/alloc/src/str.rs

@@ -90,8 +90,8 @@ impl<S: Borrow<str>> Join<&str> for [S] {
     }
 }
 
-macro_rules! spezialize_for_lengths {
-    ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => {
+macro_rules! specialize_for_lengths {
+    ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => {{
         let mut target = $target;
         let iter = $iter;
         let sep_bytes = $separator;
@@ -102,19 +102,22 @@ macro_rules! spezialize_for_lengths {
                 $num => {
                     for s in iter {
                         copy_slice_and_advance!(target, sep_bytes);
-                        copy_slice_and_advance!(target, s.borrow().as_ref());
+                        let content_bytes = s.borrow().as_ref();
+                        copy_slice_and_advance!(target, content_bytes);
                     }
                 },
             )*
             _ => {
                 // arbitrary non-zero size fallback
                 for s in iter {
                     copy_slice_and_advance!(target, sep_bytes);
-                    copy_slice_and_advance!(target, s.borrow().as_ref());
+                    let content_bytes = s.borrow().as_ref();
+                    copy_slice_and_advance!(target, content_bytes);
                 }
             }
         }
-    };
+        target
+    }}
 }
 
 macro_rules! copy_slice_and_advance {
@@ -153,30 +156,33 @@ where
     // if the `len` calculation overflows, we'll panic
     // we would have run out of memory anyway and the rest of the function requires
     // the entire Vec pre-allocated for safety
-    let len = sep_len
+    let reserved_len = sep_len
         .checked_mul(iter.len())
         .and_then(|n| {
             slice.iter().map(|s| s.borrow().as_ref().len()).try_fold(n, usize::checked_add)
         })
         .expect("attempt to join into collection with len > usize::MAX");
 
-    // crucial for safety
-    let mut result = Vec::with_capacity(len);
-    assert!(result.capacity() >= len);
+    // prepare an uninitialized buffer
+    let mut result = Vec::with_capacity(reserved_len);
+    debug_assert!(result.capacity() >= reserved_len);
 
     result.extend_from_slice(first.borrow().as_ref());
 
     unsafe {
-        {
-            let pos = result.len();
-            let target = result.get_unchecked_mut(pos..len);
-
-            // copy separator and slices over without bounds checks
-            // generate loops with hardcoded offsets for small separators
-            // massive improvements possible (~ x2)
-            spezialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4);
-        }
-        result.set_len(len);
+        let pos = result.len();
+        let target = result.get_unchecked_mut(pos..reserved_len);
+
+        // copy separator and slices over without bounds checks
+        // generate loops with hardcoded offsets for small separators
+        // massive improvements possible (~ x2)
+        let remain = specialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4);
+
+        // A weird borrow implementation may return different
+        // slices for the length calculation and the actual copy.
+        // Make sure we don't expose uninitialized bytes to the caller.
+        let result_len = reserved_len - remain.len();
+        result.set_len(result_len);
     }
     result
 }

library/alloc/src/vec/into_iter.rs

@@ -85,20 +85,29 @@ impl<T, A: Allocator> IntoIter<T, A> {
         ptr::slice_from_raw_parts_mut(self.ptr as *mut T, self.len())
     }
 
-    pub(super) fn drop_remaining(&mut self) {
-        unsafe {
-            ptr::drop_in_place(self.as_mut_slice());
-        }
-        self.ptr = self.end;
-    }
+    /// Drops remaining elements and relinquishes the backing allocation.
+    ///
+    /// This is roughly equivalent to the following, but more efficient
+    ///
+    /// ```
+    /// # let mut into_iter = Vec::<u8>::with_capacity(10).into_iter();
+    /// (&mut into_iter).for_each(core::mem::drop);
+    /// unsafe { core::ptr::write(&mut into_iter, Vec::new().into_iter()); }
+    /// ```
+    pub(super) fn forget_allocation_drop_remaining(&mut self) {
+        let remaining = self.as_raw_mut_slice();
 
-    /// Relinquishes the backing allocation, equivalent to
-    /// `ptr::write(&mut self, Vec::new().into_iter())`
-    pub(super) fn forget_allocation(&mut self) {
+        // overwrite the individual fields instead of creating a new
+        // struct and then overwriting &mut self.
+        // this creates less assembly
         self.cap = 0;
         self.buf = unsafe { NonNull::new_unchecked(RawVec::NEW.ptr()) };
         self.ptr = self.buf.as_ptr();
         self.end = self.buf.as_ptr();
+
+        unsafe {
+            ptr::drop_in_place(remaining);
+        }
     }
 }
 

library/alloc/src/vec/source_iter_marker.rs

@@ -78,9 +78,9 @@ where
         }
 
         // drop any remaining values at the tail of the source
-        src.drop_remaining();
         // but prevent drop of the allocation itself once IntoIter goes out of scope
-        src.forget_allocation();
+        // if the drop panics then we also leak any elements collected into dst_buf
+        src.forget_allocation_drop_remaining();
 
         let vec = unsafe {
             let len = dst.offset_from(dst_buf) as usize;

library/alloc/tests/str.rs

@@ -160,6 +160,36 @@ fn test_join_for_different_lengths_with_long_separator() {
     test_join!("～～～～～a～～～～～bc", ["", "a", "bc"], "～～～～～");
 }
 
+#[test]
+fn test_join_isue_80335() {
+    use core::{borrow::Borrow, cell::Cell};
+
+    struct WeirdBorrow {
+        state: Cell<bool>,
+    }
+
+    impl Default for WeirdBorrow {
+        fn default() -> Self {
+            WeirdBorrow { state: Cell::new(false) }
+        }
+    }
+
+    impl Borrow<str> for WeirdBorrow {
+        fn borrow(&self) -> &str {
+            let state = self.state.get();
+            if state {
+                "0"
+            } else {
+                self.state.set(true);
+                "123456"
+            }
+        }
+    }
+
+    let arr: [WeirdBorrow; 3] = Default::default();
+    test_join!("0-0-0", arr, "-");
+}
+
 #[test]
 #[cfg_attr(miri, ignore)] // Miri is too slow
 fn test_unsafe_slice() {

library/alloc/tests/vec.rs

@@ -1027,7 +1027,7 @@ fn test_from_iter_specialization_head_tail_drop() {
 }
 
 #[test]
-fn test_from_iter_specialization_panic_drop() {
+fn test_from_iter_specialization_panic_during_iteration_drops() {
     let drop_count: Vec<_> = (0..=2).map(|_| Rc::new(())).collect();
     let src: Vec<_> = drop_count.iter().cloned().collect();
     let iter = src.into_iter();
@@ -1050,6 +1050,42 @@ fn test_from_iter_specialization_panic_drop() {
     );
 }
 
+#[test]
+fn test_from_iter_specialization_panic_during_drop_leaks() {
+    static mut DROP_COUNTER: usize = 0;
+
+    #[derive(Debug)]
+    enum Droppable {
+        DroppedTwice(Box<i32>),
+        PanicOnDrop,
+    }
+
+    impl Drop for Droppable {
+        fn drop(&mut self) {
+            match self {
+                Droppable::DroppedTwice(_) => {
+                    unsafe {
+                        DROP_COUNTER += 1;
+                    }
+                    println!("Dropping!")
+                }
+                Droppable::PanicOnDrop => {
+                    if !std::thread::panicking() {
+                        panic!();
+                    }
+                }
+            }
+        }
+    }
+
+    let _ = std::panic::catch_unwind(AssertUnwindSafe(|| {
+        let v = vec![Droppable::DroppedTwice(Box::new(123)), Droppable::PanicOnDrop];
+        let _ = v.into_iter().take(0).collect::<Vec<_>>();
+    }));
+
+    assert_eq!(unsafe { DROP_COUNTER }, 1);
+}
+
 #[test]
 fn test_cow_from() {
     let borrowed: &[_] = &["borrowed", "(slice)"];