Merge pull request #1389 from akern40/array_macro

Expands the `array!` macro.

Author: bluss

src/dimension/ndindex.rs

@@ -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]

src/free_functions.rs

@@ -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;
@@ -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,)*],)*],)*])
     }};
@@ -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`.
 ///