Merge pull request #980 from YuhanLiin/const-generics

Const generics Improvements

Author: bluss

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
           - stable
           - beta
           - nightly
-          - 1.49.0  # MSRV
+          - 1.51.0  # MSRV
 
     steps:
       - uses: actions/checkout@v2

src/arraytraits.rs

@@ -6,20 +6,20 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::hash;
-use std::iter::FromIterator;
+use alloc::boxed::Box;
+use alloc::vec::Vec;
 use std::iter::IntoIterator;
 use std::mem;
 use std::ops::{Index, IndexMut};
-use alloc::boxed::Box;
-use alloc::vec::Vec;
+use std::{hash, mem::size_of};
+use std::{iter::FromIterator, slice};
 
 use crate::imp_prelude::*;
-use crate::iter::{Iter, IterMut};
-use crate::NdIndex;
-
-use crate::numeric_util;
-use crate::{FoldWhile, Zip};
+use crate::{
+    dimension,
+    iter::{Iter, IterMut},
+    numeric_util, FoldWhile, NdIndex, Zip,
+};
 
 #[cold]
 #[inline(never)]
@@ -323,6 +323,30 @@ where
     }
 }
 
+/// Implementation of ArrayView2::from(&S) where S is a slice to a 2D array
+///
+/// **Panics** if the product of non-zero axis lengths overflows `isize` (This can only occur if A
+/// is zero-sized because slices cannot contain more than `isize::MAX` number of bytes).
+impl<'a, A, const N: usize> From<&'a [[A; N]]> for ArrayView<'a, A, Ix2> {
+    /// Create a two-dimensional read-only array view of the data in `slice`
+    fn from(xs: &'a [[A; N]]) -> Self {
+        let cols = N;
+        let rows = xs.len();
+        let dim = Ix2(rows, cols);
+        if size_of::<A>() == 0 {
+            dimension::size_of_shape_checked(&dim)
+                .expect("Product of non-zero axis lengths must not overflow isize.");
+        }
+
+        // `cols * rows` is guaranteed to fit in `isize` because we checked that it fits in
+        // `isize::MAX`
+        unsafe {
+            let data = slice::from_raw_parts(xs.as_ptr() as *const A, cols * rows);
+            ArrayView::from_shape_ptr(dim, data.as_ptr())
+        }
+    }
+}
+
 /// Implementation of `ArrayView::from(&A)` where `A` is an array.
 impl<'a, A, S, D> From<&'a ArrayBase<S, D>> for ArrayView<'a, A, D>
 where
@@ -355,6 +379,30 @@ where
     }
 }
 
+/// Implementation of ArrayViewMut2::from(&S) where S is a slice to a 2D array
+///
+/// **Panics** if the product of non-zero axis lengths overflows `isize` (This can only occur if A
+/// is zero-sized because slices cannot contain more than `isize::MAX` number of bytes).
+impl<'a, A, const N: usize> From<&'a mut [[A; N]]> for ArrayViewMut<'a, A, Ix2> {
+    /// Create a two-dimensional read-write array view of the data in `slice`
+    fn from(xs: &'a mut [[A; N]]) -> Self {
+        let cols = N;
+        let rows = xs.len();
+        let dim = Ix2(rows, cols);
+        if size_of::<A>() == 0 {
+            dimension::size_of_shape_checked(&dim)
+                .expect("Product of non-zero axis lengths must not overflow isize.");
+        }
+
+        // `cols * rows` is guaranteed to fit in `isize` because we checked that it fits in
+        // `isize::MAX`
+        unsafe {
+            let data = slice::from_raw_parts_mut(xs.as_mut_ptr() as *mut A, cols * rows);
+            ArrayViewMut::from_shape_ptr(dim, data.as_mut_ptr())
+        }
+    }
+}
+
 /// Implementation of `ArrayViewMut::from(&mut A)` where `A` is an array.
 impl<'a, A, S, D> From<&'a mut ArrayBase<S, D>> for ArrayViewMut<'a, A, D>
 where

