Skip to content

Commit e58f0c1

Browse files
committed
Fix leak on panic in insert_many.
Fixes #208.
1 parent 6c76c41 commit e58f0c1

File tree

1 file changed

+54
-16
lines changed

1 file changed

+54
-16
lines changed

lib.rs

+54-16
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ use core::hint::unreachable_unchecked;
8181
use core::iter::{repeat, FromIterator, FusedIterator, IntoIterator};
8282
use core::mem;
8383
use core::mem::MaybeUninit;
84-
use core::ops::{self, RangeBounds};
84+
use core::ops::{self, Range, RangeBounds};
8585
use core::ptr::{self, NonNull};
8686
use core::slice::{self, SliceIndex};
8787

@@ -865,8 +865,6 @@ impl<A: Array> SmallVec<A> {
865865

866866
/// Insert multiple elements at position `index`, shifting all following elements toward the
867867
/// back.
868-
///
869-
/// Note: when the iterator panics, this can leak memory.
870868
pub fn insert_many<I: IntoIterator<Item = A::Item>>(&mut self, index: usize, iterable: I) {
871869
let iter = iterable.into_iter();
872870
if index == self.len() {
@@ -881,27 +879,41 @@ impl<A: Array> SmallVec<A> {
881879
unsafe {
882880
let old_len = self.len();
883881
assert!(index <= old_len);
884-
let mut ptr = self.as_mut_ptr().add(index);
882+
let start = self.as_mut_ptr();
883+
let mut ptr = start.add(index);
885884

886885
// Move the trailing elements.
887886
ptr::copy(ptr, ptr.add(lower_size_bound), old_len - index);
888887

889888
// In case the iterator panics, don't double-drop the items we just copied above.
890-
self.set_len(index);
889+
self.set_len(0);
890+
let mut guard = DropOnPanic {
891+
start,
892+
skip: index..lower_size_bound,
893+
len: old_len + lower_size_bound,
894+
};
891895

892896
let mut num_added = 0;
893897
for element in iter {
894898
let mut cur = ptr.add(num_added);
895899
if num_added >= lower_size_bound {
896900
// Iterator provided more elements than the hint. Move trailing items again.
897901
self.reserve(1);
898-
ptr = self.as_mut_ptr().add(index);
902+
let start = self.as_mut_ptr();
903+
ptr = start.add(index);
899904
cur = ptr.add(num_added);
900905
ptr::copy(cur, cur.add(1), old_len - index);
906+
907+
guard.start = start;
908+
guard.len += 1;
909+
guard.skip.end += 1;
901910
}
902911
ptr::write(cur, element);
912+
guard.skip.start += 1;
903913
num_added += 1;
904914
}
915+
mem::forget(guard);
916+
905917
if num_added < lower_size_bound {
906918
// Iterator provided fewer elements than the hint
907919
ptr::copy(
@@ -913,6 +925,24 @@ impl<A: Array> SmallVec<A> {
913925

914926
self.set_len(old_len + num_added);
915927
}
928+
929+
struct DropOnPanic<T> {
930+
start: *mut T,
931+
skip: Range<usize>,
932+
len: usize,
933+
}
934+
935+
impl<T> Drop for DropOnPanic<T> {
936+
fn drop(&mut self) {
937+
for i in 0..self.len {
938+
if !self.skip.contains(&i) {
939+
unsafe {
940+
ptr::drop_in_place(self.start.add(i));
941+
}
942+
}
943+
}
944+
}
945+
}
916946
}
917947

918948
/// Convert a SmallVec to a Vec, without reallocating if the SmallVec has already spilled onto
@@ -1097,6 +1127,22 @@ impl<A: Array> SmallVec<A> {
10971127
data: SmallVecData::from_heap(ptr, length),
10981128
}
10991129
}
1130+
1131+
/// Returns a raw pointer to the vector's buffer.
1132+
pub fn as_ptr(&self) -> *const A::Item {
1133+
// We shadow the slice method of the same name to avoid going through
1134+
// `deref`, which creates an intermediate reference that may place
1135+
// additional safety constraints on the contents of the slice.
1136+
self.triple().0
1137+
}
1138+
1139+
/// Returns a raw mutable pointer to the vector's buffer.
1140+
pub fn as_mut_ptr(&mut self) -> *mut A::Item {
1141+
// We shadow the slice method of the same name to avoid going through
1142+
// `deref_mut`, which creates an intermediate reference that may place
1143+
// additional safety constraints on the contents of the slice.
1144+
self.triple_mut().0
1145+
}
11001146
}
11011147

11021148
impl<A: Array> SmallVec<A>
@@ -2094,27 +2140,19 @@ mod tests {
20942140
}
20952141
}
20962142

2097-
// These boxes are leaked on purpose by panicking `insert_many`,
2098-
// so we clean them up manually to appease Miri's leak checker.
2099-
let mut box1 = Box::new(false);
2100-
let mut box2 = Box::new(false);
2101-
21022143
let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = vec![
21032144
PanicOnDoubleDrop {
2104-
dropped: unsafe { Box::from_raw(&mut *box1) },
2145+
dropped: Box::new(false),
21052146
},
21062147
PanicOnDoubleDrop {
2107-
dropped: unsafe { Box::from_raw(&mut *box2) },
2148+
dropped: Box::new(false),
21082149
},
21092150
]
21102151
.into();
21112152
let result = ::std::panic::catch_unwind(move || {
21122153
vec.insert_many(0, BadIter);
21132154
});
21142155
assert!(result.is_err());
2115-
2116-
drop(box1);
2117-
drop(box2);
21182156
}
21192157

21202158
#[test]

0 commit comments

Comments
 (0)