ferray_core/dimension/
mod.rs

1// ferray-core: Dimension trait and concrete dimension types
2//
3// These types mirror ndarray's Ix1..Ix6 and IxDyn but live in ferray-core's
4// namespace so that ndarray never appears in the public API.
5
6#[cfg(feature = "std")]
7pub mod broadcast;
8// static_shape depends on Array, Vec, format!, and the ndarray-gated methods
9// on Dimension, so it's only available when std is in scope.
10#[cfg(all(feature = "const_shapes", feature = "std"))]
11pub mod static_shape;
12
13use core::fmt;
14
15#[cfg(not(feature = "std"))]
16extern crate alloc;
17#[cfg(not(feature = "std"))]
18use alloc::vec::Vec;
19
20// We need ndarray's Dimension trait in scope for `as_array_view()` etc.
21#[cfg(feature = "std")]
22use ndarray::Dimension as NdDimension;
23
24/// Trait for types that describe the dimensionality of an array.
25///
26/// Each dimension type knows its number of axes at the type level
27/// (except [`IxDyn`] which carries it at runtime).
28pub trait Dimension: Clone + PartialEq + Eq + fmt::Debug + Send + Sync + 'static {
29    /// The number of axes, or `None` for dynamic-rank arrays.
30    const NDIM: Option<usize>;
31
32    /// The corresponding `ndarray` dimension type (private, not exposed in public API).
33    #[doc(hidden)]
34    #[cfg(feature = "std")]
35    type NdarrayDim: ndarray::Dimension;
36
37    /// Return the shape as a slice.
38    fn as_slice(&self) -> &[usize];
39
40    /// Return the shape as a mutable slice.
41    fn as_slice_mut(&mut self) -> &mut [usize];
42
43    /// Number of dimensions.
44    fn ndim(&self) -> usize {
45        self.as_slice().len()
46    }
47
48    /// Total number of elements (product of all dimension sizes).
49    fn size(&self) -> usize {
50        self.as_slice().iter().product()
51    }
52
53    /// Convert to the internal ndarray dimension type.
54    #[doc(hidden)]
55    #[cfg(feature = "std")]
56    fn to_ndarray_dim(&self) -> Self::NdarrayDim;
57
58    /// Create from the internal ndarray dimension type.
59    #[doc(hidden)]
60    #[cfg(feature = "std")]
61    fn from_ndarray_dim(dim: &Self::NdarrayDim) -> Self;
62
63    /// Construct a dimension from a slice of axis lengths.
64    ///
65    /// Returns `None` if the slice length does not match `Self::NDIM`
66    /// for fixed-rank dimensions. Always succeeds for [`IxDyn`].
67    ///
68    /// This is the inverse of [`Dimension::as_slice`].
69    fn from_dim_slice(shape: &[usize]) -> Option<Self>;
70}
71
72// ---------------------------------------------------------------------------
73// Fixed-rank dimension types
74// ---------------------------------------------------------------------------
75
76macro_rules! impl_fixed_dimension {
77    ($name:ident, $n:expr, $ndarray_ty:ty) => {
78        /// A fixed-rank dimension with
79        #[doc = concat!(stringify!($n), " axes.")]
80        #[derive(Clone, PartialEq, Eq, Hash)]
81        pub struct $name {
82            shape: [usize; $n],
83        }
84
85        impl $name {
86            /// Create a new dimension from a fixed-size array.
87            #[inline]
88            pub fn new(shape: [usize; $n]) -> Self {
89                Self { shape }
90            }
91        }
92
93        impl fmt::Debug for $name {
94            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95                write!(f, "{:?}", &self.shape[..])
96            }
97        }
98
99        impl From<[usize; $n]> for $name {
100            #[inline]
101            fn from(shape: [usize; $n]) -> Self {
102                Self::new(shape)
103            }
104        }
105
106        impl Dimension for $name {
107            const NDIM: Option<usize> = Some($n);
108
109            #[cfg(feature = "std")]
110            type NdarrayDim = $ndarray_ty;
111
112            #[inline]
113            fn as_slice(&self) -> &[usize] {
114                &self.shape
115            }
116
117            #[inline]
118            fn as_slice_mut(&mut self) -> &mut [usize] {
119                &mut self.shape
120            }
121
122            #[cfg(feature = "std")]
123            fn to_ndarray_dim(&self) -> Self::NdarrayDim {
124                // ndarray::Dim implements From<[usize; N]> for N=1..6
125                ndarray::Dim(self.shape)
126            }
127
128            #[cfg(feature = "std")]
129            fn from_ndarray_dim(dim: &Self::NdarrayDim) -> Self {
130                let view = dim.as_array_view();
131                let s = view.as_slice().expect("ndarray dim should be contiguous");
132                let mut shape = [0usize; $n];
133                shape.copy_from_slice(s);
134                Self { shape }
135            }
136
137            fn from_dim_slice(shape: &[usize]) -> Option<Self> {
138                if shape.len() != $n {
139                    return None;
140                }
141                let mut arr = [0usize; $n];
142                arr.copy_from_slice(shape);
143                Some(Self { shape: arr })
144            }
145        }
146    };
147}
148
149impl_fixed_dimension!(Ix1, 1, ndarray::Ix1);
150impl_fixed_dimension!(Ix2, 2, ndarray::Ix2);
151impl_fixed_dimension!(Ix3, 3, ndarray::Ix3);
152impl_fixed_dimension!(Ix4, 4, ndarray::Ix4);
153impl_fixed_dimension!(Ix5, 5, ndarray::Ix5);
154impl_fixed_dimension!(Ix6, 6, ndarray::Ix6);
155
156// ---------------------------------------------------------------------------
157// Ix0: scalar (0-dimensional)
158// ---------------------------------------------------------------------------
159
160/// A zero-dimensional (scalar) dimension.
161#[derive(Clone, PartialEq, Eq, Hash)]
162pub struct Ix0;
163
164impl fmt::Debug for Ix0 {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        write!(f, "[]")
167    }
168}
169
170impl Dimension for Ix0 {
171    const NDIM: Option<usize> = Some(0);
172
173    #[cfg(feature = "std")]
174    type NdarrayDim = ndarray::Ix0;
175
176    #[inline]
177    fn as_slice(&self) -> &[usize] {
178        &[]
179    }
180
181    #[inline]
182    fn as_slice_mut(&mut self) -> &mut [usize] {
183        &mut []
184    }
185
186    #[cfg(feature = "std")]
187    fn to_ndarray_dim(&self) -> Self::NdarrayDim {
188        ndarray::Dim(())
189    }
190
191    #[cfg(feature = "std")]
192    fn from_ndarray_dim(_dim: &Self::NdarrayDim) -> Self {
193        Ix0
194    }
195
196    fn from_dim_slice(shape: &[usize]) -> Option<Self> {
197        if shape.is_empty() { Some(Ix0) } else { None }
198    }
199}
200
201// ---------------------------------------------------------------------------
202// IxDyn: dynamic-rank dimension
203// ---------------------------------------------------------------------------
204
205/// A dynamic-rank dimension whose number of axes is determined at runtime.
206#[derive(Clone, PartialEq, Eq, Hash)]
207pub struct IxDyn {
208    shape: Vec<usize>,
209}
210
211impl IxDyn {
212    /// Create a new dynamic dimension from a slice.
213    pub fn new(shape: &[usize]) -> Self {
214        Self {
215            shape: shape.to_vec(),
216        }
217    }
218}
219
220impl fmt::Debug for IxDyn {
221    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
222        write!(f, "{:?}", &self.shape[..])
223    }
224}
225
226impl From<Vec<usize>> for IxDyn {
227    fn from(shape: Vec<usize>) -> Self {
228        Self { shape }
229    }
230}
231
232impl From<&[usize]> for IxDyn {
233    fn from(shape: &[usize]) -> Self {
234        Self::new(shape)
235    }
236}
237
238impl Dimension for IxDyn {
239    const NDIM: Option<usize> = None;
240
241    #[cfg(feature = "std")]
242    type NdarrayDim = ndarray::IxDyn;
243
244    #[inline]
245    fn as_slice(&self) -> &[usize] {
246        &self.shape
247    }
248
249    #[inline]
250    fn as_slice_mut(&mut self) -> &mut [usize] {
251        &mut self.shape
252    }
253
254    #[cfg(feature = "std")]
255    fn to_ndarray_dim(&self) -> Self::NdarrayDim {
256        ndarray::IxDyn(&self.shape)
257    }
258
259    #[cfg(feature = "std")]
260    fn from_ndarray_dim(dim: &Self::NdarrayDim) -> Self {
261        let view = dim.as_array_view();
262        let s = view.as_slice().expect("ndarray IxDyn should be contiguous");
263        Self { shape: s.to_vec() }
264    }
265
266    fn from_dim_slice(shape: &[usize]) -> Option<Self> {
267        Some(Self {
268            shape: shape.to_vec(),
269        })
270    }
271}
272
273/// Newtype for axis indices used throughout ferray.
274#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
275pub struct Axis(pub usize);
276
277impl Axis {
278    /// Return the axis index.
279    #[inline]
280    pub fn index(self) -> usize {
281        self.0
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    #[test]
290    fn ix1_basics() {
291        let d = Ix1::new([5]);
292        assert_eq!(d.ndim(), 1);
293        assert_eq!(d.size(), 5);
294        assert_eq!(d.as_slice(), &[5]);
295    }
296
297    #[test]
298    fn ix2_basics() {
299        let d = Ix2::new([3, 4]);
300        assert_eq!(d.ndim(), 2);
301        assert_eq!(d.size(), 12);
302    }
303
304    #[test]
305    fn ix0_basics() {
306        let d = Ix0;
307        assert_eq!(d.ndim(), 0);
308        assert_eq!(d.size(), 1);
309    }
310
311    #[test]
312    fn ixdyn_basics() {
313        let d = IxDyn::new(&[2, 3, 4]);
314        assert_eq!(d.ndim(), 3);
315        assert_eq!(d.size(), 24);
316    }
317
318    #[test]
319    fn roundtrip_ix2_ndarray() {
320        let d = Ix2::new([3, 7]);
321        let nd = d.to_ndarray_dim();
322        let d2 = Ix2::from_ndarray_dim(&nd);
323        assert_eq!(d, d2);
324    }
325
326    #[test]
327    fn roundtrip_ixdyn_ndarray() {
328        let d = IxDyn::new(&[2, 5, 3]);
329        let nd = d.to_ndarray_dim();
330        let d2 = IxDyn::from_ndarray_dim(&nd);
331        assert_eq!(d, d2);
332    }
333
334    #[test]
335    fn axis_index() {
336        let a = Axis(2);
337        assert_eq!(a.index(), 2);
338    }
339}
ferray_core/dimension/mod.rs

ferray_core/dimension/
mod.rs
