diff --git a/src/arrayvec.rs b/src/arrayvec.rs index c7c4a2c..bb6ce19 100644 --- a/src/arrayvec.rs +++ b/src/arrayvec.rs @@ -704,7 +704,7 @@ impl std::convert::TryFrom<&[T]> for ArrayVec Err(CapacityError::new(())) } else { let mut array = Self::new(); - array.extend(slice.iter().cloned()); + array.extend_from_slice(slice); Ok(array) } } @@ -931,39 +931,75 @@ impl Drop for ScopeExitGuard /// Extend the `ArrayVec` with an iterator. /// -/// Does not extract more items than there is space for. No error -/// occurs if there are more iterator elements. +/// ***Panics*** if extending the vector exceeds its capacity. impl Extend for ArrayVec { + /// Extend the `ArrayVec` with an iterator. + /// + /// ***Panics*** if extending the vector exceeds its capacity. fn extend>(&mut self, iter: I) { - let take = self.capacity() - self.len(); unsafe { - 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; - } - }; - let mut iter = iter.into_iter(); - loop { - if ptr == end_ptr { break; } - if let Some(elt) = iter.next() { - raw_ptr_write(ptr, elt); - ptr = raw_ptr_add(ptr, 1); - guard.data += 1; - } else { - break; - } + self.extend_from_iter::<_, true>(iter) + } + } +} + +#[inline(never)] +#[cold] +fn extend_panic() { + panic!("ArrayVec: capacity exceeded in extend/from_iter"); +} + +impl ArrayVec { + /// Extend the arrayvec from the iterable. + /// + /// ## 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. + pub(crate) unsafe fn extend_from_iter(&mut self, iterable: I) + where I: IntoIterator + { + 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; + } + }; + 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); + ptr.write(elt); + ptr = raw_ptr_add(ptr, 1); + guard.data += 1; + } else { + return; // success } } } + + /// Extend the ArrayVec with clones of elements from the slice; + /// the length of the slice must be <= the remaining capacity in the arrayvec. + pub(crate) fn extend_from_slice(&mut self, slice: &[T]) + where T: Clone + { + let take = self.capacity() - self.len(); + debug_assert!(slice.len() <= take); + unsafe { + let slice = if take < slice.len() { &slice[..take] } else { slice }; + self.extend_from_iter::<_, false>(slice.iter().cloned()); + } + } } /// Rawptr add but uses arithmetic distance for ZST @@ -976,19 +1012,13 @@ unsafe fn raw_ptr_add(ptr: *mut T, offset: usize) -> *mut T { } } -unsafe fn raw_ptr_write(ptr: *mut T, value: T) { - if mem::size_of::() == 0 { - /* nothing */ - } else { - ptr::write(ptr, value) - } -} - /// Create an `ArrayVec` from an iterator. /// -/// Does not extract more items than there is space for. No error -/// occurs if there are more iterator elements. +/// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity. impl iter::FromIterator for ArrayVec { + /// Create an `ArrayVec` from an iterator. + /// + /// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity. fn from_iter>(iter: I) -> Self { let mut array = ArrayVec::new(); array.extend(iter); @@ -1014,8 +1044,8 @@ impl Clone for ArrayVec self.pop(); } } else { - let rhs_elems = rhs[self.len()..].iter().cloned(); - self.extend(rhs_elems); + let rhs_elems = &rhs[self.len()..]; + self.extend_from_slice(rhs_elems); } } } diff --git a/tests/tests.rs b/tests/tests.rs index 676cd7c..b89529e 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -251,11 +251,11 @@ fn test_drop_panics() { fn test_extend() { let mut range = 0..10; - let mut array: ArrayVec<_, 5> = range.by_ref().collect(); + let mut array: ArrayVec<_, 5> = range.by_ref().take(5).collect(); assert_eq!(&array[..], &[0, 1, 2, 3, 4]); assert_eq!(range.next(), Some(5)); - array.extend(range.by_ref()); + array.extend(range.by_ref().take(0)); assert_eq!(range.next(), Some(6)); let mut array: ArrayVec<_, 10> = (0..3).collect(); @@ -264,6 +264,25 @@ fn test_extend() { assert_eq!(&array[..], &[0, 1, 2, 3, 4]); } +#[should_panic] +#[test] +fn test_extend_capacity_panic_1() { + let mut range = 0..10; + + let _: ArrayVec<_, 5> = range.by_ref().collect(); +} + +#[should_panic] +#[test] +fn test_extend_capacity_panic_2() { + let mut range = 0..10; + + let mut array: ArrayVec<_, 5> = range.by_ref().take(5).collect(); + assert_eq!(&array[..], &[0, 1, 2, 3, 4]); + assert_eq!(range.next(), Some(5)); + array.extend(range.by_ref().take(1)); +} + #[test] fn test_is_send_sync() { let data = ArrayVec::, 5>::new(); @@ -304,7 +323,7 @@ fn test_drain() { v.drain(0..7); assert_eq!(&v[..], &[]); - v.extend(0..); + v.extend(0..8); v.drain(1..4); assert_eq!(&v[..], &[0, 4, 5, 6, 7]); let u: ArrayVec<_, 3> = v.drain(1..4).rev().collect(); @@ -320,7 +339,7 @@ fn test_drain_range_inclusive() { v.drain(0..=7); assert_eq!(&v[..], &[]); - v.extend(0..); + v.extend(0..8); v.drain(1..=4); assert_eq!(&v[..], &[0, 5, 6, 7]); let u: ArrayVec<_, 3> = v.drain(1..=2).rev().collect(); @@ -436,9 +455,9 @@ fn test_into_inner_2() { } #[test] -fn test_into_inner_3_() { +fn test_into_inner_3() { let mut v = ArrayVec::::new(); - v.extend(1..); + v.extend(1..=4); assert_eq!(v.into_inner().unwrap(), [1, 2, 3, 4]); } @@ -672,11 +691,11 @@ fn test_extend_zst() { #[derive(Copy, Clone, PartialEq, Debug)] struct Z; // Zero sized type - let mut array: ArrayVec<_, 5> = range.by_ref().map(|_| Z).collect(); + let mut array: ArrayVec<_, 5> = range.by_ref().take(5).map(|_| Z).collect(); assert_eq!(&array[..], &[Z; 5]); assert_eq!(range.next(), Some(5)); - array.extend(range.by_ref().map(|_| Z)); + array.extend(range.by_ref().take(0).map(|_| Z)); assert_eq!(range.next(), Some(6)); let mut array: ArrayVec<_, 10> = (0..3).map(|_| Z).collect();