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

ndarray/iterators/
mod.rs
