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

ndarray/iterators/
mod.rs
