Merge pull request #175 from bluss/extend-panic-on-capacity-error

Panic in .extend() and from_iter on capacity exceeded

Author: bluss

src/arrayvec.rs

@@ -704,7 +704,7 @@ impl<T, const CAP: usize> std::convert::TryFrom<&[T]> for ArrayVec<T, CAP>
             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<T, Data, F> Drop for ScopeExitGuard<T, Data, F>
 
 /// 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<T, const CAP: usize> Extend<T> for ArrayVec<T, CAP> {
+    /// Extend the `ArrayVec` with an iterator.
+    /// 
+    /// ***Panics*** if extending the vector exceeds its capacity.
     fn extend<I: IntoIterator<Item=T>>(&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<T, const CAP: usize> ArrayVec<T, CAP> {
+    /// 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<I, const CHECK: bool>(&mut self, iterable: I)
+        where I: IntoIterator<Item = T>
+    {
+        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<T>(ptr: *mut T, offset: usize) -> *mut T {
     }
 }
 
-unsafe fn raw_ptr_write<T>(ptr: *mut T, value: T) {
-    if mem::size_of::<T>() == 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<T, const CAP: usize> iter::FromIterator<T> for ArrayVec<T, CAP> {
+    /// Create an `ArrayVec` from an iterator.
+    /// 
+    /// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.
     fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> Self {
         let mut array = ArrayVec::new();
         array.extend(iter);
@@ -1014,8 +1044,8 @@ impl<T, const CAP: usize> Clone for ArrayVec<T, CAP>
                 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);
         }
     }
 }

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::<Vec<i32>,  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::<i32,  4>::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();