Skip to content

Commit

Permalink
Implement as_flattened{_mut}() for Array<Array<>> (#86)
Browse files Browse the repository at this point in the history
Implementation of `as_flattened` and `as_flattened_mut` for nested
`Array` types to make it easier to operate on nested `Array`s.
Inspiration drawn from the `as_flattened{_mut}` implementations on `[[T;
N]]` (i.e. for the [slice
primitive](https://doc.rust-lang.org/core/primitive.slice.html#method.as_flattened)).

Note: 
* `Array<[T; N], U>` already implements these functions as it is
transparent to `&[[T; N]]`.
* `[Array<T, U>]` does not implement these functions and we can't
implement them this way without a new local trait. Doesn't really seem
necessary though.

Co-authored-by: jmwample <8297368+jmwample@users.noreply.github.com>
  • Loading branch information
jmwample and jmwample authored Sep 12, 2024
1 parent 49dfead commit 4e4c10e
Showing 1 changed file with 90 additions and 0 deletions.
90 changes: 90 additions & 0 deletions src/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
use crate::{Array, ArraySize};
use core::{
fmt,
mem::size_of,
slice::{Iter, IterMut},
};

Expand Down Expand Up @@ -105,3 +106,92 @@ where
self.iter_mut()
}
}

impl<T, U, V> Array<Array<T, U>, V>
where
U: ArraySize,
V: ArraySize,
{
/// Takes a `&mut Array<Array<T, N>,M>`, and flattens it to a `&mut [T]`.
///
/// # Panics
///
/// This panics if the length of the resulting slice would overflow a `usize`.
///
/// This is only possible when flattening a slice of arrays of zero-sized
/// types, and thus tends to be irrelevant in practice. If
/// `size_of::<T>() > 0`, this will never panic.
///
/// # Examples
///
/// ```
/// use hybrid_array::{Array, typenum::U3};
///
/// fn add_5_to_all(slice: &mut [i32]) {
/// for i in slice {
/// *i += 5;
/// }
/// }
///
/// let mut array: Array<Array<i32, U3>, U3> = Array([Array([1_i32, 2, 3]), Array([4, 5, 6]), Array([7, 8, 9])]);
/// add_5_to_all(array.as_flattened_mut());
/// assert_eq!(array, Array([Array([6, 7, 8]), Array([9, 10, 11]), Array([12, 13, 14])]));
/// ```
pub fn as_flattened_mut(&mut self) -> &mut [T] {
let len = if size_of::<T>() == 0 {
self.len()
.checked_mul(U::USIZE)
.expect("slice len overflow")
} else {
// SAFETY: `self.len() * N` cannot overflow because `self` is
// already in the address space.
unsafe { self.len().unchecked_mul(U::USIZE) }
};
// SAFETY: `[T]` is layout-identical to `[T; U]`
unsafe { core::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), len) }
}

/// Takes a `&Array<Array<T, N>, >>`, and flattens it to a `&[T]`.
///
/// # Panics
///
/// This panics if the length of the resulting slice would overflow a `usize`.
///
/// This is only possible when flattening a slice of arrays of zero-sized
/// types, and thus tends to be irrelevant in practice. If
/// `size_of::<T>() > 0`, this will never panic.
///
/// # Examples
///
/// ```
/// use hybrid_array::{Array, typenum::{U0, U2, U3, U5, U10}};
///
/// let a: Array<Array<usize, U3>, U2> = Array([Array([1, 2, 3]), Array([4, 5, 6])]);
/// assert_eq!(a.as_flattened(), &[1, 2, 3, 4, 5, 6]);
///
/// let b: Array<Array<usize, U2>, U3> = Array([Array([1, 2]), Array([3, 4]), Array([5, 6])]);
/// assert_eq!(a.as_flattened(), b.as_flattened());
///
/// let c: Array<[usize; 2], U3> = Array([[1, 2], [3, 4], [5, 6]]);
/// assert_eq!(a.as_flattened(), c.as_flattened());
///
/// let slice_of_empty_arrays: &Array<Array<i32, U5>, U0> = &Array::from_fn(|_| Array([1, 2, 3, 4, 5]));
/// assert!(slice_of_empty_arrays.as_flattened().is_empty());
///
/// let empty_slice_of_arrays: &Array<Array<u32, U10>, U0> = &Array([]);
/// assert!(empty_slice_of_arrays.as_flattened().is_empty());
/// ```
pub fn as_flattened(&self) -> &[T] {
let len = if size_of::<T>() == 0 {
self.len()
.checked_mul(U::USIZE)
.expect("slice len overflow")
} else {
// SAFETY: `self.len() * N` cannot overflow because `self` is
// already in the address space.
unsafe { self.len().unchecked_mul(U::USIZE) }
};
// SAFETY: `[T]` is layout-identical to `[T; U]`
unsafe { core::slice::from_raw_parts(self.as_ptr().cast(), len) }
}
}

0 comments on commit 4e4c10e

Please sign in to comment.