src/dimension/ndindex.rs

@@ -140,50 +140,6 @@ macro_rules! ndindex_with_array {
                 0
             }
         }
-
-        // implement NdIndex<IxDyn> for Dim<[Ix; 2]> and so on
-        unsafe impl NdIndex<IxDyn> for Dim<[Ix; $n]> {
-            #[inline]
-            fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
-                debug_assert_eq!(strides.ndim(), $n,
-                              "Attempted to index with {:?} in array with {} axes",
-                              self, strides.ndim());
-                stride_offset_checked(dim.ix(), strides.ix(), self.ix())
-            }
-
-            #[inline]
-            fn index_unchecked(&self, strides: &IxDyn) -> isize {
-                debug_assert_eq!(strides.ndim(), $n,
-                              "Attempted to index with {:?} in array with {} axes",
-                              self, strides.ndim());
-                $(
-                stride_offset(get!(self, $index), get!(strides, $index)) +
-                )*
-                0
-            }
-        }
-
-        // implement NdIndex<IxDyn> for [Ix; 2] and so on
-        unsafe impl NdIndex<IxDyn> for [Ix; $n] {
-            #[inline]
-            fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
-                debug_assert_eq!(strides.ndim(), $n,
-                              "Attempted to index with {:?} in array with {} axes",
-                              self, strides.ndim());
-                stride_offset_checked(dim.ix(), strides.ix(), self)
-            }
-
-            #[inline]
-            fn index_unchecked(&self, strides: &IxDyn) -> isize {
-                debug_assert_eq!(strides.ndim(), $n,
-                              "Attempted to index with {:?} in array with {} axes",
-                              self, strides.ndim());
-                $(
-                stride_offset(self[$index], get!(strides, $index)) +
-                )*
-                0
-            }
-        }
         )+
     };
 }
@@ -198,6 +154,64 @@ ndindex_with_array! {
     [6, Ix6 0 1 2 3 4 5]
 }
 
