Fix undefined behavior in `Vec::truncate()`

1 files changed, 18 insertions, 7 deletions

diff --git a/src/vec.rs b/src/vec.rs
index f2866c0c..92e27c4d 100644
--- a/src/vec.rs
+++ b/src/vec.rs
@@ -263,13 +263,24 @@ impl<T, const N: usize> Vec<T, N> {
 
     /// Shortens the vector, keeping the first `len` elements and dropping the rest.
     pub fn truncate(&mut self, len: usize) {
-        // drop any extra elements
-        while len < self.len {
-            // decrement len before the drop_in_place(), so a panic on Drop
-            // doesn't re-drop the just-failed value.
-            self.len -= 1;
-            let len = self.len;
-            unsafe { ptr::drop_in_place(self.as_mut_slice().get_unchecked_mut(len)) };
+        // This is safe because:
+        //
+        // * 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).
+        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);
+            self.len = len;
+            ptr::drop_in_place(s);
         }
     }