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 May 14, 2020
1 parent 6c76c41 commit 31a4e89
Showing 1 changed file with 56 additions and 16 deletions.
72 changes: 56 additions & 16 deletions lib.rs
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 @@ -881,27 +879,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(
Expand All @@ -913,6 +925,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
Expand Down Expand Up @@ -1097,6 +1127,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>
Expand Down Expand Up @@ -2094,27 +2140,21 @@ 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(2, BadIter);
vec.insert_many(1, BadIter);
vec.insert_many(0, BadIter);
});
assert!(result.is_err());

drop(box1);
drop(box2);
}

#[test]
Expand Down

0 comments on commit 31a4e89

Please sign in to comment.