+// implement NdIndex<IxDyn> for Dim<[Ix; 2]> and so on
+unsafe impl<const N: usize> NdIndex<IxDyn> for Dim<[Ix; N]> {
+    #[inline]
+    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
+        debug_assert_eq!(
+            strides.ndim(),
+            N,
+            "Attempted to index with {:?} in array with {} axes",
+            self,
+            strides.ndim()
+        );
+        stride_offset_checked(dim.ix(), strides.ix(), self.ix())
+    }
+
+    #[inline]
+    fn index_unchecked(&self, strides: &IxDyn) -> isize {
+        debug_assert_eq!(
+            strides.ndim(),
+            N,
+            "Attempted to index with {:?} in array with {} axes",
+            self,
+            strides.ndim()
+        );
+        (0..N)
+            .map(|i| stride_offset(get!(self, i), get!(strides, i)))
+            .sum()
+    }
+}
+
+// implement NdIndex<IxDyn> for [Ix; 2] and so on
+unsafe impl<const N: usize> NdIndex<IxDyn> for [Ix; N] {
+    #[inline]
+    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
+        debug_assert_eq!(
+            strides.ndim(),
+            N,
+            "Attempted to index with {:?} in array with {} axes",
+            self,
+            strides.ndim()
+        );
+        stride_offset_checked(dim.ix(), strides.ix(), self)
+    }
+
+    #[inline]
+    fn index_unchecked(&self, strides: &IxDyn) -> isize {
+        debug_assert_eq!(
+            strides.ndim(),
+            N,
+            "Attempted to index with {:?} in array with {} axes",
+            self,
+            strides.ndim()
+        );
+        (0..N)
+            .map(|i| stride_offset(self[i], get!(strides, i)))
+            .sum()
+    }
+}
+
 impl<'a> IntoDimension for &'a [Ix] {
     type Dim = IxDyn;
     fn into_dimension(self) -> Self::Dim {

src/free_functions.rs

@@ -6,10 +6,9 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::mem::{forget, size_of};
-use alloc::slice;
 use alloc::vec;
 use alloc::vec::Vec;
+use std::mem::{forget, size_of};
 
 use crate::imp_prelude::*;
 use crate::{dimension, ArcArray1, ArcArray2};
@@ -87,26 +86,10 @@ pub fn aview1<A>(xs: &[A]) -> ArrayView1<'_, A> {
 
 /// Create a two-dimensional array view with elements borrowing `xs`.
 ///
-/// **Panics** if the product of non-zero axis lengths overflows `isize`. (This
-/// can only occur when `V` is zero-sized.)
-pub fn aview2<A, V: FixedInitializer<Elem = A>>(xs: &[V]) -> ArrayView2<'_, A> {
-    let cols = V::len();
-    let rows = xs.len();
-    let dim = Ix2(rows, cols);
-    if size_of::<V>() == 0 {
-        dimension::size_of_shape_checked(&dim)
-            .expect("Product of non-zero axis lengths must not overflow isize.");
-    }
-    // `rows` is guaranteed to fit in `isize` because we've checked the ZST
-    // case and slices never contain > `isize::MAX` bytes. `cols` is guaranteed
-    // to fit in `isize` because `FixedInitializer` is not implemented for any
-    // array lengths > `isize::MAX`. `cols * rows` is guaranteed to fit in
-    // `isize` because we've checked the ZST case and slices never contain >
-    // `isize::MAX` bytes.
-    unsafe {
-        let data = slice::from_raw_parts(xs.as_ptr() as *const A, cols * rows);
-        ArrayView::from_shape_ptr(dim, data.as_ptr())
-    }
+/// **Panics** if the product of non-zero axis lengths overflows `isize` (This can only occur if A
+/// is zero-sized because slices cannot contain more than `isize::MAX` number of bytes).
+pub fn aview2<A, const N: usize>(xs: &[[A; N]]) -> ArrayView2<'_, A> {
+    ArrayView2::from(xs)
 }
 
 /// Create a one-dimensional read-write array view with elements borrowing `xs`.
@@ -127,16 +110,15 @@ pub fn aview_mut1<A>(xs: &mut [A]) -> ArrayViewMut1<'_, A> {
 
 /// Create a two-dimensional read-write array view with elements borrowing `xs`.
 ///
-/// **Panics** if the product of non-zero axis lengths overflows `isize`. (This
-/// can only occur when `V` is zero-sized.)
+/// **Panics** if the product of non-zero axis lengths overflows `isize` (This can only occur if A
+/// is zero-sized because slices cannot contain more than `isize::MAX` number of bytes).
 ///
 /// # Example
 ///
 /// ```
 /// use ndarray::aview_mut2;
 ///
-/// // The inner (nested) array must be of length 1 to 16, but the outer
-/// // can be of any length.
+/// // The inner (nested) and outer arrays can be of any length.
 /// let mut data = [[0.; 2]; 128];
 /// {
 ///     // Make a 128 x 2 mut array view then turn it into 2 x 128
@@ -148,57 +130,10 @@ pub fn aview_mut1<A>(xs: &mut [A]) -> ArrayViewMut1<'_, A> {
 /// // look at the start of the result
 /// assert_eq!(&data[..3], [[1., -1.], [1., -1.], [1., -1.]]);
 /// ```
-pub fn aview_mut2<A, V: FixedInitializer<Elem = A>>(xs: &mut [V]) -> ArrayViewMut2<'_, A> {
-    let cols = V::len();
-    let rows = xs.len();
-    let dim = Ix2(rows, cols);
-    if size_of::<V>() == 0 {
-        dimension::size_of_shape_checked(&dim)
-            .expect("Product of non-zero axis lengths must not overflow isize.");
-    }
-    // `rows` is guaranteed to fit in `isize` because we've checked the ZST
-    // case and slices never contain > `isize::MAX` bytes. `cols` is guaranteed
-    // to fit in `isize` because `FixedInitializer` is not implemented for any
-    // array lengths > `isize::MAX`. `cols * rows` is guaranteed to fit in
-    // `isize` because we've checked the ZST case and slices never contain >
-    // `isize::MAX` bytes.
-    unsafe {
-        let data = slice::from_raw_parts_mut(xs.as_mut_ptr() as *mut A, cols * rows);
-        ArrayViewMut::from_shape_ptr(dim, data.as_mut_ptr())
-    }
+pub fn aview_mut2<A, const N: usize>(xs: &mut [[A; N]]) -> ArrayViewMut2<'_, A> {
+    ArrayViewMut2::from(xs)
 }
 
