diff --git a/benches/extend.rs b/benches/extend.rs index ba33a932..95fb8c2d 100644 --- a/benches/extend.rs +++ b/benches/extend.rs @@ -1,68 +1,72 @@ - extern crate arrayvec; #[macro_use] extern crate bencher; use std::io::Write; +use std::hint::black_box; -use arrayvec::ArrayVec; +const CAP: usize = 1024 * 5; +use arrayvec::ArrayVec; use bencher::Bencher; -use bencher::black_box; fn extend_with_constant(b: &mut Bencher) { - let mut v = ArrayVec::::new(); - let cap = v.capacity(); + let mut v = ArrayVec::::new(); b.iter(|| { v.clear(); - let constant = black_box(1); - v.extend((0..cap).map(move |_| constant)); - v[511] + let constant = black_box(255u8); + v.extend((0..CAP).map(move |_| constant)); + black_box(v.as_ptr()); + black_box(v[CAP - 1]) }); b.bytes = v.capacity() as u64; } fn extend_with_range(b: &mut Bencher) { - let mut v = ArrayVec::::new(); - let cap = v.capacity(); + let mut v = ArrayVec::::new(); b.iter(|| { v.clear(); - let range = 0..cap; - v.extend(range.map(|x| black_box(x as _))); - v[511] + let r = black_box((0..CAP).map(|x| x as u8)); + v.extend(r); + black_box(v.as_ptr()); + black_box(v[CAP - 1]) }); b.bytes = v.capacity() as u64; } fn extend_with_slice(b: &mut Bencher) { - let mut v = ArrayVec::::new(); - let data = [1; 512]; + let data = [1u8; CAP]; + let mut v = ArrayVec::::new(); b.iter(|| { v.clear(); - let iter = data.iter().map(|&x| x); - v.extend(iter); - v[511] + let data = black_box(&data[..]); + v.extend(data.iter().copied()); + black_box(v.as_ptr()); + black_box(v[CAP - 1]) }); b.bytes = v.capacity() as u64; } fn extend_with_write(b: &mut Bencher) { - let mut v = ArrayVec::::new(); - let data = [1; 512]; + let data = [1u8; CAP]; + let mut v = ArrayVec::::new(); b.iter(|| { v.clear(); - v.write(&data[..]).ok(); - v[511] + let data = black_box(&data[..]); + v.write(data).ok(); + black_box(v[CAP - 1]) }); b.bytes = v.capacity() as u64; } fn extend_from_slice(b: &mut Bencher) { - let mut v = ArrayVec::::new(); - let data = [1; 512]; + let data = [1u8; CAP]; + let mut v = ArrayVec::::new(); b.iter(|| { v.clear(); - v.try_extend_from_slice(&data).ok(); - v[511] + let data = black_box(&data[..]); + v.try_extend_from_slice(data).ok(); + black_box(v.as_ptr()); + black_box(v[CAP - 1]); }); b.bytes = v.capacity() as u64; } diff --git a/src/arrayvec.rs b/src/arrayvec.rs index eec77328..add4fb16 100644 --- a/src/arrayvec.rs +++ b/src/arrayvec.rs @@ -1064,19 +1064,20 @@ impl<'a, T: 'a, const CAP: usize> Drop for Drain<'a, T, CAP> { } } -struct ScopeExitGuard - where F: FnMut(&Data, &mut T) -{ - value: T, - data: Data, - f: F, +/// Guard that writes a `usize` length back to a `LenUint` field on drop. +/// +/// Used to keep the `ArrayVec` length consistent if a panic occurs during +/// element-by-element writing: on panic the vector is left with the temporally +/// correct initialized length. +struct WritebackGuard<'a> { + target: &'a mut LenUint, + len: usize, } -impl Drop for ScopeExitGuard - where F: FnMut(&Data, &mut T) -{ +impl Drop for WritebackGuard<'_> { + #[inline(always)] fn drop(&mut self) { - (self.f)(&self.data, &mut self.value) + *self.target = self.len as LenUint; } } @@ -1092,7 +1093,7 @@ impl Extend for ArrayVec { #[track_caller] fn extend>(&mut self, iter: I) { unsafe { - self.extend_from_iter::<_, true>(iter) + self.extend_from_iter::<_, true>(iter.into_iter()) } } } @@ -1104,50 +1105,37 @@ fn extend_panic() { panic!("ArrayVec: capacity exceeded in extend/from_iter"); } + impl ArrayVec { - /// Extend the arrayvec from the iterable. + /// Extend the vector from the iterator. + /// + /// ***Panics*** if extending the vector exceeds its capacity. /// /// ## Safety /// - /// Unsafe because if CHECK is false, the length of the input is not checked. - /// The caller must ensure the length of the input fits in the capacity. + /// If CHECK is false, the iterator must yield at most `CAP - len()` elements. #[track_caller] - pub(crate) unsafe fn extend_from_iter(&mut self, iterable: I) - where I: IntoIterator + unsafe fn extend_from_iter(&mut self, iter: I) + where I: Iterator { - let take = self.capacity() - self.len(); let len = self.len(); - let mut ptr = raw_ptr_add(self.as_mut_ptr(), len); - let end_ptr = raw_ptr_add(ptr, take); - // Keep the length in a separate variable, write it back on scope - // exit. To help the compiler with alias analysis and stuff. - // We update the length to handle panic in the iteration of the - // user's iterator, without dropping any elements on the floor. - let mut guard = ScopeExitGuard { - value: &mut self.len, - data: len, - f: move |&len, self_len| { - **self_len = len as LenUint; - } - }; - let mut iter = iterable.into_iter(); - loop { - if let Some(elt) = iter.next() { - if ptr == end_ptr && CHECK { extend_panic(); } - debug_assert_ne!(ptr, end_ptr); - if mem::size_of::() != 0 { - ptr.write(elt); - } else { - // The ZST element has logically been moved into the vector. - // There is no memory to write, but dropping `elt` here would - // drop it once now and once again when the vector is dropped. - mem::forget(elt); - } - ptr = raw_ptr_add(ptr, 1); - guard.data += 1; - } else { - return; // success + let end = self.capacity(); + let ptr = self.as_mut_ptr(); + + // WritebackGuard updates self.len on drop (both success and panic). + let mut guard = WritebackGuard { target: &mut self.len, len: len }; + + // Take elements from the input iterator and write them into + // the arrayvec, as long as there is capacity left. + let mut index = len; + for elt in iter { + if index == end && CHECK { + extend_panic(); } + debug_assert_ne!(index, end); + ptr.add(index).write(elt); + guard.len += 1; + index += 1; } } @@ -1165,16 +1153,6 @@ impl ArrayVec { } } -/// Rawptr add but uses arithmetic distance for ZST -unsafe fn raw_ptr_add(ptr: *mut T, offset: usize) -> *mut T { - if mem::size_of::() == 0 { - // Special case for ZST - ptr.cast::().wrapping_add(offset).cast::() - } else { - ptr.add(offset) - } -} - /// Create an `ArrayVec` from an iterator. /// /// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.