Fix leak on panic in insert_many.

Fixes #208.

Author: mbrubeck

src/lib.rs

@@ -98,7 +98,7 @@ use core::hint::unreachable_unchecked;
 use core::iter::{repeat, FromIterator, FusedIterator, IntoIterator};
 use core::mem;
 use core::mem::MaybeUninit;
-use core::ops::{self, RangeBounds};
+use core::ops::{self, Range, RangeBounds};
 use core::ptr::{self, NonNull};
 use core::slice::{self, SliceIndex};
 
@@ -975,8 +975,6 @@ impl<A: Array> SmallVec<A> {
 
     /// Insert multiple elements at position `index`, shifting all following elements toward the
     /// back.
-    ///
-    /// Note: when the iterator panics, this can leak memory.
     pub fn insert_many<I: IntoIterator<Item = A::Item>>(&mut self, index: usize, iterable: I) {
         let iter = iterable.into_iter();
         if index == self.len() {
@@ -991,27 +989,41 @@ impl<A: Array> SmallVec<A> {
         unsafe {
             let old_len = self.len();
             assert!(index <= old_len);
-            let mut ptr = self.as_mut_ptr().add(index);
+            let start = self.as_mut_ptr();
+            let mut ptr = start.add(index);
 
             // Move the trailing elements.
             ptr::copy(ptr, ptr.add(lower_size_bound), old_len - index);
 
             // In case the iterator panics, don't double-drop the items we just copied above.
-            self.set_len(index);
+            self.set_len(0);
+            let mut guard = DropOnPanic {
+                start,
+                skip: index..(index + lower_size_bound),
+                len: old_len + lower_size_bound,
+            };
 
             let mut num_added = 0;
             for element in iter {
                 let mut cur = ptr.add(num_added);
                 if num_added >= lower_size_bound {
                     // Iterator provided more elements than the hint.  Move trailing items again.
                     self.reserve(1);
-                    ptr = self.as_mut_ptr().add(index);
+                    let start = self.as_mut_ptr();
+                    ptr = start.add(index);
                     cur = ptr.add(num_added);
                     ptr::copy(cur, cur.add(1), old_len - index);
+
+                    guard.start = start;
+                    guard.len += 1;
+                    guard.skip.end += 1;
                 }
                 ptr::write(cur, element);
+                guard.skip.start += 1;
                 num_added += 1;
             }
+            mem::forget(guard);
+
             if num_added < lower_size_bound {
                 // Iterator provided fewer elements than the hint
                 ptr::copy(
@@ -1023,6 +1035,24 @@ impl<A: Array> SmallVec<A> {
 
             self.set_len(old_len + num_added);
         }
+
+        struct DropOnPanic<T> {
+            start: *mut T,
+            skip: Range<usize>,
+            len: usize,
+        }
+
+        impl<T> Drop for DropOnPanic<T> {
+            fn drop(&mut self) {
+                for i in 0..self.len {
+                    if !self.skip.contains(&i) {
+                        unsafe {
+                            ptr::drop_in_place(self.start.add(i));
+                        }
+                    }
+                }
+            }
+        }
     }
 
     /// Convert a SmallVec to a Vec, without reallocating if the SmallVec has already spilled onto
@@ -1249,6 +1279,22 @@ impl<A: Array> SmallVec<A> {
             data: SmallVecData::from_heap(ptr, length),
         }
     }
+
+    /// Returns a raw pointer to the vector's buffer.
+    pub fn as_ptr(&self) -> *const A::Item {
+        // We shadow the slice method of the same name to avoid going through
+        // `deref`, which creates an intermediate reference that may place
+        // additional safety constraints on the contents of the slice.
+        self.triple().0
+    }
+
+    /// Returns a raw mutable pointer to the vector's buffer.
+    pub fn as_mut_ptr(&mut self) -> *mut A::Item {
+        // We shadow the slice method of the same name to avoid going through
+        // `deref_mut`, which creates an intermediate reference that may place
+        // additional safety constraints on the contents of the slice.
+        self.triple_mut().0
+    }
 }
 
 impl<A: Array> SmallVec<A>

src/tests.rs

@@ -1,4 +1,4 @@
-use crate::{SmallVec, smallvec};
+use crate::{smallvec, SmallVec};
 
 use std::iter::FromIterator;
 
@@ -320,52 +320,99 @@ fn test_insert_many_long_hint() {
     );
 }
 
-#[test]
 // https://github.com/servo/rust-smallvec/issues/96
-fn test_insert_many_panic() {
+mod insert_many_panic {
+    use crate::{smallvec, SmallVec};
+    use alloc::boxed::Box;
+
     struct PanicOnDoubleDrop {
         dropped: Box<bool>,
     }
 
+    impl PanicOnDoubleDrop {
+        fn new() -> Self {
+            Self {
+                dropped: Box::new(false),
+            }
+        }
+    }
+
     impl Drop for PanicOnDoubleDrop {
         fn drop(&mut self) {
             assert!(!*self.dropped, "already dropped");
             *self.dropped = true;
         }
     }
 
-    struct BadIter;
+    /// Claims to yield `hint` items, but actually yields `count`, then panics.
+    struct BadIter {
+        hint: usize,
+        count: usize,
+    }
+
     impl Iterator for BadIter {
         type Item = PanicOnDoubleDrop;
         fn size_hint(&self) -> (usize, Option<usize>) {
-            (1, None)
+            (self.hint, None)
         }
         fn next(&mut self) -> Option<Self::Item> {
-            panic!()
+            if self.count == 0 {
+                panic!()
+            }
+            self.count -= 1;
+            Some(PanicOnDoubleDrop::new())
         }
     }
 
-    // These boxes are leaked on purpose by panicking `insert_many`,
-    // so we clean them up manually to appease Miri's leak checker.
-    let mut box1 = Box::new(false);
-    let mut box2 = Box::new(false);
+    #[test]
+    fn panic_early_at_start() {
+        let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> =
+            smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),];
+        let result = ::std::panic::catch_unwind(move || {
+            vec.insert_many(0, BadIter { hint: 1, count: 0 });
+        });
+        assert!(result.is_err());
+    }
 
-    let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = vec![
-        PanicOnDoubleDrop {
-            dropped: unsafe { Box::from_raw(&mut *box1) },
-        },
-        PanicOnDoubleDrop {
-            dropped: unsafe { Box::from_raw(&mut *box2) },
-        },
-    ]
-    .into();
-    let result = ::std::panic::catch_unwind(move || {
-        vec.insert_many(0, BadIter);
-    });
-    assert!(result.is_err());
-
-    drop(box1);
-    drop(box2);
+    #[test]
+    fn panic_early_in_middle() {
+        let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> =
+            smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),];
+        let result = ::std::panic::catch_unwind(move || {
+            vec.insert_many(1, BadIter { hint: 4, count: 2 });
+        });
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn panic_early_at_end() {
+        let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> =
+            smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),];
+        let result = ::std::panic::catch_unwind(move || {
+            vec.insert_many(2, BadIter { hint: 3, count: 1 });
+        });
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn panic_late_at_start() {
+        let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> =
+            smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),];
+        let result = ::std::panic::catch_unwind(move || {
+            vec.insert_many(0, BadIter { hint: 3, count: 5 });
+        });
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn panic_late_at_end() {
+        let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> =
+            smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),];
+        let result = ::std::panic::catch_unwind(move || {
+            vec.insert_many(2, BadIter { hint: 3, count: 5 });
+        });
+        assert!(result.is_err());
+    }
 }
 
 #[test]