Make vec::IntoIter::drop and VecInner::truncate consistent.

Author: reitermarkus

src/vec/mod.rs

@@ -735,24 +735,24 @@ impl<T, LenT: LenType, S: VecStorage<T> + ?Sized> VecInner<T, LenT, S> {
 
     /// Shortens the vector, keeping the first `len` elements and dropping the rest.
     pub fn truncate(&mut self, len: usize) {
-        // This is safe because:
+        // NOTE: It is intentional that this is `>`. Changing it to `>=` may have negative
+        //       performance implications. See rust-lang/rust#78884 for more details.
+        if len > self.len() {
+            return;
+        }
+
+        // SAFETY: `self.buffer` contains initialized elements in the range `len..(len + remaining_len)`,
+        //         so it is safe to create a mutable slice of them and drop them.
         //
-        // * the slice passed to `drop_in_place` is valid; the `len > self.len` case avoids creating
-        //   an invalid slice, and
-        // * the `len` of the vector is shrunk before calling `drop_in_place`, such that no value
-        //   will be dropped twice in case `drop_in_place` were to panic once (if it panics twice,
-        //   the program aborts).
+        //         `self.len` is set before dropping the excess elements so that a panicking
+        //         `Drop` implementation does not result in inconsistent state during
+        //         unwinding that could lead to undefined behaviour.
         unsafe {
-            // Note: It's intentional that this is `>` and not `>=`.
-            //       Changing it to `>=` has negative performance
-            //       implications in some cases. See rust-lang/rust#78884 for more.
-            if len > self.len() {
-                return;
-            }
             let remaining_len = self.len() - len;
-            let s = ptr::slice_from_raw_parts_mut(self.as_mut_ptr().add(len), remaining_len);
+            let remaining =
+                ptr::slice_from_raw_parts_mut(self.as_mut_ptr().add(len), remaining_len);
             self.len = LenT::from_usize(len);
-            ptr::drop_in_place(s);
+            ptr::drop_in_place(remaining);
         }
     }
 
@@ -1572,13 +1572,21 @@ where
 impl<T, LenT: LenType, const N: usize> Drop for IntoIter<T, N, LenT> {
     fn drop(&mut self) {
         let len = self.vec.len;
-        self.vec.len = LenT::ZERO;
+
+        // SAFETY: `self.vec` contains initialized elements in the range `self.next..len`,
+        //         so it is safe to create a mutable slice of them and drop them.
+        //
+        //         `self.vec.len` is set to `0` before dropping the elements so that
+        //         a panicking `Drop` implementation does not result in inconsistent
+        //         state during unwinding that could lead to undefined behaviour.
         unsafe {
-            // Drop all the elements that have not been moved out of vec
+            // Drop all elements that have not been consumed yet.
+            let remaining_len = len - self.next;
             let remaining = slice::from_raw_parts_mut(
                 self.vec.as_mut_ptr().add(self.next.into_usize()),
-                (len - self.next).into_usize(),
+                remaining_len.into_usize(),
             );
+            self.vec.len = LenT::ZERO;
             ptr::drop_in_place(remaining);
         }
     }
@@ -1810,7 +1818,7 @@ where
 mod tests {
     use core::fmt::Write;
     use std::{
-        panic::catch_unwind,
+        panic::{catch_unwind, AssertUnwindSafe},
         sync::atomic::{AtomicI32, Ordering::Relaxed},
     };
 
@@ -2393,40 +2401,57 @@ mod tests {
         }
     }
 
+    #[derive(Debug)]
+    struct Dropper<'a>(&'a AtomicI32);
+
+    impl<'a> Dropper<'a> {
+        fn new(count: &'a AtomicI32) -> Self {
+            count.fetch_add(1, Relaxed);
+            Self(count)
+        }
+    }
+    impl Drop for Dropper<'_> {
+        fn drop(&mut self) {
+            self.0.fetch_sub(1, Relaxed);
+            panic!("Testing  panicking");
+        }
+    }
+
+    // This tests that a panic in a downstream drop implementation does not lead to an
+    // inconsistent state during unwinding that could lead to undefined behaviour.
+    // For more info see https://github.com/rust-embedded/heapless/issues/659.
     #[test]
-    fn test_use_after_free_intoiter_drop() {
-        // This tests that a panic in a downstream drop implementation does not lead to an
-        // inconsistent state during unwinding that could lead to undefined behaviour.
-        // For more info see https://github.com/rust-embedded/heapless/issues/659
+    fn test_use_after_free_clear() {
+        let count = AtomicI32::new(0);
 
-        static COUNT: AtomicI32 = AtomicI32::new(0);
+        let mut vec = Vec::<Dropper, 5>::new();
+        vec.push(Dropper::new(&count)).unwrap();
 
-        #[derive(Debug)]
-        #[allow(unused)]
-        struct Dropper();
+        catch_unwind(AssertUnwindSafe(|| {
+            vec.clear();
+        }))
+        .unwrap_err();
 
-        impl Dropper {
-            fn new() -> Self {
-                COUNT.fetch_add(1, Relaxed);
-                Self()
-            }
-            fn count() -> i32 {
-                COUNT.load(Relaxed)
-            }
-        }
-        impl Drop for Dropper {
-            fn drop(&mut self) {
-                COUNT.fetch_sub(1, Relaxed);
-                panic!("Testing  panicking");
-            }
-        }
+        assert_eq!(count.load(Relaxed), 0);
+        assert_eq!(vec.len(), 0);
+    }
+
+    // This tests that a panic in a downstream drop implementation does not lead to an
+    // inconsistent state during unwinding that could lead to undefined behaviour.
+    // For more info see https://github.com/rust-embedded/heapless/issues/659.
+    #[test]
+    fn test_use_after_free_intoiter_drop() {
+        let count = AtomicI32::new(0);
 
         let mut vec = Vec::<Dropper, 5>::new();
-        vec.push(Dropper::new()).unwrap();
-        let iterator = vec.into_iter();
+        vec.push(Dropper::new(&count)).unwrap();
+
+        catch_unwind(move || {
+            drop(vec.into_iter());
+        })
+        .unwrap_err();
 
-        catch_unwind(move || drop(iterator)).unwrap_err();
-        assert_eq!(Dropper::count(), 0);
+        assert_eq!(count.load(Relaxed), 0);
     }
 
     fn _test_variance<'a: 'b, 'b>(x: Vec<&'a (), 42>) -> Vec<&'b (), 42> {