Add support for allocating multiple items at the same time.

Author: jrmuizel

src/vec.rs

@@ -1,3 +1,4 @@
+
 use {alloc::vec::Vec, core::ptr};
 
 /// A typesafe helper that separates new value construction from
@@ -22,6 +23,26 @@ impl<'a, T> VecAllocation<'a, T> {
     }
 }
 
+pub struct VecArrayAllocation<'a, T: 'a, const N: usize > {
+    vec: &'a mut Vec<T>,
+    index: usize,
+}
+
+impl<'a, T, const N: usize> VecArrayAllocation<'a, T, N> {
+    /// Consumes self and writes the given value into the allocation.
+    // writing is safe because alloc() ensured enough capacity
+    // and `Allocation` holds a mutable borrow to prevent anyone else
+    // from breaking this invariant.
+    #[inline(always)]
+    pub fn init(self, value: [T; N]) -> usize {
+        unsafe {
+            ptr::write(self.vec.as_mut_ptr().add(self.index) as *mut [T; N], value);
+            self.vec.set_len(self.index + N);
+        }
+        self.index
+    }
+}
+
 /// An entry into a vector, similar to `std::collections::hash_map::Entry`.
 pub enum VecEntry<'a, T: 'a> {
     /// Entry has just been freshly allocated.
@@ -52,6 +73,8 @@ pub trait VecHelper<T> {
     /// Either returns an existing element, or grows the vector by one.
     /// Doesn't expect indices to be higher than the current length.
     fn entry(&mut self, index: usize) -> VecEntry<T>;
+
+    fn alloc_multiple<const N: usize>(&mut self) -> VecArrayAllocation<T, N>;
 }
 
 impl<T> VecHelper<T> for Vec<T> {
@@ -63,6 +86,12 @@ impl<T> VecHelper<T> for Vec<T> {
         VecAllocation { vec: self, index }
     }
 
+    fn alloc_multiple<const N: usize>(&mut self) -> VecArrayAllocation<T, N> {
+        self.reserve(N);
+        let index = self.len();
+        VecArrayAllocation { vec: self, index }
+    }
+
     fn entry(&mut self, index: usize) -> VecEntry<T> {
         if index < self.len() {
             VecEntry::Occupied(unsafe { self.get_unchecked_mut(index) })
@@ -94,3 +123,15 @@ fn test_zero_copy() {
     //TODO: make the expected value to be concrete
     assert_ne!(unsafe { vec[0].big[1] }, 1.0);
 }
+
+#[test]
+fn test_add_multiple() {
+    use alloc::vec;
+    let mut vec = vec![1, 2, 3];
+
+    // check that the new helper is not overwriting
+    vec.remove(0);
+    vec.alloc_multiple().init([4, 5]);
+    assert_eq!(vec, [2, 3, 4, 5])
+
+}