IntrusiveArrayBuilder

Author: novacrazy

src/impl_alloc.rs

@@ -1,6 +1,6 @@
 use alloc::{boxed::Box, vec::Vec};
 
-use crate::{ArrayLength, GenericArray, LengthError};
+use crate::{ArrayLength, GenericArray, IntrusiveArrayBuilder, LengthError};
 
 impl<T, N: ArrayLength> TryFrom<Vec<T>> for GenericArray<T, N> {
     type Error = crate::LengthError;
@@ -11,11 +11,15 @@ impl<T, N: ArrayLength> TryFrom<Vec<T>> for GenericArray<T, N> {
         }
 
         unsafe {
-            let mut destination = crate::ArrayBuilder::new();
+            let mut destination = GenericArray::uninit();
+            let mut builder = IntrusiveArrayBuilder::new(&mut destination);
 
-            destination.extend(v.into_iter());
+            builder.extend(v.into_iter());
 
-            Ok(destination.assume_init())
+            Ok({
+                builder.finish();
+                IntrusiveArrayBuilder::array_assume_init(destination)
+            })
         }
     }
 }
@@ -162,11 +166,35 @@ unsafe impl<T, N: ArrayLength> GenericSequence<T> for Box<GenericArray<T, N>> {
     where
         F: FnMut(usize) -> T,
     {
-        let mut v = Vec::with_capacity(N::USIZE);
-        for i in 0..N::USIZE {
-            v.push(f(i));
+        unsafe {
+            use core::{
+                alloc::Layout,
+                mem::{size_of, MaybeUninit},
+                ptr,
+            };
+
+            // Box::new_uninit() is nightly-only
+            let ptr: *mut GenericArray<MaybeUninit<T>, N> = if size_of::<T>() == 0 {
+                ptr::NonNull::dangling().as_ptr()
+            } else {
+                alloc::alloc::alloc(Layout::new::<GenericArray<MaybeUninit<T>, N>>()).cast()
+            };
+
+            let mut builder = IntrusiveArrayBuilder::new(&mut *ptr);
+
+            {
+                let (builder_iter, position) = builder.iter_position();
+
+                builder_iter.enumerate().for_each(|(i, dst)| {
+                    dst.write(f(i));
+                    *position += 1;
+                });
+            }
+
+            builder.finish();
+
+            Box::from_raw(ptr.cast()) // IntrusiveArrayBuilder::array_assume_init
         }
-        GenericArray::try_from_vec(v).unwrap()
     }
 }
 

src/impl_serde.rs

@@ -1,8 +1,9 @@
 //! Serde serialization/deserialization implementation
 
-use crate::{ArrayLength, GenericArray};
+use crate::{ArrayLength, GenericArray, IntrusiveArrayBuilder};
 use core::fmt;
 use core::marker::PhantomData;
+
 use serde::de::{self, SeqAccess, Visitor};
 use serde::{ser::SerializeTuple, Deserialize, Deserializer, Serialize, Serializer};
 
@@ -62,11 +63,12 @@ where
         }
 
         unsafe {
-            let mut dst = crate::ArrayBuilder::new();
+            let mut dst = GenericArray::uninit();
+            let mut builder = IntrusiveArrayBuilder::new(&mut dst);
 
-            let (dst_iter, position) = dst.iter_position();
+            let (build_iter, position) = builder.iter_position();
 
-            for dst in dst_iter {
+            for dst in build_iter {
                 match seq.next_element()? {
                     Some(el) => {
                         dst.write(el);
@@ -81,7 +83,10 @@ where
                     return Err(de::Error::invalid_length(*position + 1, &self));
                 }
 
-                return Ok(dst.assume_init());
+                return Ok({
+                    builder.finish();
+                    IntrusiveArrayBuilder::array_assume_init(dst)
+                });
             }
 
             Err(de::Error::invalid_length(*position, &self))

src/internal.rs

@@ -1,3 +1,5 @@
+#![allow(dead_code)] // ArrayBuilder is soft-deprecated internally
+
 use crate::*;
 
 pub trait Sealed {}
@@ -8,6 +10,10 @@ impl<T> Sealed for [T; 0] {}
 ///
 /// You MUST increment the position while iterating to mark off created elements,
 /// which will be dropped if `into_inner` is not called.
+///
+/// This is soft-deprecated in favor of [`IntrusiveArrayBuilder`] due to Rust's
+/// lack of return-value optimization causing issues moving the array out of the struct.
+/// Still works fine for smaller arrays, though.
 pub struct ArrayBuilder<T, N: ArrayLength> {
     array: GenericArray<MaybeUninit<T>, N>,
     position: usize,
@@ -96,6 +102,97 @@ impl<T, N: ArrayLength> Drop for ArrayBuilder<T, N> {
     }
 }
 
+/// Similar to [`ArrayBuilder`] but uses a reference to a pre-allocated array, be
+/// it on the stack or heap.
+pub struct IntrusiveArrayBuilder<'a, T, N: ArrayLength> {
+    array: &'a mut GenericArray<MaybeUninit<T>, N>,
+    position: usize,
+}
+
+impl<'a, T, N: ArrayLength> IntrusiveArrayBuilder<'a, T, N> {
+    /// Begin building an array
+    #[inline(always)]
+    pub fn new(array: &'a mut GenericArray<MaybeUninit<T>, N>) -> IntrusiveArrayBuilder<T, N> {
+        IntrusiveArrayBuilder { array, position: 0 }
+    }
+
+    /// Consume an iterator, `.zip`-ing it to fill some or all of the array. This does not check if the
+    /// iterator had extra elements or too few elements.
+    ///
+    /// This makes no attempt to continue where a previous `extend` leaves off. Therefore, it should
+    /// only be used once per `ArrayBuilder`.
+    #[inline(always)]
+    pub unsafe fn extend(&mut self, source: impl Iterator<Item = T>) {
+        let (destination, position) = (self.array.iter_mut(), &mut self.position);
+
+        destination.zip(source).for_each(|(dst, src)| {
+            dst.write(src);
+            *position += 1;
+        });
+    }
+
+    /// Returns true if the write position equals the array size
+    #[inline(always)]
+    pub fn is_full(&self) -> bool {
+        self.position == N::USIZE
+    }
+
+    /// Creates a mutable iterator for writing to the array elements.
+    ///
+    /// You MUST increment the position value (given as a mutable reference) as you iterate
+    /// to mark how many elements have been created.
+    ///
+    /// ```
+    /// #[cfg(feature = "internals")]
+    /// # {
+    /// # use generic_array::{GenericArray, internals::IntrusiveArrayBuilder, typenum::U5};
+    /// # struct SomeType;
+    /// fn make_some_struct() -> SomeType { SomeType }
+    /// unsafe {
+    ///     let mut array = GenericArray::uninit();
+    ///     let mut builder = IntrusiveArrayBuilder::<SomeType, U5>::new(&mut array);
+    ///     let (dst_iter, position) = builder.iter_position();
+    ///     for dst in dst_iter {
+    ///         dst.write(make_some_struct());
+    ///         // MUST be done AFTER ownership of the value has been given to `dst.write`
+    ///         *position += 1;
+    ///     }
+    ///     let your_array = { builder.finish(); IntrusiveArrayBuilder::array_assume_init(array) };
+    /// }
+    /// # }
+    /// ```
+    #[inline(always)]
+    pub unsafe fn iter_position(&mut self) -> (slice::IterMut<MaybeUninit<T>>, &mut usize) {
+        (self.array.iter_mut(), &mut self.position)
+    }
+
+    /// When done writing (assuming all elements have been written to),
+    /// get the inner array.
+    #[inline(always)]
+    pub unsafe fn finish(self) {
+        debug_assert!(self.is_full());
+        mem::forget(self)
+    }
+
+    /// Similar to [`GenericArray::assume_init`] but not `const` and optimizes better.
+    #[inline(always)]
+    pub unsafe fn array_assume_init(array: GenericArray<MaybeUninit<T>, N>) -> GenericArray<T, N> {
+        ptr::read(&array as *const _ as *const MaybeUninit<GenericArray<T, N>>).assume_init()
+    }
+}
+
+impl<'a, T, N: ArrayLength> Drop for IntrusiveArrayBuilder<'a, T, N> {
+    fn drop(&mut self) {
+        unsafe {
+            ptr::drop_in_place(
+                // Same cast as MaybeUninit::slice_assume_init_mut
+                self.array.get_unchecked_mut(..self.position) as *mut [MaybeUninit<T>]
+                    as *mut [T],
+            );
+        }
+    }
+}
+
 /// **UNSAFE**: Consumes an array one element at a time.
 ///
 /// You MUST increment the position while iterating and any leftover elements

src/iter.rs

@@ -15,18 +15,6 @@ pub struct GenericArrayIter<T, N: ArrayLength> {
     index_back: usize,
 }
 
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    fn send<I: Send>(_iter: I) {}
-
-    #[test]
-    fn test_send_iter() {
-        send(GenericArray::from([1, 2, 3, 4]).into_iter());
-    }
-}
-
 impl<T, N: ArrayLength> GenericArrayIter<T, N> {
     /// Returns the remaining items of this iterator as a slice
     #[inline(always)]
@@ -242,3 +230,15 @@ impl<T, N: ArrayLength> ExactSizeIterator for GenericArrayIter<T, N> {
 impl<T, N: ArrayLength> FusedIterator for GenericArrayIter<T, N> {}
 
 // TODO: Implement `TrustedLen` when stabilized
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn send<I: Send>(_iter: I) {}
+
+    #[test]
+    fn test_send_iter() {
+        send(GenericArray::from([1, 2, 3, 4]).into_iter());
+    }
+}

src/lib.rs

@@ -8,7 +8,7 @@
 //! Before Rust 1.51, arrays `[T; N]` were problematic in that they couldn't be
 //! generic with respect to the length `N`, so this wouldn't work:
 //!
-//! ```rust{compile_fail}
+//! ```compile_fail
 //! struct Foo<N> {
 //!     data: [i32; N],
 //! }
@@ -133,7 +133,7 @@ pub mod functional;
 pub mod sequence;
 
 mod internal;
-use internal::{ArrayBuilder, ArrayConsumer, Sealed};
+use internal::{ArrayConsumer, IntrusiveArrayBuilder, Sealed};
 
 // re-export to allow doc_auto_cfg to handle it
 #[cfg(feature = "internals")]
@@ -142,8 +142,10 @@ pub mod internals {
     //!
     //! These are used internally for building and consuming generic arrays. When used correctly,
     //! they can ensure elements are correctly dropped if something panics while using them.
+    //!
+    //! The API of these is not guarenteed to be stable, as they are not intended for general use.
 
-    pub use crate::internal::{ArrayBuilder, ArrayConsumer};
+    pub use crate::internal::{ArrayBuilder, ArrayConsumer, IntrusiveArrayBuilder};
 }
 
 use self::functional::*;
@@ -510,18 +512,20 @@ where
         F: FnMut(usize) -> T,
     {
         unsafe {
-            let mut destination = ArrayBuilder::new();
+            let mut array = GenericArray::<T, N>::uninit();
+            let mut builder = IntrusiveArrayBuilder::new(&mut array);
 
             {
-                let (destination_iter, position) = destination.iter_position();
+                let (builder_iter, position) = builder.iter_position();
 
-                destination_iter.enumerate().for_each(|(i, dst)| {
+                builder_iter.enumerate().for_each(|(i, dst)| {
                     dst.write(f(i));
                     *position += 1;
                 });
             }
 
-            destination.assume_init()
+            builder.finish();
+            IntrusiveArrayBuilder::array_assume_init(array)
         }
     }
 
@@ -968,15 +972,19 @@ impl<T, N: ArrayLength> GenericArray<T, N> {
         }
 
         unsafe {
-            let mut destination = ArrayBuilder::new();
+            let mut array = GenericArray::uninit();
+            let mut builder = IntrusiveArrayBuilder::new(&mut array);
 
-            destination.extend(&mut iter);
+            builder.extend(&mut iter);
 
-            if !destination.is_full() || iter.next().is_some() {
+            if !builder.is_full() || iter.next().is_some() {
                 return Err(LengthError);
             }
 
-            Ok(destination.assume_init())
+            Ok({
+                builder.finish();
+                IntrusiveArrayBuilder::array_assume_init(array)
+            })
         }
     }
 }