Skip to content

Commit

Permalink
Merge pull request #1389 from akern40/array_macro
Browse files Browse the repository at this point in the history
Expands the `array!` macro.
  • Loading branch information
bluss authored Jul 23, 2024
2 parents 024dccb + e2facb7 commit 6cb2368
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 56 deletions.
16 changes: 16 additions & 0 deletions src/dimension/ndindex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,22 @@ unsafe impl NdIndex<Ix5> for (Ix, Ix, Ix, Ix, Ix)
}
}

unsafe impl NdIndex<Ix6> for (Ix, Ix, Ix, Ix, Ix, Ix)
{
#[inline]
fn index_checked(&self, dim: &Ix6, strides: &Ix6) -> Option<isize>
{
dim.stride_offset_checked(strides, &self.into_dimension())
}
#[inline]
fn index_unchecked(&self, strides: &Ix6) -> isize
{
zip(strides.ix(), self.into_dimension().ix())
.map(|(&s, &i)| stride_offset(i, s))
.sum()
}
}

unsafe impl NdIndex<Ix1> for Ix
{
#[inline]
Expand Down
121 changes: 65 additions & 56 deletions src/free_functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,15 @@
use alloc::vec;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[allow(unused_imports)]
use std::compile_error;
use std::mem::{forget, size_of};
use std::ptr::NonNull;

use crate::imp_prelude::*;
use crate::{dimension, ArcArray1, ArcArray2};

/// Create an **[`Array`]** with one, two or
/// three dimensions.
/// Create an **[`Array`]** with one, two, three, four, five, or six dimensions.
///
/// ```
/// use ndarray::array;
Expand All @@ -28,17 +29,49 @@ use crate::{dimension, ArcArray1, ArcArray2};
/// let a3 = array![[[1, 2], [3, 4]],
/// [[5, 6], [7, 8]]];
///
/// let a4 = array![[[[1, 2, 3, 4]]]];
///
/// let a5 = array![[[[[1, 2, 3, 4, 5]]]]];
///
/// let a6 = array![[[[[[1, 2, 3, 4, 5, 6]]]]]];
///
/// assert_eq!(a1.shape(), &[4]);
/// assert_eq!(a2.shape(), &[2, 2]);
/// assert_eq!(a3.shape(), &[2, 2, 2]);
/// assert_eq!(a4.shape(), &[1, 1, 1, 4]);
/// assert_eq!(a5.shape(), &[1, 1, 1, 1, 5]);
/// assert_eq!(a6.shape(), &[1, 1, 1, 1, 1, 6]);
/// ```
///
/// This macro uses `vec![]`, and has the same ownership semantics;
/// elements are moved into the resulting `Array`.
///
/// Use `array![...].into_shared()` to create an `ArcArray`.
///
/// Attempts to crate 7D+ arrays with this macro will lead to
/// a compiler error, since the difference between a 7D array
/// of i32 and a 6D array of `[i32; 3]` is ambiguous. Higher-dim
/// arrays can be created with [`ArrayD`].
///
/// ```compile_fail
/// use ndarray::array;
/// let a7 = array![[[[[[[1, 2, 3]]]]]]];
/// // error: Arrays of 7 dimensions or more (or ndarrays of Rust arrays) cannot be constructed with the array! macro.
/// ```
#[macro_export]
macro_rules! array {
($([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
compile_error!("Arrays of 7 dimensions or more (or ndarrays of Rust arrays) cannot be constructed with the array! macro.");
}};
($([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
$crate::Array6::from(vec![$([$([$([$([$([$($x,)*],)*],)*],)*],)*],)*])
}};
($([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
$crate::Array5::from(vec![$([$([$([$([$($x,)*],)*],)*],)*],)*])
}};
($([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
$crate::Array4::from(vec![$([$([$([$($x,)*],)*],)*],)*])
}};
($([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*) => {{
$crate::Array3::from(vec![$([$([$($x,)*],)*],)*])
}};
Expand Down Expand Up @@ -233,63 +266,39 @@ pub fn arr2<A: Clone, const N: usize>(xs: &[[A; N]]) -> Array2<A>
Array2::from(xs.to_vec())
}

impl<A, const N: usize> From<Vec<[A; N]>> for Array2<A>
{
/// Converts the `Vec` of arrays to an owned 2-D array.
///
/// **Panics** if the product of non-zero axis lengths overflows `isize`.
fn from(mut xs: Vec<[A; N]>) -> Self
{
let dim = Ix2(xs.len(), N);
let ptr = xs.as_mut_ptr();
let cap = xs.capacity();
let expand_len =
dimension::size_of_shape_checked(&dim).expect("Product of non-zero axis lengths must not overflow isize.");
forget(xs);
unsafe {
let v = if size_of::<A>() == 0 {
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
} else if N == 0 {
Vec::new()
} else {
// Guaranteed not to overflow in this case since A is non-ZST
// and Vec never allocates more than isize bytes.
let expand_cap = cap * N;
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_cap)
};
ArrayBase::from_shape_vec_unchecked(dim, v)
macro_rules! impl_from_nested_vec {
($arr_type:ty, $ix_type:tt, $($n:ident),+) => {
impl<A, $(const $n: usize),+> From<Vec<$arr_type>> for Array<A, $ix_type>
{
fn from(mut xs: Vec<$arr_type>) -> Self
{
let dim = $ix_type(xs.len(), $($n),+);
let ptr = xs.as_mut_ptr();
let cap = xs.capacity();
let expand_len = dimension::size_of_shape_checked(&dim)
.expect("Product of non-zero axis lengths must not overflow isize.");
forget(xs);
unsafe {
let v = if size_of::<A>() == 0 {
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
} else if $($n == 0 ||)+ false {
Vec::new()
} else {
let expand_cap = cap $(* $n)+;
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_cap)
};
ArrayBase::from_shape_vec_unchecked(dim, v)
}
}
}
}
};
}

impl<A, const N: usize, const M: usize> From<Vec<[[A; M]; N]>> for Array3<A>
{
/// Converts the `Vec` of arrays to an owned 3-D array.
///
/// **Panics** if the product of non-zero axis lengths overflows `isize`.
fn from(mut xs: Vec<[[A; M]; N]>) -> Self
{
let dim = Ix3(xs.len(), N, M);
let ptr = xs.as_mut_ptr();
let cap = xs.capacity();
let expand_len =
dimension::size_of_shape_checked(&dim).expect("Product of non-zero axis lengths must not overflow isize.");
forget(xs);
unsafe {
let v = if size_of::<A>() == 0 {
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
} else if N == 0 || M == 0 {
Vec::new()
} else {
// Guaranteed not to overflow in this case since A is non-ZST
// and Vec never allocates more than isize bytes.
let expand_cap = cap * N * M;
Vec::from_raw_parts(ptr as *mut A, expand_len, expand_cap)
};
ArrayBase::from_shape_vec_unchecked(dim, v)
}
}
}
impl_from_nested_vec!([A; N], Ix2, N);
impl_from_nested_vec!([[A; M]; N], Ix3, N, M);
impl_from_nested_vec!([[[A; L]; M]; N], Ix4, N, M, L);
impl_from_nested_vec!([[[[A; K]; L]; M]; N], Ix5, N, M, L, K);
impl_from_nested_vec!([[[[[A; J]; K]; L]; M]; N], Ix6, N, M, L, K, J);

/// Create a two-dimensional array with elements from `xs`.
///
Expand Down

0 comments on commit 6cb2368

Please sign in to comment.