From 50ee7da25d1521fb35093601cadd7c66b559ad5b Mon Sep 17 00:00:00 2001
From: Matt Brubeck <mbrubeck@limpet.net>
Date: Wed, 22 Apr 2020 13:34:45 -0700
Subject: [PATCH] Fix leak on panic in `insert_many`.

Fixes #208.
---
 lib.rs | 48 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/lib.rs b/lib.rs
index ed83a75..6db096b 100644
--- a/lib.rs
+++ b/lib.rs
@@ -81,7 +81,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};
 
@@ -865,8 +865,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() {
@@ -887,7 +885,12 @@ impl<A: Array> SmallVec<A> {
             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 {
+                ptr: self.as_mut_ptr(),
+                skip: index..lower_size_bound,
+                len: old_len + lower_size_bound,
+            };
 
             let mut num_added = 0;
             for element in iter {
@@ -898,10 +901,17 @@ impl<A: Array> SmallVec<A> {
                     ptr = self.as_mut_ptr().add(index);
                     cur = ptr.add(num_added);
                     ptr::copy(cur, cur.add(1), old_len - index);
+
+                    guard.ptr = self.as_mut_ptr();
+                    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(
@@ -913,6 +923,24 @@ impl<A: Array> SmallVec<A> {
 
             self.set_len(old_len + num_added);
         }
+
+        struct DropOnPanic<T> {
+            ptr: *mut T,
+            len: usize,
+            skip: Range<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.ptr.add(i));
+                        }
+                    }
+                }
+            }
+        }
     }
 
     /// Convert a SmallVec to a Vec, without reallocating if the SmallVec has already spilled onto
@@ -2094,17 +2122,12 @@ mod tests {
             }
         }
 
-        // 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);
-
         let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = vec![
             PanicOnDoubleDrop {
-                dropped: unsafe { Box::from_raw(&mut *box1) },
+                dropped: Box::new(false),
             },
             PanicOnDoubleDrop {
-                dropped: unsafe { Box::from_raw(&mut *box2) },
+                dropped: Box::new(false),
             },
         ]
         .into();
@@ -2112,9 +2135,6 @@ mod tests {
             vec.insert_many(0, BadIter);
         });
         assert!(result.is_err());
-
-        drop(box1);
-        drop(box2);
     }
 
     #[test]