Add ArrayVec::splice

src/arrayvec.rs

@@ -23,6 +23,7 @@ use serde::{Serialize, Deserialize, Serializer, Deserializer};
 use crate::LenUint;
 use crate::errors::CapacityError;
 use crate::arrayvec_impl::ArrayVecImpl;
+use crate::splice::Splice;
 use crate::utils::MakeMaybeUninit;
 
 /// A vector with a fixed capacity.
@@ -41,8 +42,8 @@ use crate::utils::MakeMaybeUninit;
 /// available. The ArrayVec can be converted into a by value iterator.
 pub struct ArrayVec<T, const CAP: usize> {
     // the `len` first elements of the array are initialized
-    xs: [MaybeUninit<T>; CAP],
-    len: LenUint,
+    pub(super) xs: [MaybeUninit<T>; CAP],
+    pub(super) len: LenUint,
 }
 
 impl<T, const CAP: usize> Drop for ArrayVec<T, CAP> {
@@ -253,7 +254,7 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
 
 
     /// Get pointer to where element at `index` would be
-    unsafe fn get_unchecked_ptr(&mut self, index: usize) -> *mut T {
+    pub(super) unsafe fn get_unchecked_ptr(&mut self, index: usize) -> *mut T {
         self.as_mut_ptr().add(index)
     }
 
@@ -579,6 +580,33 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
         Ok(())
     }
 
+    /// Creates a splicing iterator that replaces the specified range in the vector with the given `replace_with` iterator and yields the removed items. `replace_with` does not need to be the same length as `range`.
+    ///
+    /// `range` is removed even if the iterator is not consumed until the end.
+    ///
+    /// It is unspecified how many elements are removed from the vector if the `Splice` value is leaked.
+    ///
+    /// The input iterator `replace_with` is only consumed when the `Splice` value is dropped.
+    ///
+    /// ```
+    /// use std::iter::FromIterator;
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut vec: ArrayVec<usize, 4> = ArrayVec::from_iter((0..4));
+    /// let elements_popped: Vec<_> = vec.splice(1..3, vec![7, 9]).into_iter().collect();
+    /// assert_eq!(&vec[..], &[0, 7, 9, 3]);
+    /// assert_eq!(&elements_popped[..], &[1, 2]);
+    /// ```
+    ///
+    /// ***Panics*** if splicing the vector exceeds its capacity.
+    pub fn splice<R, I>(&mut self, range: R, replace_with: I) -> Splice<'_, I::IntoIter, CAP>
+        where
+            R: RangeBounds<usize>,
+            I: IntoIterator<Item = T>,
+        {
+            Splice { drain: self.drain(range), replace_with: replace_with.into_iter() }
+        }
+
     /// 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
     /// removed even if the iterator is not consumed until the end.
@@ -925,12 +953,12 @@ where
 /// A draining iterator for `ArrayVec`.
 pub struct Drain<'a, T: 'a, const CAP: usize> {
     /// Index of tail to preserve
-    tail_start: usize,
+    pub(super) tail_start: usize,
     /// Length of tail
-    tail_len: usize,
+    pub(super) tail_len: usize,
     /// Current remaining range to remove
-    iter: slice::Iter<'a, T>,
-    vec: *mut ArrayVec<T, CAP>,
+    pub(super) iter: slice::Iter<'a, T>,
+    pub(super) vec: *mut ArrayVec<T, CAP>,
 }
 
 unsafe impl<'a, T: Sync, const CAP: usize> Sync for Drain<'a, T, CAP> {}

src/lib.rs

@@ -48,6 +48,7 @@ macro_rules! assert_capacity_limit_const {
 
 mod arrayvec_impl;
 mod arrayvec;
+mod splice;
 mod array_string;
 mod char;
 mod errors;

src/splice.rs

@@ -0,0 +1,97 @@
+use core::ptr;
+use core::slice;
+
+use crate::Drain;
+
+/// A splicing iterator adapted for `ArrayVec` from `Vec`.
+pub struct Splice<
+    'a,
+    I: Iterator + 'a,
+    const CAP: usize
+> {
+    pub(super) drain: Drain<'a, I::Item, CAP>,
+    pub(super) replace_with: I,
+}
+
+impl<I: Iterator, const CAP: usize> Iterator for Splice<'_, I, CAP> {
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.drain.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.drain.size_hint()
+    }
+}
+
+impl<I: Iterator, const CAP: usize> DoubleEndedIterator for Splice<'_, I, CAP> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.drain.next_back()
+    }
+}
+
+impl<I: Iterator, const CAP: usize> ExactSizeIterator for Splice<'_, I, CAP> {}
+
+impl<I: Iterator, const CAP: usize> Drop for Splice<'_, I, CAP> {
+    fn drop(&mut self) {
+        self.drain.by_ref().for_each(drop);
+
+        unsafe {
+            if self.drain.tail_len == 0 {
+                let target_vec = &mut *self.drain.vec;
+                target_vec.extend(self.replace_with.by_ref());
+                return;
+            }
+
+            // First fill the range left by drain().
+            if !self.drain.fill(&mut self.replace_with) {
+                return;
+            }
+
+            // There may be more elements. Use the lower bound as an estimate.
+            // FIXME: Is the upper bound a better guess? Or something else?
+            let (lower_bound, _upper_bound) = self.replace_with.size_hint();
+            if lower_bound > 0 {
+                if !self.drain.fill(&mut self.replace_with) {
+                    return;
+                }
+            }
+
+            // Collect any remaining elements.
+            // This is a zero-length vector which does not allocate if `lower_bound` was exact.
+            let mut collected = self.replace_with.by_ref().collect::<Vec<I::Item>>().into_iter();
+            // Now we have an exact count.
+            if collected.len() > 0 {
+                let filled = self.drain.fill(&mut collected);
+                debug_assert!(filled);
+                debug_assert_eq!(collected.len(), 0);
+            }
+        }
+        // Let `Drain::drop` move the tail back if necessary and restore `vec.len`.
+    }
+}
+
+/// Private helper methods for `Splice::drop`
+impl<T, const CAP: usize> Drain<'_, T, CAP> {
+    /// The range from `self.vec.len` to `self.tail_start` contains elements
+    /// that have been moved out.
+    /// Fill that range as much as possible with new elements from the `replace_with` iterator.
+    /// Returns `true` if we filled the entire range. (`replace_with.next()` didn’t return `None`.)
+    unsafe fn fill<I: Iterator<Item = T>>(&mut self, replace_with: &mut I) -> bool {
+        let vec = &mut *self.vec;
+        let range_start = vec.len as usize;
+        let range_end = self.tail_start;
+        let range_slice = slice::from_raw_parts_mut(vec.get_unchecked_ptr(range_start), range_end - range_start);
+
+        for place in range_slice {
+            if let Some(new_item) = replace_with.next() {
+                ptr::write(place, new_item);
+                vec.len += 1;
+            } else {
+                return false;
+            }
+        }
+        true
+    }
+}

