Skip to content

Commit

Permalink
Fix leak on panic in insert_many.
Browse files Browse the repository at this point in the history
Fixes #208.
  • Loading branch information
mbrubeck committed Apr 22, 2020
1 parent 6c76c41 commit 50ee7da
Showing 1 changed file with 34 additions and 14 deletions.
48 changes: 34 additions & 14 deletions lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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() {
Expand All @@ -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 {
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -2094,27 +2122,19 @@ 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();
let result = ::std::panic::catch_unwind(move || {
vec.insert_many(0, BadIter);
});
assert!(result.is_err());

drop(box1);
drop(box2);
}

#[test]
Expand Down

0 comments on commit 50ee7da

Please sign in to comment.