Skip to content

Commit

Permalink
IntrusiveArrayBuilder
Browse files Browse the repository at this point in the history
  • Loading branch information
novacrazy committed Jul 4, 2024
1 parent 2bc2034 commit 13ddbf3
Show file tree
Hide file tree
Showing 5 changed files with 174 additions and 36 deletions.
44 changes: 36 additions & 8 deletions src/impl_alloc.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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)
})
}
}
}
Expand Down Expand Up @@ -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()
}
}

Expand Down
15 changes: 10 additions & 5 deletions src/impl_serde.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand Down Expand Up @@ -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);
Expand All @@ -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))
Expand Down
97 changes: 97 additions & 0 deletions src/internal.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#![allow(dead_code)] // ArrayBuilder is soft-deprecated internally

use crate::*;

pub trait Sealed {}
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down
24 changes: 12 additions & 12 deletions src/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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());
}
}
30 changes: 19 additions & 11 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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],
//! }
Expand Down Expand Up @@ -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")]
Expand All @@ -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::*;
Expand Down Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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)
})
}
}
}
Expand Down

0 comments on commit 13ddbf3

Please sign in to comment.