ndarray/zip/
ndproducer.rs

1use crate::imp_prelude::*;
2use crate::Layout;
3use crate::NdIndex;
4#[cfg(not(feature = "std"))]
5use alloc::vec::Vec;
6
7/// Argument conversion into a producer.
8///
9/// Slices and vectors can be used (equivalent to 1-dimensional array views).
10///
11/// This trait is like `IntoIterator` for `NdProducers` instead of iterators.
12pub trait IntoNdProducer
13{
14    /// The element produced per iteration.
15    type Item;
16    /// Dimension type of the producer
17    type Dim: Dimension;
18    type Output: NdProducer<Dim = Self::Dim, Item = Self::Item>;
19    /// Convert the value into an `NdProducer`.
20    fn into_producer(self) -> Self::Output;
21}
22
23impl<P> IntoNdProducer for P
24where P: NdProducer
25{
26    type Item = P::Item;
27    type Dim = P::Dim;
28    type Output = Self;
29    fn into_producer(self) -> Self::Output
30    {
31        self
32    }
33}
34
35/// A producer of an n-dimensional set of elements;
36/// for example an array view, mutable array view or an iterator
37/// that yields chunks.
38///
39/// Producers are used as a arguments to [`Zip`](crate::Zip) and
40/// [`azip!()`].
41///
42/// # Comparison to `IntoIterator`
43///
44/// Most `NdProducers` are *iterable* (implement `IntoIterator`) but not directly
45/// iterators. This separation is needed because the producer represents
46/// a multidimensional set of items, it can be split along a particular axis for
47/// parallelization, and it has no fixed correspondence to a sequence.
48///
49/// The natural exception is one dimensional producers, like `AxisIter`, which
50/// implement `Iterator` directly
51/// (`AxisIter` traverses a one dimensional sequence, along an axis, while
52/// *producing* multidimensional items).
53///
54/// See also [`IntoNdProducer`]
55pub trait NdProducer
56{
57    /// The element produced per iteration.
58    type Item;
59    // Internal use / Pointee type
60    /// Dimension type
61    type Dim: Dimension;
62
63    // The pointer Ptr is used by an array view to simply point to the
64    // current element. It doesn't have to be a pointer (see Indices).
65    // Its main function is that it can be incremented with a particular
66    // stride (= along a particular axis)
67    #[doc(hidden)]
68    /// Pointer or stand-in for pointer
69    type Ptr: Offset<Stride = Self::Stride>;
70    #[doc(hidden)]
71    /// Pointer stride
72    type Stride: Copy;
73
74    #[doc(hidden)]
75    fn layout(&self) -> Layout;
76    /// Return the shape of the producer.
77    fn raw_dim(&self) -> Self::Dim;
78    #[doc(hidden)]
79    fn equal_dim(&self, dim: &Self::Dim) -> bool
80    {
81        self.raw_dim() == *dim
82    }
83    #[doc(hidden)]
84    fn as_ptr(&self) -> Self::Ptr;
85    #[doc(hidden)]
86    unsafe fn as_ref(&self, ptr: Self::Ptr) -> Self::Item;
87    #[doc(hidden)]
88    unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr;
89    #[doc(hidden)]
90    fn stride_of(&self, axis: Axis) -> <Self::Ptr as Offset>::Stride;
91    #[doc(hidden)]
92    fn contiguous_stride(&self) -> Self::Stride;
93    #[doc(hidden)]
94    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
95    where Self: Sized;
96
97    private_decl! {}
98}
99
100pub trait Offset: Copy
101{
102    type Stride: Copy;
103    unsafe fn stride_offset(self, s: Self::Stride, index: usize) -> Self;
104    private_decl! {}
105}
106
107impl<T> Offset for *const T
108{
109    type Stride = isize;
110    unsafe fn stride_offset(self, s: Self::Stride, index: usize) -> Self
111    {
112        self.offset(s * (index as isize))
113    }
114    private_impl! {}
115}
116
117impl<T> Offset for *mut T
118{
119    type Stride = isize;
120    unsafe fn stride_offset(self, s: Self::Stride, index: usize) -> Self
121    {
122        self.offset(s * (index as isize))
123    }
124    private_impl! {}
125}
126
127/// An array reference is an n-dimensional producer of element references
128/// (like ArrayView).
129impl<'a, A: 'a, S, D> IntoNdProducer for &'a ArrayBase<S, D>
130where
131    D: Dimension,
132    S: Data<Elem = A>,
133{
134    type Item = &'a A;
135    type Dim = D;
136    type Output = ArrayView<'a, A, D>;
137    fn into_producer(self) -> Self::Output
138    {
139        self.view()
140    }
141}
142
143/// A mutable array reference is an n-dimensional producer of mutable element
144/// references (like ArrayViewMut).
145impl<'a, A: 'a, S, D> IntoNdProducer for &'a mut ArrayBase<S, D>
146where
147    D: Dimension,
148    S: DataMut<Elem = A>,
149{
150    type Item = &'a mut A;
151    type Dim = D;
152    type Output = ArrayViewMut<'a, A, D>;
153    fn into_producer(self) -> Self::Output
154    {
155        self.view_mut()
156    }
157}
158
159/// A slice is a one-dimensional producer
160impl<'a, A: 'a> IntoNdProducer for &'a [A]
161{
162    type Item = <Self::Output as NdProducer>::Item;
163    type Dim = Ix1;
164    type Output = ArrayView1<'a, A>;
165    fn into_producer(self) -> Self::Output
166    {
167        <_>::from(self)
168    }
169}
170
171/// A mutable slice is a mutable one-dimensional producer
172impl<'a, A: 'a> IntoNdProducer for &'a mut [A]
173{
174    type Item = <Self::Output as NdProducer>::Item;
175    type Dim = Ix1;
176    type Output = ArrayViewMut1<'a, A>;
177    fn into_producer(self) -> Self::Output
178    {
179        <_>::from(self)
180    }
181}
182
183/// A one-dimensional array is a one-dimensional producer
184impl<'a, A: 'a, const N: usize> IntoNdProducer for &'a [A; N]
185{
186    type Item = <Self::Output as NdProducer>::Item;
187    type Dim = Ix1;
188    type Output = ArrayView1<'a, A>;
189    fn into_producer(self) -> Self::Output
190    {
191        <_>::from(self)
192    }
193}
194
195/// A mutable one-dimensional array is a mutable one-dimensional producer
196impl<'a, A: 'a, const N: usize> IntoNdProducer for &'a mut [A; N]
197{
198    type Item = <Self::Output as NdProducer>::Item;
199    type Dim = Ix1;
200    type Output = ArrayViewMut1<'a, A>;
201    fn into_producer(self) -> Self::Output
202    {
203        <_>::from(self)
204    }
205}
206
207/// A Vec is a one-dimensional producer
208impl<'a, A: 'a> IntoNdProducer for &'a Vec<A>
209{
210    type Item = <Self::Output as NdProducer>::Item;
211    type Dim = Ix1;
212    type Output = ArrayView1<'a, A>;
213    fn into_producer(self) -> Self::Output
214    {
215        <_>::from(self)
216    }
217}
218
219/// A mutable Vec is a mutable one-dimensional producer
220impl<'a, A: 'a> IntoNdProducer for &'a mut Vec<A>
221{
222    type Item = <Self::Output as NdProducer>::Item;
223    type Dim = Ix1;
224    type Output = ArrayViewMut1<'a, A>;
225    fn into_producer(self) -> Self::Output
226    {
227        <_>::from(self)
228    }
229}
230
231impl<'a, A, D: Dimension> NdProducer for ArrayView<'a, A, D>
232{
233    type Item = &'a A;
234    type Dim = D;
235    type Ptr = *mut A;
236    type Stride = isize;
237
238    private_impl! {}
239
240    fn raw_dim(&self) -> Self::Dim
241    {
242        self.raw_dim()
243    }
244
245    fn equal_dim(&self, dim: &Self::Dim) -> bool
246    {
247        self.dim.equal(dim)
248    }
249
250    fn as_ptr(&self) -> *mut A
251    {
252        self.as_ptr() as _
253    }
254
255    fn layout(&self) -> Layout
256    {
257        self.layout_impl()
258    }
259
260    unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item
261    {
262        &*ptr
263    }
264
265    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A
266    {
267        self.ptr.as_ptr().offset(i.index_unchecked(&self.strides))
268    }
269
270    fn stride_of(&self, axis: Axis) -> isize
271    {
272        self.stride_of(axis)
273    }
274
275    #[inline(always)]
276    fn contiguous_stride(&self) -> Self::Stride
277    {
278        1
279    }
280
281    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
282    {
283        self.split_at(axis, index)
284    }
285}
286
287impl<'a, A, D: Dimension> NdProducer for ArrayViewMut<'a, A, D>
288{
289    type Item = &'a mut A;
290    type Dim = D;
291    type Ptr = *mut A;
292    type Stride = isize;
293
294    private_impl! {}
295
296    fn raw_dim(&self) -> Self::Dim
297    {
298        self.raw_dim()
299    }
300
301    fn equal_dim(&self, dim: &Self::Dim) -> bool
302    {
303        self.dim.equal(dim)
304    }
305
306    fn as_ptr(&self) -> *mut A
307    {
308        self.as_ptr() as _
309    }
310
311    fn layout(&self) -> Layout
312    {
313        self.layout_impl()
314    }
315
316    unsafe fn as_ref(&self, ptr: *mut A) -> Self::Item
317    {
318        &mut *ptr
319    }
320
321    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A
322    {
323        self.ptr.as_ptr().offset(i.index_unchecked(&self.strides))
324    }
325
326    fn stride_of(&self, axis: Axis) -> isize
327    {
328        self.stride_of(axis)
329    }
330
331    #[inline(always)]
332    fn contiguous_stride(&self) -> Self::Stride
333    {
334        1
335    }
336
337    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
338    {
339        self.split_at(axis, index)
340    }
341}
342
343impl<A, D: Dimension> NdProducer for RawArrayView<A, D>
344{
345    type Item = *const A;
346    type Dim = D;
347    type Ptr = *const A;
348    type Stride = isize;
349
350    private_impl! {}
351
352    fn raw_dim(&self) -> Self::Dim
353    {
354        self.raw_dim()
355    }
356
357    fn equal_dim(&self, dim: &Self::Dim) -> bool
358    {
359        self.dim.equal(dim)
360    }
361
362    fn as_ptr(&self) -> *const A
363    {
364        self.as_ptr()
365    }
366
367    fn layout(&self) -> Layout
368    {
369        self.layout_impl()
370    }
371
372    unsafe fn as_ref(&self, ptr: *const A) -> *const A
373    {
374        ptr
375    }
376
377    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *const A
378    {
379        self.ptr.as_ptr().offset(i.index_unchecked(&self.strides))
380    }
381
382    fn stride_of(&self, axis: Axis) -> isize
383    {
384        self.stride_of(axis)
385    }
386
387    #[inline(always)]
388    fn contiguous_stride(&self) -> Self::Stride
389    {
390        1
391    }
392
393    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
394    {
395        self.split_at(axis, index)
396    }
397}
398
399impl<A, D: Dimension> NdProducer for RawArrayViewMut<A, D>
400{
401    type Item = *mut A;
402    type Dim = D;
403    type Ptr = *mut A;
404    type Stride = isize;
405
406    private_impl! {}
407
408    fn raw_dim(&self) -> Self::Dim
409    {
410        self.raw_dim()
411    }
412
413    fn equal_dim(&self, dim: &Self::Dim) -> bool
414    {
415        self.dim.equal(dim)
416    }
417
418    fn as_ptr(&self) -> *mut A
419    {
420        self.as_ptr() as _
421    }
422
423    fn layout(&self) -> Layout
424    {
425        self.layout_impl()
426    }
427
428    unsafe fn as_ref(&self, ptr: *mut A) -> *mut A
429    {
430        ptr
431    }
432
433    unsafe fn uget_ptr(&self, i: &Self::Dim) -> *mut A
434    {
435        self.ptr.as_ptr().offset(i.index_unchecked(&self.strides))
436    }
437
438    fn stride_of(&self, axis: Axis) -> isize
439    {
440        self.stride_of(axis)
441    }
442
443    #[inline(always)]
444    fn contiguous_stride(&self) -> Self::Stride
445    {
446        1
447    }
448
449    fn split_at(self, axis: Axis, index: usize) -> (Self, Self)
450    {
451        self.split_at(axis, index)
452    }
453}
ndarray/zip/ndproducer.rs

ndarray/zip/
ndproducer.rs
