diff --git a/lib.rs b/lib.rs index 5ee40b7..b0ed9ae 100644 --- a/lib.rs +++ b/lib.rs @@ -771,24 +771,33 @@ impl SmallVec { unsafe { let old_len = self.len(); assert!(index <= old_len); - let ptr = self.as_mut_ptr().offset(index as isize); + let mut ptr = self.as_mut_ptr().offset(index as isize); + + // Move the trailing elements. ptr::copy(ptr, ptr.offset(lower_size_bound as isize), old_len - index); - for (off, element) in iter.enumerate() { - if off < lower_size_bound { - ptr::write(ptr.offset(off as isize), element); - let len = self.len() + 1; - self.set_len(len); - } else { - // Iterator provided more elements than the hint. - assert!(index + off >= index); // Protect against overflow. - self.insert(index + off, element); + + // In case the iterator panics, don't double-drop the items we just copied above. + self.set_len(index); + + let mut num_added = 0; + for element in iter { + let mut cur = ptr.offset(num_added as isize); + 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().offset(index as isize); + cur = ptr.offset(num_added as isize); + ptr::copy(cur, cur.offset(1), old_len - index); } + ptr::write(cur, element); + num_added += 1; } - let num_added = self.len() - old_len; if num_added < lower_size_bound { // Iterator provided fewer elements than the hint ptr::copy(ptr.offset(lower_size_bound as isize), ptr.offset(num_added as isize), old_len - index); } + + self.set_len(old_len + num_added); } } @@ -937,19 +946,17 @@ impl SmallVec where A::Item: Clone { if n > A::size() { vec![elem; n].into() } else { + let mut v = SmallVec::::new(); unsafe { - let mut arr: A = ::std::mem::uninitialized(); - let ptr = arr.ptr_mut(); + let (ptr, len_ptr, _) = v.triple_mut(); + let mut local_len = SetLenOnDrop::new(len_ptr); for i in 0..n as isize { ::std::ptr::write(ptr.offset(i), elem.clone()); - } - - SmallVec { - capacity: n, - data: SmallVecData::from_inline(arr), + local_len.increment_len(1); } } + v } } } @@ -1337,6 +1344,33 @@ pub unsafe trait Array { fn ptr_mut(&mut self) -> *mut Self::Item; } +/// Set the length of the vec when the `SetLenOnDrop` value goes out of scope. +/// +/// Copied from https://github.com/rust-lang/rust/pull/36355 +struct SetLenOnDrop<'a> { + len: &'a mut usize, + local_len: usize, +} + +impl<'a> SetLenOnDrop<'a> { + #[inline] + fn new(len: &'a mut usize) -> Self { + SetLenOnDrop { local_len: *len, len: len } + } + + #[inline] + fn increment_len(&mut self, increment: usize) { + self.local_len += increment; + } +} + +impl<'a> Drop for SetLenOnDrop<'a> { + #[inline] + fn drop(&mut self) { + *self.len = self.local_len; + } +} + macro_rules! impl_array( ($($size:expr),+) => { $( @@ -1645,6 +1679,37 @@ mod tests { assert_eq!(&v.iter().map(|v| *v).collect::>(), &[0, 5, 6, 1, 2, 3]); } + #[test] + // https://github.com/servo/rust-smallvec/issues/96 + fn test_insert_many_panic() { + struct PanicOnDoubleDrop { + dropped: Box + } + + impl Drop for PanicOnDoubleDrop { + fn drop(&mut self) { + assert!(!*self.dropped, "already dropped"); + *self.dropped = true; + } + } + + struct BadIter; + impl Iterator for BadIter { + type Item = PanicOnDoubleDrop; + fn size_hint(&self) -> (usize, Option) { (1, None) } + fn next(&mut self) -> Option { panic!() } + } + + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = vec![ + PanicOnDoubleDrop { dropped: Box::new(false) }, + PanicOnDoubleDrop { dropped: Box::new(false) }, + ].into(); + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(0, BadIter); + }); + assert!(result.is_err()); + } + #[test] #[should_panic] fn test_invalid_grow() {