ndarray/iterators/
mod.rs

1// Copyright 2014-2016 bluss and ndarray developers.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9#[macro_use]
10mod macros;
11mod chunks;
12pub mod iter;
13mod lanes;
14mod windows;
15
16use std::iter::FromIterator;
17use std::marker::PhantomData;
18use std::ptr;
19use alloc::vec::Vec;
20
21use crate::Ix1;
22
23use super::{ArrayBase, ArrayView, ArrayViewMut, Axis, Data, NdProducer, RemoveAxis};
24use super::{Dimension, Ix, Ixs};
25
26pub use self::chunks::{ExactChunks, ExactChunksIter, ExactChunksIterMut, ExactChunksMut};
27pub use self::lanes::{Lanes, LanesMut};
28pub use self::windows::Windows;
29
30use std::slice::{self, Iter as SliceIter, IterMut as SliceIterMut};
31
32/// Base for iterators over all axes.
33///
34/// Iterator element type is `*mut A`.
35pub struct Baseiter<A, D> {
36    ptr: *mut A,
37    dim: D,
38    strides: D,
39    index: Option<D>,
40}
41
42impl<A, D: Dimension> Baseiter<A, D> {
43    /// Creating a Baseiter is unsafe because shape and stride parameters need
44    /// to be correct to avoid performing an unsafe pointer offset while
45    /// iterating.
46    #[inline]
47    pub unsafe fn new(ptr: *mut A, len: D, stride: D) -> Baseiter<A, D> {
48        Baseiter {
49            ptr,
50            index: len.first_index(),
51            dim: len,
52            strides: stride,
53        }
54    }
55}
56
57impl<A, D: Dimension> Iterator for Baseiter<A, D> {
58    type Item = *mut A;
59
60    #[inline]
61    fn next(&mut self) -> Option<*mut A> {
62        let index = match self.index {
63            None => return None,
64            Some(ref ix) => ix.clone(),
65        };
66        let offset = D::stride_offset(&index, &self.strides);
67        self.index = self.dim.next_for(index);
68        unsafe { Some(self.ptr.offset(offset)) }
69    }
70
71    fn size_hint(&self) -> (usize, Option<usize>) {
72        let len = self.len();
73        (len, Some(len))
74    }
75
76    fn fold<Acc, G>(mut self, init: Acc, mut g: G) -> Acc
77    where
78        G: FnMut(Acc, *mut A) -> Acc,
79    {
80        let ndim = self.dim.ndim();
81        debug_assert_ne!(ndim, 0);
82        let mut accum = init;
83        while let Some(mut index) = self.index {
84            let stride = self.strides.last_elem() as isize;
85            let elem_index = index.last_elem();
86            let len = self.dim.last_elem();
87            let offset = D::stride_offset(&index, &self.strides);
88            unsafe {
89                let row_ptr = self.ptr.offset(offset);
90                let mut i = 0;
91                let i_end = len - elem_index;
92                while i < i_end {
93                    accum = g(accum, row_ptr.offset(i as isize * stride));
94                    i += 1;
95                }
96            }
97            index.set_last_elem(len - 1);
98            self.index = self.dim.next_for(index);
99        }
100        accum
101    }
102}
103
104impl<'a, A, D: Dimension> ExactSizeIterator for Baseiter<A, D> {
105    fn len(&self) -> usize {
106        match self.index {
107            None => 0,
108            Some(ref ix) => {
109                let gone = self
110                    .dim
111                    .default_strides()
112                    .slice()
113                    .iter()
114                    .zip(ix.slice().iter())
115                    .fold(0, |s, (&a, &b)| s + a as usize * b as usize);
116                self.dim.size() - gone
117            }
118        }
119    }
120}
121
122impl<A> DoubleEndedIterator for Baseiter<A, Ix1> {
123    #[inline]
124    fn next_back(&mut self) -> Option<*mut A> {
125        let index = match self.index {
126            None => return None,
127            Some(ix) => ix,
128        };
129        self.dim[0] -= 1;
130        let offset = <_>::stride_offset(&self.dim, &self.strides);
131        if index == self.dim {
132            self.index = None;
133        }
134
135        unsafe { Some(self.ptr.offset(offset)) }
136    }
137
138    fn nth_back(&mut self, n: usize) -> Option<*mut A> {
139        let index = self.index?;
140        let len = self.dim[0] - index[0];
141        if n < len {
142            self.dim[0] -= n + 1;
143            let offset = <_>::stride_offset(&self.dim, &self.strides);
144            if index == self.dim {
145                self.index = None;
146            }
147            unsafe { Some(self.ptr.offset(offset)) }
148        } else {
149            self.index = None;
150            None
151        }
152    }
153
154    fn rfold<Acc, G>(mut self, init: Acc, mut g: G) -> Acc
155    where
156        G: FnMut(Acc, *mut A) -> Acc,
157    {
158        let mut accum = init;
159        if let Some(index) = self.index {
160            let elem_index = index[0];
161            unsafe {
162                // self.dim[0] is the current length
163                while self.dim[0] > elem_index {
164                    self.dim[0] -= 1;
165                    accum = g(
166                        accum,
167                        self.ptr
168                            .offset(Ix1::stride_offset(&self.dim, &self.strides)),
169                    );
170                }
171            }
172        }
173        accum
174    }
175}
176
177clone_bounds!(
178    [A, D: Clone]
179    Baseiter[A, D] {
180        @copy {
181            ptr,
182        }
183        dim,
184        strides,
185        index,
186    }
187);
188
189clone_bounds!(
190    ['a, A, D: Clone]
191    ElementsBase['a, A, D] {
192        @copy {
193            life,
194        }
195        inner,
196    }
197);
198
199impl<'a, A, D: Dimension> ElementsBase<'a, A, D> {
200    pub fn new(v: ArrayView<'a, A, D>) -> Self {
201        ElementsBase {
202            inner: v.into_base_iter(),
203            life: PhantomData,
204        }
205    }
206}
207
208impl<'a, A, D: Dimension> Iterator for ElementsBase<'a, A, D> {
209    type Item = &'a A;
210    #[inline]
211    fn next(&mut self) -> Option<&'a A> {
212        self.inner.next().map(|p| unsafe { &*p })
213    }
214
215    fn size_hint(&self) -> (usize, Option<usize>) {
216        self.inner.size_hint()
217    }
218
219    fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
220    where
221        G: FnMut(Acc, Self::Item) -> Acc,
222    {
223        unsafe { self.inner.fold(init, move |acc, ptr| g(acc, &*ptr)) }
224    }
225}
226
227impl<'a, A> DoubleEndedIterator for ElementsBase<'a, A, Ix1> {
228    #[inline]
229    fn next_back(&mut self) -> Option<&'a A> {
230        self.inner.next_back().map(|p| unsafe { &*p })
231    }
232
233    fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
234    where
235        G: FnMut(Acc, Self::Item) -> Acc,
236    {
237        unsafe { self.inner.rfold(init, move |acc, ptr| g(acc, &*ptr)) }
238    }
239}
240
241impl<'a, A, D> ExactSizeIterator for ElementsBase<'a, A, D>
242where
243    D: Dimension,
244{
245    fn len(&self) -> usize {
246        self.inner.len()
247    }
248}
249
250macro_rules! either {
251    ($value:expr, $inner:pat => $result:expr) => {
252        match $value {
253            ElementsRepr::Slice($inner) => $result,
254            ElementsRepr::Counted($inner) => $result,
255        }
256    };
257}
258
259macro_rules! either_mut {
260    ($value:expr, $inner:ident => $result:expr) => {
261        match $value {
262            ElementsRepr::Slice(ref mut $inner) => $result,
263            ElementsRepr::Counted(ref mut $inner) => $result,
264        }
265    };
266}
267
268clone_bounds!(
269    ['a, A, D: Clone]
270    Iter['a, A, D] {
271        @copy {
272        }
273        inner,
274    }
275);
276
277impl<'a, A, D> Iter<'a, A, D>
278where
279    D: Dimension,
280{
281    pub(crate) fn new(self_: ArrayView<'a, A, D>) -> Self {
282        Iter {
283            inner: if let Some(slc) = self_.to_slice() {
284                ElementsRepr::Slice(slc.iter())
285            } else {
286                ElementsRepr::Counted(self_.into_elements_base())
287            },
288        }
289    }
290}
291
292impl<'a, A, D> IterMut<'a, A, D>
293where
294    D: Dimension,
295{
296    pub(crate) fn new(self_: ArrayViewMut<'a, A, D>) -> Self {
297        IterMut {
298            inner: match self_.try_into_slice() {
299                Ok(x) => ElementsRepr::Slice(x.iter_mut()),
300                Err(self_) => ElementsRepr::Counted(self_.into_elements_base()),
301            },
302        }
303    }
304}
305
306#[derive(Clone)]
307pub enum ElementsRepr<S, C> {
308    Slice(S),
309    Counted(C),
310}
311
312/// An iterator over the elements of an array.
313///
314/// Iterator element type is `&'a A`.
315///
316/// See [`.iter()`](../struct.ArrayBase.html#method.iter) for more information.
317pub struct Iter<'a, A, D> {
318    inner: ElementsRepr<SliceIter<'a, A>, ElementsBase<'a, A, D>>,
319}
320
321/// Counted read only iterator
322pub struct ElementsBase<'a, A, D> {
323    inner: Baseiter<A, D>,
324    life: PhantomData<&'a A>,
325}
326
327/// An iterator over the elements of an array (mutable).
328///
329/// Iterator element type is `&'a mut A`.
330///
331/// See [`.iter_mut()`](../struct.ArrayBase.html#method.iter_mut) for more information.
332pub struct IterMut<'a, A, D> {
333    inner: ElementsRepr<SliceIterMut<'a, A>, ElementsBaseMut<'a, A, D>>,
334}
335
336/// An iterator over the elements of an array.
337///
338/// Iterator element type is `&'a mut A`.
339pub struct ElementsBaseMut<'a, A, D> {
340    inner: Baseiter<A, D>,
341    life: PhantomData<&'a mut A>,
342}
343
344impl<'a, A, D: Dimension> ElementsBaseMut<'a, A, D> {
345    pub fn new(v: ArrayViewMut<'a, A, D>) -> Self {
346        ElementsBaseMut {
347            inner: v.into_base_iter(),
348            life: PhantomData,
349        }
350    }
351}
352
353/// An iterator over the indexes and elements of an array.
354///
355/// See [`.indexed_iter()`](../struct.ArrayBase.html#method.indexed_iter) for more information.
356#[derive(Clone)]
357pub struct IndexedIter<'a, A, D>(ElementsBase<'a, A, D>);
358/// An iterator over the indexes and elements of an array (mutable).
359///
360/// See [`.indexed_iter_mut()`](../struct.ArrayBase.html#method.indexed_iter_mut) for more information.
361pub struct IndexedIterMut<'a, A, D>(ElementsBaseMut<'a, A, D>);
362
363impl<'a, A, D> IndexedIter<'a, A, D>
364where
365    D: Dimension,
366{
367    pub(crate) fn new(x: ElementsBase<'a, A, D>) -> Self {
368        IndexedIter(x)
369    }
370}
371
372impl<'a, A, D> IndexedIterMut<'a, A, D>
373where
374    D: Dimension,
375{
376    pub(crate) fn new(x: ElementsBaseMut<'a, A, D>) -> Self {
377        IndexedIterMut(x)
378    }
379}
380
381impl<'a, A, D: Dimension> Iterator for Iter<'a, A, D> {
382    type Item = &'a A;
383    #[inline]
384    fn next(&mut self) -> Option<&'a A> {
385        either_mut!(self.inner, iter => iter.next())
386    }
387
388    fn size_hint(&self) -> (usize, Option<usize>) {
389        either!(self.inner, ref iter => iter.size_hint())
390    }
391
392    fn fold<Acc, G>(self, init: Acc, g: G) -> Acc
393    where
394        G: FnMut(Acc, Self::Item) -> Acc,
395    {
396        either!(self.inner, iter => iter.fold(init, g))
397    }
398
399    fn nth(&mut self, n: usize) -> Option<Self::Item> {
400        either_mut!(self.inner, iter => iter.nth(n))
401    }
402
403    fn collect<B>(self) -> B
404    where
405        B: FromIterator<Self::Item>,
406    {
407        either!(self.inner, iter => iter.collect())
408    }
409
410    fn all<F>(&mut self, f: F) -> bool
411    where
412        F: FnMut(Self::Item) -> bool,
413    {
414        either_mut!(self.inner, iter => iter.all(f))
415    }
416
417    fn any<F>(&mut self, f: F) -> bool
418    where
419        F: FnMut(Self::Item) -> bool,
420    {
421        either_mut!(self.inner, iter => iter.any(f))
422    }
423
424    fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
425    where
426        P: FnMut(&Self::Item) -> bool,
427    {
428        either_mut!(self.inner, iter => iter.find(predicate))
429    }
430
431    fn find_map<B, F>(&mut self, f: F) -> Option<B>
432    where
433        F: FnMut(Self::Item) -> Option<B>,
434    {
435        either_mut!(self.inner, iter => iter.find_map(f))
436    }
437
438    fn count(self) -> usize {
439        either!(self.inner, iter => iter.count())
440    }
441
442    fn last(self) -> Option<Self::Item> {
443        either!(self.inner, iter => iter.last())
444    }
445
446    fn position<P>(&mut self, predicate: P) -> Option<usize>
447    where
448        P: FnMut(Self::Item) -> bool,
449    {
450        either_mut!(self.inner, iter => iter.position(predicate))
451    }
452}
453
454impl<'a, A> DoubleEndedIterator for Iter<'a, A, Ix1> {
455    #[inline]
456    fn next_back(&mut self) -> Option<&'a A> {
457        either_mut!(self.inner, iter => iter.next_back())
458    }
459
460    fn nth_back(&mut self, n: usize) -> Option<&'a A> {
461        either_mut!(self.inner, iter => iter.nth_back(n))
462    }
463
464    fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
465    where
466        G: FnMut(Acc, Self::Item) -> Acc,
467    {
468        either!(self.inner, iter => iter.rfold(init, g))
469    }
470}
471
472impl<'a, A, D> ExactSizeIterator for Iter<'a, A, D>
473where
474    D: Dimension,
475{
476    fn len(&self) -> usize {
477        either!(self.inner, ref iter => iter.len())
478    }
479}
480
481impl<'a, A, D: Dimension> Iterator for IndexedIter<'a, A, D> {
482    type Item = (D::Pattern, &'a A);
483    #[inline]
484    fn next(&mut self) -> Option<Self::Item> {
485        let index = match self.0.inner.index {
486            None => return None,
487            Some(ref ix) => ix.clone(),
488        };
489        match self.0.next() {
490            None => None,
491            Some(elem) => Some((index.into_pattern(), elem)),
492        }
493    }
494
495    fn size_hint(&self) -> (usize, Option<usize>) {
496        self.0.size_hint()
497    }
498}
499
500impl<'a, A, D> ExactSizeIterator for IndexedIter<'a, A, D>
501where
502    D: Dimension,
503{
504    fn len(&self) -> usize {
505        self.0.inner.len()
506    }
507}
508
509impl<'a, A, D: Dimension> Iterator for IterMut<'a, A, D> {
510    type Item = &'a mut A;
511    #[inline]
512    fn next(&mut self) -> Option<&'a mut A> {
513        either_mut!(self.inner, iter => iter.next())
514    }
515
516    fn size_hint(&self) -> (usize, Option<usize>) {
517        either!(self.inner, ref iter => iter.size_hint())
518    }
519
520    fn fold<Acc, G>(self, init: Acc, g: G) -> Acc
521    where
522        G: FnMut(Acc, Self::Item) -> Acc,
523    {
524        either!(self.inner, iter => iter.fold(init, g))
525    }
526
527    fn nth(&mut self, n: usize) -> Option<Self::Item> {
528        either_mut!(self.inner, iter => iter.nth(n))
529    }
530
531    fn collect<B>(self) -> B
532    where
533        B: FromIterator<Self::Item>,
534    {
535        either!(self.inner, iter => iter.collect())
536    }
537
538    fn all<F>(&mut self, f: F) -> bool
539    where
540        F: FnMut(Self::Item) -> bool,
541    {
542        either_mut!(self.inner, iter => iter.all(f))
543    }
544
545    fn any<F>(&mut self, f: F) -> bool
546    where
547        F: FnMut(Self::Item) -> bool,
548    {
549        either_mut!(self.inner, iter => iter.any(f))
550    }
551
552    fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
553    where
554        P: FnMut(&Self::Item) -> bool,
555    {
556        either_mut!(self.inner, iter => iter.find(predicate))
557    }
558
559    fn find_map<B, F>(&mut self, f: F) -> Option<B>
560    where
561        F: FnMut(Self::Item) -> Option<B>,
562    {
563        either_mut!(self.inner, iter => iter.find_map(f))
564    }
565
566    fn count(self) -> usize {
567        either!(self.inner, iter => iter.count())
568    }
569
570    fn last(self) -> Option<Self::Item> {
571        either!(self.inner, iter => iter.last())
572    }
573
574    fn position<P>(&mut self, predicate: P) -> Option<usize>
575    where
576        P: FnMut(Self::Item) -> bool,
577    {
578        either_mut!(self.inner, iter => iter.position(predicate))
579    }
580}
581
582impl<'a, A> DoubleEndedIterator for IterMut<'a, A, Ix1> {
583    #[inline]
584    fn next_back(&mut self) -> Option<&'a mut A> {
585        either_mut!(self.inner, iter => iter.next_back())
586    }
587
588    fn nth_back(&mut self, n: usize) -> Option<&'a mut A> {
589        either_mut!(self.inner, iter => iter.nth_back(n))
590    }
591
592    fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
593    where
594        G: FnMut(Acc, Self::Item) -> Acc,
595    {
596        either!(self.inner, iter => iter.rfold(init, g))
597    }
598}
599
600impl<'a, A, D> ExactSizeIterator for IterMut<'a, A, D>
601where
602    D: Dimension,
603{
604    fn len(&self) -> usize {
605        either!(self.inner, ref iter => iter.len())
606    }
607}
608
609impl<'a, A, D: Dimension> Iterator for ElementsBaseMut<'a, A, D> {
610    type Item = &'a mut A;
611    #[inline]
612    fn next(&mut self) -> Option<&'a mut A> {
613        self.inner.next().map(|p| unsafe { &mut *p })
614    }
615
616    fn size_hint(&self) -> (usize, Option<usize>) {
617        self.inner.size_hint()
618    }
619
620    fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
621    where
622        G: FnMut(Acc, Self::Item) -> Acc,
623    {
624        unsafe { self.inner.fold(init, move |acc, ptr| g(acc, &mut *ptr)) }
625    }
626}
627
628impl<'a, A> DoubleEndedIterator for ElementsBaseMut<'a, A, Ix1> {
629    #[inline]
630    fn next_back(&mut self) -> Option<&'a mut A> {
631        self.inner.next_back().map(|p| unsafe { &mut *p })
632    }
633
634    fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
635    where
636        G: FnMut(Acc, Self::Item) -> Acc,
637    {
638        unsafe { self.inner.rfold(init, move |acc, ptr| g(acc, &mut *ptr)) }
639    }
640}
641
642impl<'a, A, D> ExactSizeIterator for ElementsBaseMut<'a, A, D>
643where
644    D: Dimension,
645{
646    fn len(&self) -> usize {
647        self.inner.len()
648    }
649}
650
651impl<'a, A, D: Dimension> Iterator for IndexedIterMut<'a, A, D> {
652    type Item = (D::Pattern, &'a mut A);
653    #[inline]
654    fn next(&mut self) -> Option<Self::Item> {
655        let index = match self.0.inner.index {
656            None => return None,
657            Some(ref ix) => ix.clone(),
658        };
659        match self.0.next() {
660            None => None,
661            Some(elem) => Some((index.into_pattern(), elem)),
662        }
663    }
664
665    fn size_hint(&self) -> (usize, Option<usize>) {
666        self.0.size_hint()
667    }
668}
669
670impl<'a, A, D> ExactSizeIterator for IndexedIterMut<'a, A, D>
671where
672    D: Dimension,
673{
674    fn len(&self) -> usize {
675        self.0.inner.len()
676    }
677}
678
679/// An iterator that traverses over all axes but one, and yields a view for
680/// each lane along that axis.
681///
682/// See [`.lanes()`](../struct.ArrayBase.html#method.lanes) for more information.
683pub struct LanesIter<'a, A, D> {
684    inner_len: Ix,
685    inner_stride: Ixs,
686    iter: Baseiter<A, D>,
687    life: PhantomData<&'a A>,
688}
689
690clone_bounds!(
691    ['a, A, D: Clone]
692    LanesIter['a, A, D] {
693        @copy {
694            inner_len,
695            inner_stride,
696            life,
697        }
698        iter,
699    }
700);
701
702impl<'a, A, D> Iterator for LanesIter<'a, A, D>
703where
704    D: Dimension,
705{
706    type Item = ArrayView<'a, A, Ix1>;
707    fn next(&mut self) -> Option<Self::Item> {
708        self.iter.next().map(|ptr| unsafe {
709            ArrayView::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
710        })
711    }
712
713    fn size_hint(&self) -> (usize, Option<usize>) {
714        self.iter.size_hint()
715    }
716}
717
718impl<'a, A, D> ExactSizeIterator for LanesIter<'a, A, D>
719where
720    D: Dimension,
721{
722    fn len(&self) -> usize {
723        self.iter.len()
724    }
725}
726
727// NOTE: LanesIterMut is a mutable iterator and must not expose aliasing
728// pointers. Due to this we use an empty slice for the raw data (it's unused
729// anyway).
730/// An iterator that traverses over all dimensions but the innermost,
731/// and yields each inner row (mutable).
732///
733/// See [`.lanes_mut()`](../struct.ArrayBase.html#method.lanes_mut)
734/// for more information.
735pub struct LanesIterMut<'a, A, D> {
736    inner_len: Ix,
737    inner_stride: Ixs,
738    iter: Baseiter<A, D>,
739    life: PhantomData<&'a mut A>,
740}
741
742impl<'a, A, D> Iterator for LanesIterMut<'a, A, D>
743where
744    D: Dimension,
745{
746    type Item = ArrayViewMut<'a, A, Ix1>;
747    fn next(&mut self) -> Option<Self::Item> {
748        self.iter.next().map(|ptr| unsafe {
749            ArrayViewMut::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
750        })
751    }
752
753    fn size_hint(&self) -> (usize, Option<usize>) {
754        self.iter.size_hint()
755    }
756}
757
758impl<'a, A, D> ExactSizeIterator for LanesIterMut<'a, A, D>
759where
760    D: Dimension,
761{
762    fn len(&self) -> usize {
763        self.iter.len()
764    }
765}
766
767#[derive(Debug)]
768pub struct AxisIterCore<A, D> {
769    /// Index along the axis of the value of `.next()`, relative to the start
770    /// of the axis.
771    index: Ix,
772    /// (Exclusive) upper bound on `index`. Initially, this is equal to the
773    /// length of the axis.
774    end: Ix,
775    /// Stride along the axis (offset between consecutive pointers).
776    stride: Ixs,
777    /// Shape of the iterator's items.
778    inner_dim: D,
779    /// Strides of the iterator's items.
780    inner_strides: D,
781    /// Pointer corresponding to `index == 0`.
782    ptr: *mut A,
783}
784
785clone_bounds!(
786    [A, D: Clone]
787    AxisIterCore[A, D] {
788        @copy {
789            index,
790            end,
791            stride,
792            ptr,
793        }
794        inner_dim,
795        inner_strides,
796    }
797);
798
799impl<A, D: Dimension> AxisIterCore<A, D> {
800    /// Constructs a new iterator over the specified axis.
801    fn new<S, Di>(v: ArrayBase<S, Di>, axis: Axis) -> Self
802    where
803        Di: RemoveAxis<Smaller = D>,
804        S: Data<Elem = A>,
805    {
806        AxisIterCore {
807            index: 0,
808            end: v.len_of(axis),
809            stride: v.stride_of(axis),
810            inner_dim: v.dim.remove_axis(axis),
811            inner_strides: v.strides.remove_axis(axis),
812            ptr: v.ptr.as_ptr(),
813        }
814    }
815
816    #[inline]
817    unsafe fn offset(&self, index: usize) -> *mut A {
818        debug_assert!(
819            index < self.end,
820            "index={}, end={}, stride={}",
821            index,
822            self.end,
823            self.stride
824        );
825        self.ptr.offset(index as isize * self.stride)
826    }
827
828    /// Splits the iterator at `index`, yielding two disjoint iterators.
829    ///
830    /// `index` is relative to the current state of the iterator (which is not
831    /// necessarily the start of the axis).
832    ///
833    /// **Panics** if `index` is strictly greater than the iterator's remaining
834    /// length.
835    fn split_at(self, index: usize) -> (Self, Self) {
836        assert!(index <= self.len());
837        let mid = self.index + index;
838        let left = AxisIterCore {
839            index: self.index,
840            end: mid,
841            stride: self.stride,
842            inner_dim: self.inner_dim.clone(),
843            inner_strides: self.inner_strides.clone(),
844            ptr: self.ptr,
845        };
846        let right = AxisIterCore {
847            index: mid,
848            end: self.end,
849            stride: self.stride,
850            inner_dim: self.inner_dim,
851            inner_strides: self.inner_strides,
852            ptr: self.ptr,
853        };
854        (left, right)
855    }
856
857    /// Does the same thing as `.next()` but also returns the index of the item
858    /// relative to the start of the axis.
859    fn next_with_index(&mut self) -> Option<(usize, *mut A)> {
860        let index = self.index;
861        self.next().map(|ptr| (index, ptr))
862    }
863
864    /// Does the same thing as `.next_back()` but also returns the index of the
865    /// item relative to the start of the axis.
866    fn next_back_with_index(&mut self) -> Option<(usize, *mut A)> {
867        self.next_back().map(|ptr| (self.end, ptr))
868    }
869}
870
871impl<A, D> Iterator for AxisIterCore<A, D>
872where
873    D: Dimension,
874{
875    type Item = *mut A;
876
877    fn next(&mut self) -> Option<Self::Item> {
878        if self.index >= self.end {
879            None
880        } else {
881            let ptr = unsafe { self.offset(self.index) };
882            self.index += 1;
883            Some(ptr)
884        }
885    }
886
887    fn size_hint(&self) -> (usize, Option<usize>) {
888        let len = self.len();
889        (len, Some(len))
890    }
891}
892
893impl<A, D> DoubleEndedIterator for AxisIterCore<A, D>
894where
895    D: Dimension,
896{
897    fn next_back(&mut self) -> Option<Self::Item> {
898        if self.index >= self.end {
899            None
900        } else {
901            let ptr = unsafe { self.offset(self.end - 1) };
902            self.end -= 1;
903            Some(ptr)
904        }
905    }
906}
907
908impl<A, D> ExactSizeIterator for AxisIterCore<A, D>
909where
910    D: Dimension,
911{
912    fn len(&self) -> usize {
913        self.end - self.index
914    }
915}
916
917/// An iterator that traverses over an axis and
918/// and yields each subview.
919///
920/// The outermost dimension is `Axis(0)`, created with `.outer_iter()`,
921/// but you can traverse arbitrary dimension with `.axis_iter()`.
922///
923/// For example, in a 3 × 5 × 5 array, with `axis` equal to `Axis(2)`,
924/// the iterator element is a 3 × 5 subview (and there are 5 in total).
925///
926/// Iterator element type is `ArrayView<'a, A, D>`.
927///
928/// See [`.outer_iter()`](../struct.ArrayBase.html#method.outer_iter)
929/// or [`.axis_iter()`](../struct.ArrayBase.html#method.axis_iter)
930/// for more information.
931#[derive(Debug)]
932pub struct AxisIter<'a, A, D> {
933    iter: AxisIterCore<A, D>,
934    life: PhantomData<&'a A>,
935}
936
937clone_bounds!(
938    ['a, A, D: Clone]
939    AxisIter['a, A, D] {
940        @copy {
941            life,
942        }
943        iter,
944    }
945);
946
947impl<'a, A, D: Dimension> AxisIter<'a, A, D> {
948    /// Creates a new iterator over the specified axis.
949    pub(crate) fn new<Di>(v: ArrayView<'a, A, Di>, axis: Axis) -> Self
950    where
951        Di: RemoveAxis<Smaller = D>,
952    {
953        AxisIter {
954            iter: AxisIterCore::new(v, axis),
955            life: PhantomData,
956        }
957    }
958
959    /// Splits the iterator at `index`, yielding two disjoint iterators.
960    ///
961    /// `index` is relative to the current state of the iterator (which is not
962    /// necessarily the start of the axis).
963    ///
964    /// **Panics** if `index` is strictly greater than the iterator's remaining
965    /// length.
966    pub fn split_at(self, index: usize) -> (Self, Self) {
967        let (left, right) = self.iter.split_at(index);
968        (
969            AxisIter {
970                iter: left,
971                life: self.life,
972            },
973            AxisIter {
974                iter: right,
975                life: self.life,
976            },
977        )
978    }
979}
980
981impl<'a, A, D> Iterator for AxisIter<'a, A, D>
982where
983    D: Dimension,
984{
985    type Item = ArrayView<'a, A, D>;
986
987    fn next(&mut self) -> Option<Self::Item> {
988        self.iter.next().map(|ptr| unsafe { self.as_ref(ptr) })
989    }
990
991    fn size_hint(&self) -> (usize, Option<usize>) {
992        self.iter.size_hint()
993    }
994}
995
996impl<'a, A, D> DoubleEndedIterator for AxisIter<'a, A, D>
997where
998    D: Dimension,
999{
1000    fn next_back(&mut self) -> Option<Self::Item> {
1001        self.iter.next_back().map(|ptr| unsafe { self.as_ref(ptr) })
1002    }
1003}
1004
1005impl<'a, A, D> ExactSizeIterator for AxisIter<'a, A, D>
1006where
1007    D: Dimension,
1008{
1009    fn len(&self) -> usize {
1010        self.iter.len()
1011    }
1012}
1013
1014/// An iterator that traverses over an axis and
1015/// and yields each subview (mutable)
1016///
1017/// The outermost dimension is `Axis(0)`, created with `.outer_iter()`,
1018/// but you can traverse arbitrary dimension with `.axis_iter()`.
1019///
1020/// For example, in a 3 × 5 × 5 array, with `axis` equal to `Axis(2)`,
1021/// the iterator element is a 3 × 5 subview (and there are 5 in total).
1022///
1023/// Iterator element type is `ArrayViewMut<'a, A, D>`.
1024///
1025/// See [`.outer_iter_mut()`](../struct.ArrayBase.html#method.outer_iter_mut)
1026/// or [`.axis_iter_mut()`](../struct.ArrayBase.html#method.axis_iter_mut)
1027/// for more information.
1028pub struct AxisIterMut<'a, A, D> {
1029    iter: AxisIterCore<A, D>,
1030    life: PhantomData<&'a mut A>,
1031}
1032
1033impl<'a, A, D: Dimension> AxisIterMut<'a, A, D> {
1034    /// Creates a new iterator over the specified axis.
1035    pub(crate) fn new<Di>(v: ArrayViewMut<'a, A, Di>, axis: Axis) -> Self
1036    where
1037        Di: RemoveAxis<Smaller = D>,
1038    {
1039        AxisIterMut {
1040            iter: AxisIterCore::new(v, axis),
1041            life: PhantomData,
1042        }
1043    }
1044
1045    /// Splits the iterator at `index`, yielding two disjoint iterators.
1046    ///
1047    /// `index` is relative to the current state of the iterator (which is not
1048    /// necessarily the start of the axis).
1049    ///
1050    /// **Panics** if `index` is strictly greater than the iterator's remaining
1051    /// length.
1052    pub fn split_at(self, index: usize) -> (Self, Self) {
1053        let (left, right) = self.iter.split_at(index);
1054        (
1055            AxisIterMut {
1056                iter: left,
1057                life: self.life,
1058            },
1059            AxisIterMut {
1060                iter: right,
1061                life: self.life,
1062            },
1063        )
1064    }
1065}
1066
1067impl<'a, A, D> Iterator for AxisIterMut<'a, A, D>
1068where
1069    D: Dimension,
1070{
1071    type Item = ArrayViewMut<'a, A, D>;
1072
1073    fn next(&mut self) -> Option<Self::Item> {
1074        self.iter.next().map(|ptr| unsafe { self.as_ref(ptr) })
1075    }
1076
1077    fn size_hint(&self) -> (usize, Option<usize>) {
1078        self.iter.size_hint()
1079    }
1080}
1081
1082impl<'a, A, D> DoubleEndedIterator for AxisIterMut<'a, A, D>
1083where
1084    D: Dimension,
1085{
1086    fn next_back(&mut self) -> Option<Self::Item> {
1087        self.iter.next_back().map(|ptr| unsafe { self.as_ref(ptr) })
1088    }
1089}
1090
1091impl<'a, A, D> ExactSizeIterator for AxisIterMut<'a, A, D>
1092where
1093    D: Dimension,
1094{
1095    fn len(&self) -> usize {
1096        self.iter.len()
1097    }
1098}
1099
1100impl<'a, A, D: Dimension> NdProducer for AxisIter<'a, A, D> {
1101    type Item = <Self as Iterator>::Item;
1102    type Dim = Ix1;
1103    type Ptr = *mut A;
1104    type Stride = isize;
1105
1106    #[doc(hidden)]
1107    fn layout(&self) -> crate::Layout {
1108        crate::Layout::one_dimensional()
1109    }
1110    #[doc(hidden)]
1111    fn raw_dim(&self) -> Self::Dim {
1112        Ix1(self.len())
1113    }
1114    #[doc(hidden)]
1115    fn as_ptr(&self) -> Self::Ptr {
1116        if self.len() > 0 {
1117            // `self.iter.index` is guaranteed to be in-bounds if any of the
1118            // iterator remains (i.e. if `self.len() > 0`).
1119            unsafe { self.iter.offset(self.iter.index) }
1120        } else {
1121            // In this case, `self.iter.index` may be past the end, so we must
1122            // not call `.offset()`. It's okay to return a dangling pointer
1123            // because it will never be used in the length 0 case.
1124            std::ptr::NonNull::dangling().as_ptr()
1125        }
1126    }
1127
1128    fn contiguous_stride(&self) -> isize {
1129        self.iter.stride
1130    }
1131
1132    #[doc(hidden)]
1133    unsafe fn as_ref(&self, ptr: Self::Ptr) -> Self::Item {
1134        ArrayView::new_(
1135            ptr,
1136            self.iter.inner_dim.clone(),
1137            self.iter.inner_strides.clone(),
1138        )
1139    }
1140    #[doc(hidden)]
1141    unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
1142        self.iter.offset(self.iter.index + i[0])
1143    }
1144
1145    #[doc(hidden)]
1146    fn stride_of(&self, _axis: Axis) -> isize {
1147        self.contiguous_stride()
1148    }
1149
1150    #[doc(hidden)]
1151    fn split_at(self, _axis: Axis, index: usize) -> (Self, Self) {
1152        self.split_at(index)
1153    }
1154    private_impl! {}
1155}
1156
1157impl<'a, A, D: Dimension> NdProducer for AxisIterMut<'a, A, D> {
1158    type Item = <Self as Iterator>::Item;
1159    type Dim = Ix1;
1160    type Ptr = *mut A;
1161    type Stride = isize;
1162
1163    #[doc(hidden)]
1164    fn layout(&self) -> crate::Layout {
1165        crate::Layout::one_dimensional()
1166    }
1167    #[doc(hidden)]
1168    fn raw_dim(&self) -> Self::Dim {
1169        Ix1(self.len())
1170    }
1171    #[doc(hidden)]
1172    fn as_ptr(&self) -> Self::Ptr {
1173        if self.len() > 0 {
1174            // `self.iter.index` is guaranteed to be in-bounds if any of the
1175            // iterator remains (i.e. if `self.len() > 0`).
1176            unsafe { self.iter.offset(self.iter.index) }
1177        } else {
1178            // In this case, `self.iter.index` may be past the end, so we must
1179            // not call `.offset()`. It's okay to return a dangling pointer
1180            // because it will never be used in the length 0 case.
1181            std::ptr::NonNull::dangling().as_ptr()
1182        }
1183    }
1184
1185    fn contiguous_stride(&self) -> isize {
1186        self.iter.stride
1187    }
1188
1189    #[doc(hidden)]
1190    unsafe fn as_ref(&self, ptr: Self::Ptr) -> Self::Item {
1191        ArrayViewMut::new_(
1192            ptr,
1193            self.iter.inner_dim.clone(),
1194            self.iter.inner_strides.clone(),
1195        )
1196    }
1197    #[doc(hidden)]
1198    unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
1199        self.iter.offset(self.iter.index + i[0])
1200    }
1201
1202    #[doc(hidden)]
1203    fn stride_of(&self, _axis: Axis) -> isize {
1204        self.contiguous_stride()
1205    }
1206
1207    #[doc(hidden)]
1208    fn split_at(self, _axis: Axis, index: usize) -> (Self, Self) {
1209        self.split_at(index)
1210    }
1211    private_impl! {}
1212}
1213
1214/// An iterator that traverses over the specified axis
1215/// and yields views of the specified size on this axis.
1216///
1217/// For example, in a 2 × 8 × 3 array, if the axis of iteration
1218/// is 1 and the chunk size is 2, the yielded elements
1219/// are 2 × 2 × 3 views (and there are 4 in total).
1220///
1221/// Iterator element type is `ArrayView<'a, A, D>`.
1222///
1223/// See [`.axis_chunks_iter()`](../struct.ArrayBase.html#method.axis_chunks_iter) for more information.
1224pub struct AxisChunksIter<'a, A, D> {
1225    iter: AxisIterCore<A, D>,
1226    /// Index of the partial chunk (the chunk smaller than the specified chunk
1227    /// size due to the axis length not being evenly divisible). If the axis
1228    /// length is evenly divisible by the chunk size, this index is larger than
1229    /// the maximum valid index.
1230    partial_chunk_index: usize,
1231    /// Dimension of the partial chunk.
1232    partial_chunk_dim: D,
1233    life: PhantomData<&'a A>,
1234}
1235
1236clone_bounds!(
1237    ['a, A, D: Clone]
1238    AxisChunksIter['a, A, D] {
1239        @copy {
1240            life,
1241            partial_chunk_index,
1242        }
1243        iter,
1244        partial_chunk_dim,
1245    }
1246);
1247
1248/// Computes the information necessary to construct an iterator over chunks
1249/// along an axis, given a `view` of the array, the `axis` to iterate over, and
1250/// the chunk `size`.
1251///
1252/// Returns an axis iterator with the correct stride to move between chunks,
1253/// the number of chunks, and the shape of the last chunk.
1254///
1255/// **Panics** if `size == 0`.
1256fn chunk_iter_parts<A, D: Dimension>(
1257    v: ArrayView<'_, A, D>,
1258    axis: Axis,
1259    size: usize,
1260) -> (AxisIterCore<A, D>, usize, D) {
1261    assert_ne!(size, 0, "Chunk size must be nonzero.");
1262    let axis_len = v.len_of(axis);
1263    let n_whole_chunks = axis_len / size;
1264    let chunk_remainder = axis_len % size;
1265    let iter_len = if chunk_remainder == 0 {
1266        n_whole_chunks
1267    } else {
1268        n_whole_chunks + 1
1269    };
1270    let stride = if n_whole_chunks == 0 {
1271        // This case avoids potential overflow when `size > axis_len`.
1272        0
1273    } else {
1274        v.stride_of(axis) * size as isize
1275    };
1276
1277    let axis = axis.index();
1278    let mut inner_dim = v.dim.clone();
1279    inner_dim[axis] = size;
1280
1281    let mut partial_chunk_dim = v.dim;
1282    partial_chunk_dim[axis] = chunk_remainder;
1283    let partial_chunk_index = n_whole_chunks;
1284
1285    let iter = AxisIterCore {
1286        index: 0,
1287        end: iter_len,
1288        stride,
1289        inner_dim,
1290        inner_strides: v.strides,
1291        ptr: v.ptr.as_ptr(),
1292    };
1293
1294    (iter, partial_chunk_index, partial_chunk_dim)
1295}
1296
1297impl<'a, A, D: Dimension> AxisChunksIter<'a, A, D> {
1298    pub(crate) fn new(v: ArrayView<'a, A, D>, axis: Axis, size: usize) -> Self {
1299        let (iter, partial_chunk_index, partial_chunk_dim) = chunk_iter_parts(v, axis, size);
1300        AxisChunksIter {
1301            iter,
1302            partial_chunk_index,
1303            partial_chunk_dim,
1304            life: PhantomData,
1305        }
1306    }
1307}
1308
1309macro_rules! chunk_iter_impl {
1310    ($iter:ident, $array:ident) => {
1311        impl<'a, A, D> $iter<'a, A, D>
1312        where
1313            D: Dimension,
1314        {
1315            fn get_subview(&self, index: usize, ptr: *mut A) -> $array<'a, A, D> {
1316                if index != self.partial_chunk_index {
1317                    unsafe {
1318                        $array::new_(
1319                            ptr,
1320                            self.iter.inner_dim.clone(),
1321                            self.iter.inner_strides.clone(),
1322                        )
1323                    }
1324                } else {
1325                    unsafe {
1326                        $array::new_(
1327                            ptr,
1328                            self.partial_chunk_dim.clone(),
1329                            self.iter.inner_strides.clone(),
1330                        )
1331                    }
1332                }
1333            }
1334
1335            /// Splits the iterator at index, yielding two disjoint iterators.
1336            ///
1337            /// `index` is relative to the current state of the iterator (which is not
1338            /// necessarily the start of the axis).
1339            ///
1340            /// **Panics** if `index` is strictly greater than the iterator's remaining
1341            /// length.
1342            pub fn split_at(self, index: usize) -> (Self, Self) {
1343                let (left, right) = self.iter.split_at(index);
1344                (
1345                    Self {
1346                        iter: left,
1347                        partial_chunk_index: self.partial_chunk_index,
1348                        partial_chunk_dim: self.partial_chunk_dim.clone(),
1349                        life: self.life,
1350                    },
1351                    Self {
1352                        iter: right,
1353                        partial_chunk_index: self.partial_chunk_index,
1354                        partial_chunk_dim: self.partial_chunk_dim,
1355                        life: self.life,
1356                    },
1357                )
1358            }
1359        }
1360
1361        impl<'a, A, D> Iterator for $iter<'a, A, D>
1362        where
1363            D: Dimension,
1364        {
1365            type Item = $array<'a, A, D>;
1366
1367            fn next(&mut self) -> Option<Self::Item> {
1368                self.iter
1369                    .next_with_index()
1370                    .map(|(index, ptr)| self.get_subview(index, ptr))
1371            }
1372
1373            fn size_hint(&self) -> (usize, Option<usize>) {
1374                self.iter.size_hint()
1375            }
1376        }
1377
1378        impl<'a, A, D> DoubleEndedIterator for $iter<'a, A, D>
1379        where
1380            D: Dimension,
1381        {
1382            fn next_back(&mut self) -> Option<Self::Item> {
1383                self.iter
1384                    .next_back_with_index()
1385                    .map(|(index, ptr)| self.get_subview(index, ptr))
1386            }
1387        }
1388
1389        impl<'a, A, D> ExactSizeIterator for $iter<'a, A, D> where D: Dimension {}
1390    };
1391}
1392
1393/// An iterator that traverses over the specified axis
1394/// and yields mutable views of the specified size on this axis.
1395///
1396/// For example, in a 2 × 8 × 3 array, if the axis of iteration
1397/// is 1 and the chunk size is 2, the yielded elements
1398/// are 2 × 2 × 3 views (and there are 4 in total).
1399///
1400/// Iterator element type is `ArrayViewMut<'a, A, D>`.
1401///
1402/// See [`.axis_chunks_iter_mut()`](../struct.ArrayBase.html#method.axis_chunks_iter_mut)
1403/// for more information.
1404pub struct AxisChunksIterMut<'a, A, D> {
1405    iter: AxisIterCore<A, D>,
1406    partial_chunk_index: usize,
1407    partial_chunk_dim: D,
1408    life: PhantomData<&'a mut A>,
1409}
1410
1411impl<'a, A, D: Dimension> AxisChunksIterMut<'a, A, D> {
1412    pub(crate) fn new(v: ArrayViewMut<'a, A, D>, axis: Axis, size: usize) -> Self {
1413        let (iter, partial_chunk_index, partial_chunk_dim) =
1414            chunk_iter_parts(v.into_view(), axis, size);
1415        AxisChunksIterMut {
1416            iter,
1417            partial_chunk_index,
1418            partial_chunk_dim,
1419            life: PhantomData,
1420        }
1421    }
1422}
1423
1424chunk_iter_impl!(AxisChunksIter, ArrayView);
1425chunk_iter_impl!(AxisChunksIterMut, ArrayViewMut);
1426
1427send_sync_read_only!(Iter);
1428send_sync_read_only!(IndexedIter);
1429send_sync_read_only!(LanesIter);
1430send_sync_read_only!(AxisIter);
1431send_sync_read_only!(AxisChunksIter);
1432send_sync_read_only!(ElementsBase);
1433
1434send_sync_read_write!(IterMut);
1435send_sync_read_write!(IndexedIterMut);
1436send_sync_read_write!(LanesIterMut);
1437send_sync_read_write!(AxisIterMut);
1438send_sync_read_write!(AxisChunksIterMut);
1439send_sync_read_write!(ElementsBaseMut);
1440
1441/// (Trait used internally) An iterator that we trust
1442/// to deliver exactly as many items as it said it would.
1443///
1444/// The iterator must produce exactly the number of elements it reported or
1445/// diverge before reaching the end.
1446pub unsafe trait TrustedIterator {}
1447
1448use crate::indexes::IndicesIterF;
1449use crate::iter::IndicesIter;
1450#[cfg(feature = "std")]
1451use crate::{geomspace::Geomspace, linspace::Linspace, logspace::Logspace};
1452#[cfg(feature = "std")]
1453unsafe impl<F> TrustedIterator for Linspace<F> {}
1454#[cfg(feature = "std")]
1455unsafe impl<F> TrustedIterator for Geomspace<F> {}
1456#[cfg(feature = "std")]
1457unsafe impl<F> TrustedIterator for Logspace<F> {}
1458unsafe impl<'a, A, D> TrustedIterator for Iter<'a, A, D> {}
1459unsafe impl<'a, A, D> TrustedIterator for IterMut<'a, A, D> {}
1460unsafe impl<I> TrustedIterator for std::iter::Cloned<I> where I: TrustedIterator {}
1461unsafe impl<I, F> TrustedIterator for std::iter::Map<I, F> where I: TrustedIterator {}
1462unsafe impl<'a, A> TrustedIterator for slice::Iter<'a, A> {}
1463unsafe impl<'a, A> TrustedIterator for slice::IterMut<'a, A> {}
1464unsafe impl TrustedIterator for ::std::ops::Range<usize> {}
1465// FIXME: These indices iter are dubious -- size needs to be checked up front.
1466unsafe impl<D> TrustedIterator for IndicesIter<D> where D: Dimension {}
1467unsafe impl<D> TrustedIterator for IndicesIterF<D> where D: Dimension {}
1468
1469/// Like Iterator::collect, but only for trusted length iterators
1470pub fn to_vec<I>(iter: I) -> Vec<I::Item>
1471where
1472    I: TrustedIterator + ExactSizeIterator,
1473{
1474    to_vec_mapped(iter, |x| x)
1475}
1476
1477/// Like Iterator::collect, but only for trusted length iterators
1478pub fn to_vec_mapped<I, F, B>(iter: I, mut f: F) -> Vec<B>
1479where
1480    I: TrustedIterator + ExactSizeIterator,
1481    F: FnMut(I::Item) -> B,
1482{
1483    // Use an `unsafe` block to do this efficiently.
1484    // We know that iter will produce exactly .size() elements,
1485    // and the loop can vectorize if it's clean (without branch to grow the vector).
1486    let (size, _) = iter.size_hint();
1487    let mut result = Vec::with_capacity(size);
1488    let mut out_ptr = result.as_mut_ptr();
1489    let mut len = 0;
1490    iter.fold((), |(), elt| unsafe {
1491        ptr::write(out_ptr, f(elt));
1492        len += 1;
1493        result.set_len(len);
1494        out_ptr = out_ptr.offset(1);
1495    });
1496    debug_assert_eq!(size, result.len());
1497    result
1498}