ndarray/
shape_builder.rs

1use crate::dimension::IntoDimension;
2use crate::order::Order;
3use crate::Dimension;
4
5/// A contiguous array shape of n dimensions.
6///
7/// Either c- or f- memory ordered (*c* a.k.a *row major* is the default).
8#[derive(Copy, Clone, Debug)]
9pub struct Shape<D>
10{
11    /// Shape (axis lengths)
12    pub(crate) dim: D,
13    /// Strides can only be C or F here
14    pub(crate) strides: Strides<Contiguous>,
15}
16
17#[derive(Copy, Clone, Debug)]
18pub(crate) enum Contiguous {}
19
20impl<D> Shape<D>
21{
22    pub(crate) fn is_c(&self) -> bool
23    {
24        matches!(self.strides, Strides::C)
25    }
26}
27
28/// An array shape of n dimensions in c-order, f-order or custom strides.
29#[derive(Copy, Clone, Debug)]
30pub struct StrideShape<D>
31{
32    pub(crate) dim: D,
33    pub(crate) strides: Strides<D>,
34}
35
36impl<D> StrideShape<D>
37where D: Dimension
38{
39    /// Return a reference to the dimension
40    pub fn raw_dim(&self) -> &D
41    {
42        &self.dim
43    }
44    /// Return the size of the shape in number of elements
45    pub fn size(&self) -> usize
46    {
47        self.dim.size()
48    }
49}
50
51/// Stride description
52#[derive(Copy, Clone, Debug)]
53pub(crate) enum Strides<D>
54{
55    /// Row-major ("C"-order)
56    C,
57    /// Column-major ("F"-order)
58    F,
59    /// Custom strides
60    Custom(D),
61}
62
63impl<D> Strides<D>
64{
65    /// Return strides for `dim` (computed from dimension if c/f, else return the custom stride)
66    pub(crate) fn strides_for_dim(self, dim: &D) -> D
67    where D: Dimension
68    {
69        match self {
70            Strides::C => dim.default_strides(),
71            Strides::F => dim.fortran_strides(),
72            Strides::Custom(c) => {
73                debug_assert_eq!(
74                    c.ndim(),
75                    dim.ndim(),
76                    "Custom strides given with {} dimensions, expected {}",
77                    c.ndim(),
78                    dim.ndim()
79                );
80                c
81            }
82        }
83    }
84
85    #[inline]
86    pub(crate) fn is_custom(&self) -> bool
87    {
88        matches!(*self, Strides::Custom(_))
89    }
90}
91
92/// A trait for `Shape` and `D where D: Dimension` that allows
93/// customizing the memory layout (strides) of an array shape.
94///
95/// This trait is used together with array constructor methods like
96/// `Array::from_shape_vec`.
97pub trait ShapeBuilder
98{
99    type Dim: Dimension;
100    type Strides;
101
102    fn into_shape_with_order(self) -> Shape<Self::Dim>;
103    fn f(self) -> Shape<Self::Dim>;
104    fn set_f(self, is_f: bool) -> Shape<Self::Dim>;
105    fn strides(self, strides: Self::Strides) -> StrideShape<Self::Dim>;
106}
107
108impl<D> From<D> for Shape<D>
109where D: Dimension
110{
111    /// Create a `Shape` from `dimension`, using the default memory layout.
112    fn from(dimension: D) -> Shape<D>
113    {
114        dimension.into_shape_with_order()
115    }
116}
117
118impl<T, D> From<T> for StrideShape<D>
119where
120    D: Dimension,
121    T: ShapeBuilder<Dim = D>,
122{
123    fn from(value: T) -> Self
124    {
125        let shape = value.into_shape_with_order();
126        let st = if shape.is_c() { Strides::C } else { Strides::F };
127        StrideShape {
128            strides: st,
129            dim: shape.dim,
130        }
131    }
132}
133
134impl<T> ShapeBuilder for T
135where T: IntoDimension
136{
137    type Dim = T::Dim;
138    type Strides = T;
139    fn into_shape_with_order(self) -> Shape<Self::Dim>
140    {
141        Shape {
142            dim: self.into_dimension(),
143            strides: Strides::C,
144        }
145    }
146    fn f(self) -> Shape<Self::Dim>
147    {
148        self.set_f(true)
149    }
150    fn set_f(self, is_f: bool) -> Shape<Self::Dim>
151    {
152        self.into_shape_with_order().set_f(is_f)
153    }
154    fn strides(self, st: T) -> StrideShape<Self::Dim>
155    {
156        self.into_shape_with_order().strides(st.into_dimension())
157    }
158}
159
160impl<D> ShapeBuilder for Shape<D>
161where D: Dimension
162{
163    type Dim = D;
164    type Strides = D;
165
166    fn into_shape_with_order(self) -> Shape<D>
167    {
168        self
169    }
170
171    fn f(self) -> Self
172    {
173        self.set_f(true)
174    }
175
176    fn set_f(mut self, is_f: bool) -> Self
177    {
178        self.strides = if !is_f { Strides::C } else { Strides::F };
179        self
180    }
181
182    fn strides(self, st: D) -> StrideShape<D>
183    {
184        StrideShape {
185            dim: self.dim,
186            strides: Strides::Custom(st),
187        }
188    }
189}
190
191impl<D> Shape<D>
192where D: Dimension
193{
194    /// Return a reference to the dimension
195    pub fn raw_dim(&self) -> &D
196    {
197        &self.dim
198    }
199    /// Return the size of the shape in number of elements
200    pub fn size(&self) -> usize
201    {
202        self.dim.size()
203    }
204}
205
206/// Array shape argument with optional order parameter
207///
208/// Shape or array dimension argument, with optional [`Order`] parameter.
209///
210/// This is an argument conversion trait that is used to accept an array shape and
211/// (optionally) an ordering argument.
212///
213/// See for example [`.to_shape()`](crate::ArrayBase::to_shape).
214pub trait ShapeArg
215{
216    type Dim: Dimension;
217    fn into_shape_and_order(self) -> (Self::Dim, Option<Order>);
218}
219
220impl<T> ShapeArg for T
221where T: IntoDimension
222{
223    type Dim = T::Dim;
224
225    fn into_shape_and_order(self) -> (Self::Dim, Option<Order>)
226    {
227        (self.into_dimension(), None)
228    }
229}
230
231impl<T> ShapeArg for (T, Order)
232where T: IntoDimension
233{
234    type Dim = T::Dim;
235
236    fn into_shape_and_order(self) -> (Self::Dim, Option<Order>)
237    {
238        (self.0.into_dimension(), Some(self.1))
239    }
240}
ndarray/shape_builder.rs

ndarray/
shape_builder.rs
