diff --git a/src/lib.rs b/src/lib.rs index 203a154..cf79ad8 100644 --- a/src/lib.rs +++ b/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 SmallVec { /// 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>(&mut self, index: usize, iterable: I) { let iter = iterable.into_iter(); if index == self.len() { @@ -991,13 +989,19 @@ impl SmallVec { 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 { @@ -1005,13 +1009,21 @@ impl SmallVec { 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 SmallVec { self.set_len(old_len + num_added); } + + struct DropOnPanic { + start: *mut T, + skip: Range, + len: usize, + } + + impl Drop for DropOnPanic { + 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 SmallVec { 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 SmallVec diff --git a/src/tests.rs b/src/tests.rs index f88e341..6197a6e 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -345,17 +345,12 @@ fn test_insert_many_panic() { } } - // 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(); @@ -363,9 +358,6 @@ fn test_insert_many_panic() { vec.insert_many(0, BadIter); }); assert!(result.is_err()); - - drop(box1); - drop(box2); } #[test]