Merge pull request #101 from Thomasdezeeuw/extend_from_slice

Add ArrayVec.extend_from_slice and capacity_left

Author: bluss

benches/extend.rs

@@ -40,5 +40,16 @@ fn extend_with_slice(b: &mut Bencher) {
     b.bytes = v.capacity() as u64;
 }
 
-benchmark_group!(benches, extend_with_constant, extend_with_range, extend_with_slice);
+fn extend_with_slice_fn(b: &mut Bencher) {
+    let mut v = ArrayVec::<[u8; 512]>::new();
+    let data = [1; 512];
+    b.iter(|| {
+        v.clear();
+        black_box(v.extend_from_slice(&data));
+        v[0]
+    });
+    b.bytes = v.capacity() as u64;
+}
+
+benchmark_group!(benches, extend_with_constant, extend_with_range, extend_with_slice, extend_with_slice_fn);
 benchmark_main!(benches);

src/lib.rs

@@ -172,6 +172,19 @@ impl<A: Array> ArrayVec<A> {
     /// ```
     pub fn is_full(&self) -> bool { self.len() == self.capacity() }
 
+    /// Returns the capacity left in the `ArrayVec`.
+    ///
+    /// ```
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut array = ArrayVec::from([1, 2, 3]);
+    /// array.pop();
+    /// assert_eq!(array.capacity_left(), 1);
+    /// ```
+    pub fn capacity_left(&self) -> usize {
+        self.capacity() - self.len()
+    }
+
     /// Push `element` to the end of the vector.
     ///
     /// ***Panics*** if the vector is already full.
@@ -523,6 +536,39 @@ impl<A: Array> ArrayVec<A> {
         self.len = Index::from(length);
     }
 
+    /// Copy and appends all elements in a slice to the `ArrayVec`.
+    ///
+    /// ```
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut vec: ArrayVec<[usize; 10]> = ArrayVec::new();
+    /// vec.push(1);
+    /// vec.extend_from_slice(&[2, 3]);
+    /// assert_eq!(&vec[..], &[1, 2, 3]);
+    /// ```
+    ///
+    /// # Panics
+    ///
+    /// This method will panic if the capacity left (see [`capacity_left`]) is
+    /// smaller then the length of the provided slice.
+    ///
+    /// [`capacity_left`]: #method.capacity_left
+    pub fn extend_from_slice(&mut self, other: &[A::Item])
+        where A::Item: Copy,
+    {
+        if self.capacity_left() < other.len() {
+            panic!("ArrayVec::extend_from_slice: slice is larger then capacity left");
+        }
+
+        let self_len = self.len();
+        let other_len = other.len();
+
+        unsafe {
+            let dst = self.xs.as_mut_ptr().offset(self_len as isize);
+            ptr::copy_nonoverlapping(other.as_ptr(), dst, other_len);
+            self.set_len(self_len + other_len);
+        }
+    }
 
     /// Create a draining iterator that removes the specified range in the vector
     /// and yields the removed items from start to end. The element range is
@@ -1033,16 +1079,9 @@ impl<A: Array> Ord for ArrayVec<A> where A::Item: Ord {
 /// Requires `features="std"`.
 impl<A: Array<Item=u8>> io::Write for ArrayVec<A> {
     fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        unsafe {
-            let len = self.len();
-            let mut tail = slice::from_raw_parts_mut(self.get_unchecked_mut(len),
-                                                     A::capacity() - len);
-            let result = tail.write(data);
-            if let Ok(written) = result {
-                self.set_len(len + written);
-            }
-            result
-        }
+        let len = cmp::min(self.capacity_left(), data.len());
+        self.extend_from_slice(&data[..len]);
+        Ok(len)
     }
     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }

tests/tests.rs

@@ -27,6 +27,31 @@ fn test_simple() {
     assert_eq!(sum_len, 8);
 }
 
+#[test]
+fn test_capacity_left() {
+    let mut vec: ArrayVec<[usize; 4]> = ArrayVec::new();
+    assert_eq!(vec.capacity_left(), 4);
+    vec.push(1);
+    assert_eq!(vec.capacity_left(), 3);
+    vec.push(2);
+    assert_eq!(vec.capacity_left(), 2);
+    vec.push(3);
+    assert_eq!(vec.capacity_left(), 1);
+    vec.push(4);
+    assert_eq!(vec.capacity_left(), 0);
+}
+
+#[test]
+fn test_extend_from_slice() {
+    let mut vec: ArrayVec<[usize; 10]> = ArrayVec::new();
+
+    vec.extend_from_slice(&[1, 2, 3]);
+    assert_eq!(vec.len(), 3);
+    assert_eq!(&vec[..], &[1, 2, 3]);
+    assert_eq!(vec.pop(), Some(3));
+    assert_eq!(&vec[..], &[1, 2]);
+}
+
 #[test]
 fn test_u16_index() {
     const N: usize = 4096;