-/// Fixed-size array used for array initialization
-#[allow(clippy::missing_safety_doc)] // Should not be implemented downstream and to be deprecated.
-pub unsafe trait FixedInitializer {
-    type Elem;
-    fn as_init_slice(&self) -> &[Self::Elem];
-    fn len() -> usize;
-}
-
-macro_rules! impl_arr_init {
-    (__impl $n: expr) => (
-        unsafe impl<T> FixedInitializer for [T;  $n] {
-            type Elem = T;
-            fn as_init_slice(&self) -> &[T] { self }
-            fn len() -> usize { $n }
-        }
-    );
-    () => ();
-    ($n: expr, $($m:expr,)*) => (
-        impl_arr_init!(__impl $n);
-        impl_arr_init!($($m,)*);
-    )
-
-}
-
-// For implementors: If you ever implement `FixedInitializer` for array lengths
-// > `isize::MAX` (e.g. once Rust adds const generics), you must update
-// `aview2` and `aview_mut2` to perform the necessary checks. In particular,
-// the assumption that `cols` can never exceed `isize::MAX` would be incorrect.
-// (Consider e.g. `let xs: &[[i32; ::std::usize::MAX]] = &[]`.)
-impl_arr_init!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,);
-
 /// Create a two-dimensional array with elements from `xs`.
 ///
 /// ```
@@ -210,22 +145,16 @@ impl_arr_init!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,);
 ///     a.shape() == [2, 3]
 /// );
 /// ```
