From 13ddbf31272826c0ea89a63e349734f6f22ad7c8 Mon Sep 17 00:00:00 2001 From: novacrazy Date: Thu, 4 Jul 2024 11:45:20 -0500 Subject: [PATCH] IntrusiveArrayBuilder --- src/impl_alloc.rs | 44 +++++++++++++++++---- src/impl_serde.rs | 15 +++++--- src/internal.rs | 97 +++++++++++++++++++++++++++++++++++++++++++++++ src/iter.rs | 24 ++++++------ src/lib.rs | 30 +++++++++------ 5 files changed, 174 insertions(+), 36 deletions(-) diff --git a/src/impl_alloc.rs b/src/impl_alloc.rs index 10b0fb8ac..d072e6d10 100644 --- a/src/impl_alloc.rs +++ b/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 TryFrom> for GenericArray { type Error = crate::LengthError; @@ -11,11 +11,15 @@ impl TryFrom> for GenericArray { } 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 GenericSequence for Box> { 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, N> = if size_of::() == 0 { + ptr::NonNull::dangling().as_ptr() + } else { + alloc::alloc::alloc(Layout::new::, 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() } } diff --git a/src/impl_serde.rs b/src/impl_serde.rs index 4b012bb4c..5ca7ad79e 100644 --- a/src/impl_serde.rs +++ b/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)) diff --git a/src/internal.rs b/src/internal.rs index 7897d1fe8..44d0ce08f 100644 --- a/src/internal.rs +++ b/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 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 { array: GenericArray, N>, position: usize, @@ -96,6 +102,97 @@ impl Drop for ArrayBuilder { } } +/// 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, 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, N>) -> IntrusiveArrayBuilder { + 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) { + 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::::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>, &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, N>) -> GenericArray { + ptr::read(&array as *const _ as *const MaybeUninit>).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] + as *mut [T], + ); + } + } +} + /// **UNSAFE**: Consumes an array one element at a time. /// /// You MUST increment the position while iterating and any leftover elements diff --git a/src/iter.rs b/src/iter.rs index 29d40e4e5..fb590399d 100644 --- a/src/iter.rs +++ b/src/iter.rs @@ -15,18 +15,6 @@ pub struct GenericArrayIter { index_back: usize, } -#[cfg(test)] -mod test { - use super::*; - - fn send(_iter: I) {} - - #[test] - fn test_send_iter() { - send(GenericArray::from([1, 2, 3, 4]).into_iter()); - } -} - impl GenericArrayIter { /// Returns the remaining items of this iterator as a slice #[inline(always)] @@ -242,3 +230,15 @@ impl ExactSizeIterator for GenericArrayIter { impl FusedIterator for GenericArrayIter {} // TODO: Implement `TrustedLen` when stabilized + +#[cfg(test)] +mod test { + use super::*; + + fn send(_iter: I) {} + + #[test] + fn test_send_iter() { + send(GenericArray::from([1, 2, 3, 4]).into_iter()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 0e427adb8..116ed30d1 100644 --- a/src/lib.rs +++ b/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 { //! 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::::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 GenericArray { } 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) + }) } } }