tests/tests.rs

@@ -65,6 +65,78 @@ fn test_extend_from_slice_error() {
     assert_matches!(res, Err(_));
 }
 
+#[test]
+fn test_splice_works_when_it_should() {
+    // Comparing behavior with std::vec::Vec to make sure it's analogous
+    let mut vec: Vec<_> = (0..5).collect();
+    let mut array: ArrayVec<_,  5> = (0..5).collect();
+    assert_eq!(&vec[..], &[0, 1, 2, 3, 4]);
+    assert_eq!(&array[..], &[0, 1, 2, 3, 4]);
+    // Typical case
+    let vec_popped: Vec<_> = vec.splice( 1..4, vec![11, 12, 13]).into_iter().collect();
+    assert_eq!(&vec[..], &[0, 11, 12, 13, 4]);
+    assert_eq!(&vec_popped, &[1, 2, 3]);
+    let array_popped: Vec<_> = array.splice( 1..4, vec![11, 12, 13]).into_iter().collect();
+    assert_eq!(&array[..], &vec[..]);
+    assert_eq!(&array_popped, &vec_popped);
+    // `replace_with` shorter than `range`
+    let vec_popped: Vec<_> = vec.splice( 2..5, vec![21, 22]).into_iter().collect();
+    assert_eq!(&vec[..], &[0, 11, 21, 22]);
+    assert_eq!(&vec_popped, &[12, 13, 4]);
+    let array_popped: Vec<_> = array.splice( 2..5, vec![21, 22]).into_iter().collect();
+    assert_eq!(&array[..], &vec[..]);
+    assert_eq!(&array_popped, &vec_popped);
+    // `range` shorter than `replace_with`
+    let vec_popped: Vec<_> = vec.splice( 3..4, vec![31, 32]).into_iter().collect();
+    assert_eq!(&vec[..], &[0, 11, 21, 31, 32]);
+    assert_eq!(&vec_popped, &[22]);
+    let array_popped: Vec<_> = array.splice( 3..4, vec![31, 32]).into_iter().collect();
+    assert_eq!(&array[..], &vec[..]);
+    assert_eq!(&array_popped, &vec_popped);
+    //`replace_with` shorter than open `range`
+    let vec_popped: Vec<_> = vec.splice( 1.., vec![41, 42]).into_iter().collect();
+    assert_eq!(&vec[..], &[0, 41, 42]);
+    assert_eq!(&vec_popped, &[11, 21, 31, 32]);
+    let array_popped: Vec<_> = array.splice( 1.., vec![41, 42]).into_iter().collect();
+    assert_eq!(&array[..], &vec[..]);
+    assert_eq!(&array_popped, &vec_popped);
+    // `range` shorter than open `replace_with`
+    let vec_popped: Vec<_> = vec.splice( 3.., vec![51, 52]).into_iter().collect();
+    assert_eq!(&vec[..], &[0, 41, 42, 51, 52]);
+    assert_eq!(&vec_popped, &[]);
+    let array_popped: Vec<_> = array.splice( 3.., vec![51, 52]).into_iter().collect();
+    assert_eq!(&array[..], &vec[..]);
+    assert_eq!(&array_popped, &vec_popped);
+}
+
+#[test]
+#[should_panic]
+fn test_splice_fails_on_overflowing_range() {
+    // Comparing behavior with std::vec::Vec to make sure it's analogous
+    let mut vec: Vec<_> = (0..5).collect();
+    let mut array: ArrayVec<_,  5> = (0..5).collect();
+    let vec_popped: Vec<_> = vec.splice( 4..6, vec![11]).into_iter().collect();
+    // This is actually less than 5!
+    // Maybe this would be fine, but we don't want to allow using overflowing ranges with ArrayVec anyway
+    assert_eq!(&vec[..], &[0, 1, 2, 11]);
+    assert_eq!(&vec_popped, &[3, 4]);
+    array.splice(4..6, vec![11]);
+}
+
+#[test]
+#[should_panic]
+fn test_splice_fails_if_operation_would_enlarge_vec() {
+    // Comparing behavior with std::vec::Vec to make sure it's analogous
+    let mut vec: Vec<_> = (0..5).collect();
+    let mut array: ArrayVec<_,  5> = (0..5).collect();
+    let vec_popped: Vec<_> = vec.splice( 4..5, vec![11, 12, 13]).into_iter().collect();
+    // This is now 6 instead of 5!
+    // No way for this ArrayVec to fit this in
+    assert_eq!(&vec[..], &[0, 1, 2, 11, 12, 13]);
+    assert_eq!(&vec_popped, &[3, 4]);
+    array.splice(4..5, vec![11, 12, 13]);
+}
+
 #[test]
 fn test_try_from_slice_error() {
     use arrayvec::ArrayVec;
@@ -790,4 +862,4 @@ fn test_arraystring_zero_filled_has_some_sanity_checks() {
     let string = ArrayString::<4>::zero_filled();
     assert_eq!(string.as_str(), "\0\0\0\0");
     assert_eq!(string.len(), 4);
-}
+}