diff --git a/src/arraylike.rs b/src/arraylike.rs index 1bef17e8e..cf0c51df6 100644 --- a/src/arraylike.rs +++ b/src/arraylike.rs @@ -22,46 +22,161 @@ use crate::{ /// and other types that conceptually act like arrays. It's designed to make your functions /// more flexible by letting them handle a wide range of types without extra boilerplate. /// -/// ## More Details -/// The key idea of `ArrayLike` is to bridge the gap between native ndarray types -/// (e.g., [`Array2`](crate::Array2), [`ArrayView`](crate::ArrayView)) and other -/// data structures like scalars or slices. It enables treating all these types -/// as "array-like" objects with a common set of operations. +/// Like other parts of the `ndarray` crate, `ArrayLike` only works with scalars that implement +/// [`ScalarOperand`]. /// /// # Example /// ``` -/// use ndarray::{array, Array, Array2, ArrayLike, DimMax}; +/// use core::ops::Mul; +/// use ndarray::{array, Array, ArrayLike, DimMax}; /// -/// fn multiply(left: &T, right: &G) -> Array>::Output> +/// fn multiply(left: T, right: G) -> Array>::Output> /// where -/// T: ArrayLike, -/// G: ArrayLike, +/// T: ArrayLike, +/// G: ArrayLike, +/// // Bounds to enable multiplication +/// T::Elem: Clone + Mul, +/// G::Elem: Clone, /// T::Dim: DimMax, -/// A: Mul, /// { -/// left.as_array() * right.as_array() +/// &left.view() * &right.view() /// } /// -/// let rows = array![[1], [2]]; -/// let cols = array![3, 4]; -/// assert_eq!(multiply(rows, col), array![[3, 4], [6, 8]]); +/// let left = array![1, 2]; +/// let right = vec![3, 4]; +/// // Array-vector multiplication +/// assert_eq!(multiply(&left, &right), array![3, 8]); +/// // Array-scalar multiplication +/// assert_eq!(multiply(&left, 3), array![3, 6]); /// ``` +/// +/// # `ArrayLike` vs [`ArrayRef`] +/// Both `ArrayLike` and `ArrayRef` provide a kind of unifying abstraction for `ndarray`, +/// and both are useful for writing functions with `ndarray`. However, they should not be +/// used interchangeably. `ArrayLike` is ideal when you want to write generic functions +/// that work with anything that "looks" like a multidimensional array, even if it isn't +/// strictly an `ndarray` type. When you know that a given variable or argument will be an +/// `ndarray` type, use `ArrayRef` instead. pub trait ArrayLike { + /// The dimensionality of the underlying array-like data structure. type Dim: Dimension; + + /// The element type of the underlying array-like data structure. type Elem; + /// Get a read-only view of the underlying array-like data structure. + /// + /// This method should never re-allocate the underlying data. fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim>; + /// Get the shape and strides of the underlying array-like data structure. + /// + /// # Example fn dim(&self) -> Self::Dim; + /// If the underlying object only has one element, return it; otherwise return `None`. + /// + /// This method allows for optimizations when the `ArrayLike` value is actually a + /// scalar, in which case one can avoid allocations and aid the compiler by not + /// turning it into a full view. + /// + /// # Example + /// ```rust + /// use ndarray::{array, ArrayLike}; + /// + /// let arr = array![1, 2, 3]; + /// let arr_single = array![1]; + /// let scalar = 1; + /// + /// matches!(arr.as_elem(), None); + /// matches!(arr.as_elem(), Some(1)); + /// matches!(scalar.as_elem(), Some(1)); + /// ``` + /// + /// # For Implementors: + /// Array-like objects that can contain multiple elements are free to return `Some(_)` + /// if and only if a runtime check determines there is only one element in the container. fn as_elem(&self) -> Option<&Self::Elem>; } +/// A trait for mutable array-like objects. +/// +/// This extends [`ArrayLike`] by providing mutable access to the underlying data. +/// Use it when you need to modify the contents of an array-like object. +/// +/// ## More Details +/// `ArrayLikeMut` is designed for types that can provide mutable access to their elements. +/// For example, mutable slices and arrays implement this trait, but immutable views or +/// read-only data structures won't. +/// +/// # Examples +/// ``` +/// use core::ops::MulAssign; +/// use ndarray::{array, ArrayLike, ArrayLikeMut, DimMax}; +/// +/// fn multiply_assign(left: &mut T, right: &G) +/// where +/// T: ArrayLikeMut, +/// G: ArrayLike, +/// // Bounds to enable multiplication +/// T::Elem: Clone + MulAssign, +/// G::Elem: Clone, +/// // Ensure that the broadcast is still assignable to the left side +/// T::Dim: DimMax, +/// { +/// *left.view_mut() *= &right.view(); +/// } +/// +/// let mut left = array![1, 2]; +/// let right = array![3, 4]; +/// +/// multiply_assign(&mut left, &right); +/// assert_eq!(left, array![3, 8]); +/// +/// multiply_assign(&mut left, &2); +/// assert_eq!(left, array![6, 16]); +/// ``` pub trait ArrayLikeMut: ArrayLike { + /// Get a mutable view of the underlying array-like data structure. + /// + /// This method should never re-allocate the underlying data. fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim>; + /// If the underlying object only has one element, return a mutable reference; otherwise return `None`. + /// + /// This method allows for optimizations when the `ArrayLike` value is actually a + /// scalar, in which case one can avoid allocations and aid the compiler by not + /// turning it into a full view. + /// + /// # Example + /// ```rust + /// use ndarray::{array, ArrayLike, ArrayLikeMut}; + /// use num_traits::Zero; + /// + /// fn assign_sum(mut left: T, right: G) + /// where + /// T: ArrayLikeMut, + /// G: ArrayLike, + /// // Bounds to enable sum + /// T::Elem: Zero + Clone, + /// { + /// if let Some(e) = left.as_elem_mut() { + /// *e = right.view().sum(); + /// } + /// } + /// + /// + /// let arr = array![1, 2, 3]; + /// let mut arr_single = array![1]; + /// assign_sum(&mut arr_single, arr); + /// assert_eq!(arr_single[0], 6); + /// ``` + /// + /// # For Implementors: + /// Array-like objects that can contain multiple elements are free to return `Some(_)` + /// if and only if a runtime check determines there is only one element in the container. fn as_elem_mut(&mut self) -> Option<&mut Self::Elem>; } @@ -130,6 +245,58 @@ where D: Dimension } } +impl ArrayLike for &ArrayRef +where D: Dimension +{ + type Dim = D; + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + (*self).view() + } + + fn dim(&self) -> Self::Dim + { + self.raw_dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.dim.size() == 1 { + self.first() + } else { + None + } + } +} + +impl ArrayLike for &mut ArrayRef +where D: Dimension +{ + type Dim = D; + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + (**self).view() + } + + fn dim(&self) -> Self::Dim + { + self.raw_dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.dim.size() == 1 { + self.first() + } else { + None + } + } +} + impl ArrayLikeMut for ArrayRef where D: Dimension { @@ -148,6 +315,24 @@ where D: Dimension } } +impl ArrayLikeMut for &mut ArrayRef +where D: Dimension +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + ArrayRef::view_mut(self) + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + if self.dim.size() == 1 { + self.first_mut() + } else { + None + } + } +} + impl ArrayLike for ArrayBase where S: Data, @@ -158,7 +343,63 @@ where fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> { - self.view() + ArrayRef::view(self) + } + + fn dim(&self) -> Self::Dim + { + self.raw_dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.dim.size() == 1 { + self.first() + } else { + None + } + } +} + +impl ArrayLike for &ArrayBase +where + S: Data, + D: Dimension, +{ + type Dim = D; + type Elem = S::Elem; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + ArrayRef::view(self) + } + + fn dim(&self) -> Self::Dim + { + self.raw_dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.dim.size() == 1 { + self.first() + } else { + None + } + } +} + +impl ArrayLike for &mut ArrayBase +where + S: Data, + D: Dimension, +{ + type Dim = D; + type Elem = S::Elem; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + ArrayRef::view(self) } fn dim(&self) -> Self::Dim @@ -183,7 +424,27 @@ where { fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> { - self.view_mut() + ArrayRef::view_mut(self) + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + if self.dim.size() == 1 { + self.first_mut() + } else { + None + } + } +} + +impl ArrayLikeMut for &mut ArrayBase +where + S: DataMut, + D: Dimension, +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + ArrayRef::view_mut(self) } fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> @@ -222,6 +483,58 @@ impl ArrayLike for [A] } } +impl ArrayLike for &[A] +{ + type Dim = Ix1; + + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + aview1(self) + } + + fn dim(&self) -> Self::Dim + { + Ix1(self.len()) + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.len() == 1 { + Some(&self[0]) + } else { + None + } + } +} + +impl ArrayLike for &mut [A] +{ + type Dim = Ix1; + + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + aview1(self) + } + + fn dim(&self) -> Self::Dim + { + Ix1(self.len()) + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if self.len() == 1 { + Some(&self[0]) + } else { + None + } + } +} + impl ArrayLikeMut for [A] { fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> @@ -239,6 +552,23 @@ impl ArrayLikeMut for [A] } } +impl ArrayLikeMut for &mut [A] +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + aview_mut1(self) + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + if self.len() == 1 { + Some(&mut self[0]) + } else { + None + } + } +} + impl ArrayLike for Vec { type Dim = Ix1; @@ -261,6 +591,50 @@ impl ArrayLike for Vec } } +impl ArrayLike for &Vec +{ + type Dim = Ix1; + + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + (&**self).view() + } + + fn dim(&self) -> Self::Dim + { + (&**self).dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + (&**self).as_elem() + } +} + +impl ArrayLike for &mut Vec +{ + type Dim = Ix1; + + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + (&**self).view() + } + + fn dim(&self) -> Self::Dim + { + (&**self).dim() + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + (&**self).as_elem() + } +} + impl ArrayLikeMut for Vec { fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> @@ -274,26 +648,160 @@ impl ArrayLikeMut for Vec } } +impl ArrayLikeMut for &mut Vec +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + (&mut **self).view_mut() + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + (&mut **self).as_elem_mut() + } +} + +impl ArrayLike for [A; N] +{ + type Dim = Ix1; + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + ArrayView::from(self) + } + + fn dim(&self) -> Self::Dim + { + Ix1(N) + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if N == 1 { + Some(&self[0]) + } else { + None + } + } +} + +impl ArrayLike for &[A; N] +{ + type Dim = Ix1; + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + ArrayView::from(self) + } + + fn dim(&self) -> Self::Dim + { + Ix1(N) + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if N == 1 { + Some(&self[0]) + } else { + None + } + } +} + +impl ArrayLike for &mut [A; N] +{ + type Dim = Ix1; + type Elem = A; + + fn view(&self) -> ArrayView<'_, Self::Elem, Self::Dim> + { + ArrayView::from(self) + } + + fn dim(&self) -> Self::Dim + { + Ix1(N) + } + + fn as_elem(&self) -> Option<&Self::Elem> + { + if N == 1 { + Some(&self[0]) + } else { + None + } + } +} + +impl ArrayLikeMut for [A; N] +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + ArrayViewMut::from(self) + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + if N == 1 { + Some(&mut self[0]) + } else { + None + } + } +} + +impl ArrayLikeMut for &mut [A; N] +{ + fn view_mut(&mut self) -> ArrayViewMut<'_, Self::Elem, Self::Dim> + { + ArrayViewMut::from(self) + } + + fn as_elem_mut(&mut self) -> Option<&mut Self::Elem> + { + if N == 1 { + Some(&mut self[0]) + } else { + None + } + } +} + #[cfg(test)] mod tests { - - use core::ops::Mul; + use core::ops::{Mul, MulAssign}; use crate::{array, Array, ArrayLike, DimMax}; - fn multiply(left: &T, right: &G) -> Array>::Output> + use super::ArrayLikeMut; + + fn multiply(left: T, right: G) -> Array>::Output> where T: ArrayLike, - G: ArrayLike, + G: ArrayLike, // Bounds to enable multiplication - T::Elem: Clone + Mul, + T::Elem: Clone + Mul, G::Elem: Clone, T::Dim: DimMax, { - let left = &*left.view(); - let right = &*right.view(); - left * right + &left.view() * &right.view() + } + + fn multiply_assign(mut left: T, right: G) + where + T: ArrayLikeMut, + G: ArrayLike, + // Bounds to enable multiplication + T::Elem: Clone + MulAssign, + G::Elem: Clone, + // Ensure that the broadcast is still assignable to the left side + T::Dim: DimMax, + { + *left.view_mut() *= &right.view(); } #[test] @@ -302,6 +810,19 @@ mod tests let left = array![1, 2]; let right = array![3, 4]; assert_eq!(multiply(&left, &right), array![3, 8]); - assert_eq!(multiply(&left, &3), array![3, 6]); + assert_eq!(multiply(&left, 3), array![3, 6]); + } + + #[test] + fn test_multiply_assign() + { + let mut left = array![1, 2]; + let right = array![3, 4]; + + multiply_assign(&mut left, &right); + assert_eq!(left, array![3, 8]); + + multiply_assign(&mut left, 2); + assert_eq!(left, array![6, 16]); } } diff --git a/src/lib.rs b/src/lib.rs index 6f24fe0f5..a9f53c147 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -208,7 +208,7 @@ mod zip; mod dimension; mod arraylike; -pub use crate::arraylike::ArrayLike; +pub use crate::arraylike::{ArrayLike, ArrayLikeMut}; pub use crate::zip::{FoldWhile, IntoNdProducer, NdProducer, Zip};