-pub fn arr2<A: Clone, V: FixedInitializer<Elem = A>>(xs: &[V]) -> Array2<A>
-where
-    V: Clone,
-{
+pub fn arr2<A: Clone, const N: usize>(xs: &[[A; N]]) -> Array2<A> {
     Array2::from(xs.to_vec())
 }
 
-impl<A, V> From<Vec<V>> for Array2<A>
-where
-    V: FixedInitializer<Elem = A>,
-{
+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<V>) -> Self {
-        let dim = Ix2(xs.len(), V::len());
+    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)
@@ -234,29 +163,25 @@ where
         unsafe {
             let v = if size_of::<A>() == 0 {
                 Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
-            } else if V::len() == 0 {
+            } 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 * V::len();
+                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, V, U> From<Vec<V>> for Array3<A>
-where
-    V: FixedInitializer<Elem = U>,
-    U: FixedInitializer<Elem = A>,
-{
+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<V>) -> Self {
-        let dim = Ix3(xs.len(), V::len(), U::len());
+    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)
@@ -265,12 +190,12 @@ where
         unsafe {
             let v = if size_of::<A>() == 0 {
                 Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
-            } else if V::len() == 0 || U::len() == 0 {
+            } 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 * V::len() * U::len();
+                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)
@@ -280,7 +205,7 @@ where
 
 /// Create a two-dimensional array with elements from `xs`.
 ///
-pub fn rcarr2<A: Clone, V: Clone + FixedInitializer<Elem = A>>(xs: &[V]) -> ArcArray2<A> {
+pub fn rcarr2<A: Clone, const N: usize>(xs: &[[A; N]]) -> ArcArray2<A> {
     arr2(xs).into_shared()
 }
 
@@ -301,23 +226,11 @@ pub fn rcarr2<A: Clone, V: Clone + FixedInitializer<Elem = A>>(xs: &[V]) -> ArcA
 ///     a.shape() == [3, 2, 2]
 /// );
 /// ```
-pub fn arr3<A: Clone, V: FixedInitializer<Elem = U>, U: FixedInitializer<Elem = A>>(
-    xs: &[V],
-) -> Array3<A>
-where
-    V: Clone,
-    U: Clone,
-{
+pub fn arr3<A: Clone, const N: usize, const M: usize>(xs: &[[[A; M]; N]]) -> Array3<A> {
     Array3::from(xs.to_vec())
 }
 
 /// Create a three-dimensional array with elements from `xs`.
-pub fn rcarr3<A: Clone, V: FixedInitializer<Elem = U>, U: FixedInitializer<Elem = A>>(
-    xs: &[V],
-) -> ArcArray<A, Ix3>
-where
-    V: Clone,
-    U: Clone,
-{
+pub fn rcarr3<A: Clone, const N: usize, const M: usize>(xs: &[[[A; M]; N]]) -> ArcArray<A, Ix3> {
     arr3(xs).into_shared()
 }

src/zip/ndproducer.rs

@@ -1,4 +1,3 @@
-
 use crate::imp_prelude::*;
 use crate::Layout;
 use crate::NdIndex;
@@ -168,6 +167,26 @@ impl<'a, A: 'a> IntoNdProducer for &'a mut [A] {
     }
 }
 
+/// A one-dimensional array is a one-dimensional producer
+impl<'a, A: 'a, const N: usize> IntoNdProducer for &'a [A; N] {
+    type Item = <Self::Output as NdProducer>::Item;
+    type Dim = Ix1;
+    type Output = ArrayView1<'a, A>;
+    fn into_producer(self) -> Self::Output {
+        <_>::from(self)
+    }
+}
+
+/// A mutable one-dimensional array is a mutable one-dimensional producer
+impl<'a, A: 'a, const N: usize> IntoNdProducer for &'a mut [A; N] {
+    type Item = <Self::Output as NdProducer>::Item;
+    type Dim = Ix1;
+    type Output = ArrayViewMut1<'a, A>;
+    fn into_producer(self) -> Self::Output {
+        <_>::from(self)
+    }
+}
+
 /// A Vec is a one-dimensional producer
 impl<'a, A: 'a> IntoNdProducer for &'a Vec<A> {
     type Item = <Self::Output as NdProducer>::Item;
@@ -399,4 +418,3 @@ impl<A, D: Dimension> NdProducer for RawArrayViewMut<A, D> {
         self.split_at(axis, index)
     }
 }
-

tests/array.rs

@@ -731,7 +731,7 @@ fn diag() {
     let a = arr2(&[[1., 2., 3.0f32], [0., 0., 0.]]);
     let d = a.view().into_diag();
     assert_eq!(d.dim(), 2);
-    let d = arr2::<f32, _>(&[[]]).into_diag();
+    let d = arr2::<f32, 0>(&[[]]).into_diag();
     assert_eq!(d.dim(), 0);
     let d = ArcArray::<f32, _>::zeros(()).into_diag();
     assert_eq!(d.dim(), 1);
@@ -960,7 +960,7 @@ fn zero_axes() {
     a.map_inplace(|_| panic!());
     a.for_each(|_| panic!());
     println!("{:?}", a);
-    let b = arr2::<f32, _>(&[[], [], [], []]);
+    let b = arr2::<f32, 0>(&[[], [], [], []]);
     println!("{:?}\n{:?}", b.shape(), b);
 
     // we can even get a subarray of b