1#[macro_use]
10mod macros;
11mod chunks;
12pub mod iter;
13mod lanes;
14mod windows;
15
16use std::iter::FromIterator;
17use std::marker::PhantomData;
18use std::ptr;
19use alloc::vec::Vec;
20
21use crate::Ix1;
22
23use super::{ArrayBase, ArrayView, ArrayViewMut, Axis, Data, NdProducer, RemoveAxis};
24use super::{Dimension, Ix, Ixs};
25
26pub use self::chunks::{ExactChunks, ExactChunksIter, ExactChunksIterMut, ExactChunksMut};
27pub use self::lanes::{Lanes, LanesMut};
28pub use self::windows::Windows;
29
30use std::slice::{self, Iter as SliceIter, IterMut as SliceIterMut};
31
32pub struct Baseiter<A, D> {
36 ptr: *mut A,
37 dim: D,
38 strides: D,
39 index: Option<D>,
40}
41
42impl<A, D: Dimension> Baseiter<A, D> {
43 #[inline]
47 pub unsafe fn new(ptr: *mut A, len: D, stride: D) -> Baseiter<A, D> {
48 Baseiter {
49 ptr,
50 index: len.first_index(),
51 dim: len,
52 strides: stride,
53 }
54 }
55}
56
57impl<A, D: Dimension> Iterator for Baseiter<A, D> {
58 type Item = *mut A;
59
60 #[inline]
61 fn next(&mut self) -> Option<*mut A> {
62 let index = match self.index {
63 None => return None,
64 Some(ref ix) => ix.clone(),
65 };
66 let offset = D::stride_offset(&index, &self.strides);
67 self.index = self.dim.next_for(index);
68 unsafe { Some(self.ptr.offset(offset)) }
69 }
70
71 fn size_hint(&self) -> (usize, Option<usize>) {
72 let len = self.len();
73 (len, Some(len))
74 }
75
76 fn fold<Acc, G>(mut self, init: Acc, mut g: G) -> Acc
77 where
78 G: FnMut(Acc, *mut A) -> Acc,
79 {
80 let ndim = self.dim.ndim();
81 debug_assert_ne!(ndim, 0);
82 let mut accum = init;
83 while let Some(mut index) = self.index {
84 let stride = self.strides.last_elem() as isize;
85 let elem_index = index.last_elem();
86 let len = self.dim.last_elem();
87 let offset = D::stride_offset(&index, &self.strides);
88 unsafe {
89 let row_ptr = self.ptr.offset(offset);
90 let mut i = 0;
91 let i_end = len - elem_index;
92 while i < i_end {
93 accum = g(accum, row_ptr.offset(i as isize * stride));
94 i += 1;
95 }
96 }
97 index.set_last_elem(len - 1);
98 self.index = self.dim.next_for(index);
99 }
100 accum
101 }
102}
103
104impl<'a, A, D: Dimension> ExactSizeIterator for Baseiter<A, D> {
105 fn len(&self) -> usize {
106 match self.index {
107 None => 0,
108 Some(ref ix) => {
109 let gone = self
110 .dim
111 .default_strides()
112 .slice()
113 .iter()
114 .zip(ix.slice().iter())
115 .fold(0, |s, (&a, &b)| s + a as usize * b as usize);
116 self.dim.size() - gone
117 }
118 }
119 }
120}
121
122impl<A> DoubleEndedIterator for Baseiter<A, Ix1> {
123 #[inline]
124 fn next_back(&mut self) -> Option<*mut A> {
125 let index = match self.index {
126 None => return None,
127 Some(ix) => ix,
128 };
129 self.dim[0] -= 1;
130 let offset = <_>::stride_offset(&self.dim, &self.strides);
131 if index == self.dim {
132 self.index = None;
133 }
134
135 unsafe { Some(self.ptr.offset(offset)) }
136 }
137
138 fn nth_back(&mut self, n: usize) -> Option<*mut A> {
139 let index = self.index?;
140 let len = self.dim[0] - index[0];
141 if n < len {
142 self.dim[0] -= n + 1;
143 let offset = <_>::stride_offset(&self.dim, &self.strides);
144 if index == self.dim {
145 self.index = None;
146 }
147 unsafe { Some(self.ptr.offset(offset)) }
148 } else {
149 self.index = None;
150 None
151 }
152 }
153
154 fn rfold<Acc, G>(mut self, init: Acc, mut g: G) -> Acc
155 where
156 G: FnMut(Acc, *mut A) -> Acc,
157 {
158 let mut accum = init;
159 if let Some(index) = self.index {
160 let elem_index = index[0];
161 unsafe {
162 while self.dim[0] > elem_index {
164 self.dim[0] -= 1;
165 accum = g(
166 accum,
167 self.ptr
168 .offset(Ix1::stride_offset(&self.dim, &self.strides)),
169 );
170 }
171 }
172 }
173 accum
174 }
175}
176
177clone_bounds!(
178 [A, D: Clone]
179 Baseiter[A, D] {
180 @copy {
181 ptr,
182 }
183 dim,
184 strides,
185 index,
186 }
187);
188
189clone_bounds!(
190 ['a, A, D: Clone]
191 ElementsBase['a, A, D] {
192 @copy {
193 life,
194 }
195 inner,
196 }
197);
198
199impl<'a, A, D: Dimension> ElementsBase<'a, A, D> {
200 pub fn new(v: ArrayView<'a, A, D>) -> Self {
201 ElementsBase {
202 inner: v.into_base_iter(),
203 life: PhantomData,
204 }
205 }
206}
207
208impl<'a, A, D: Dimension> Iterator for ElementsBase<'a, A, D> {
209 type Item = &'a A;
210 #[inline]
211 fn next(&mut self) -> Option<&'a A> {
212 self.inner.next().map(|p| unsafe { &*p })
213 }
214
215 fn size_hint(&self) -> (usize, Option<usize>) {
216 self.inner.size_hint()
217 }
218
219 fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
220 where
221 G: FnMut(Acc, Self::Item) -> Acc,
222 {
223 unsafe { self.inner.fold(init, move |acc, ptr| g(acc, &*ptr)) }
224 }
225}
226
227impl<'a, A> DoubleEndedIterator for ElementsBase<'a, A, Ix1> {
228 #[inline]
229 fn next_back(&mut self) -> Option<&'a A> {
230 self.inner.next_back().map(|p| unsafe { &*p })
231 }
232
233 fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
234 where
235 G: FnMut(Acc, Self::Item) -> Acc,
236 {
237 unsafe { self.inner.rfold(init, move |acc, ptr| g(acc, &*ptr)) }
238 }
239}
240
241impl<'a, A, D> ExactSizeIterator for ElementsBase<'a, A, D>
242where
243 D: Dimension,
244{
245 fn len(&self) -> usize {
246 self.inner.len()
247 }
248}
249
250macro_rules! either {
251 ($value:expr, $inner:pat => $result:expr) => {
252 match $value {
253 ElementsRepr::Slice($inner) => $result,
254 ElementsRepr::Counted($inner) => $result,
255 }
256 };
257}
258
259macro_rules! either_mut {
260 ($value:expr, $inner:ident => $result:expr) => {
261 match $value {
262 ElementsRepr::Slice(ref mut $inner) => $result,
263 ElementsRepr::Counted(ref mut $inner) => $result,
264 }
265 };
266}
267
268clone_bounds!(
269 ['a, A, D: Clone]
270 Iter['a, A, D] {
271 @copy {
272 }
273 inner,
274 }
275);
276
277impl<'a, A, D> Iter<'a, A, D>
278where
279 D: Dimension,
280{
281 pub(crate) fn new(self_: ArrayView<'a, A, D>) -> Self {
282 Iter {
283 inner: if let Some(slc) = self_.to_slice() {
284 ElementsRepr::Slice(slc.iter())
285 } else {
286 ElementsRepr::Counted(self_.into_elements_base())
287 },
288 }
289 }
290}
291
292impl<'a, A, D> IterMut<'a, A, D>
293where
294 D: Dimension,
295{
296 pub(crate) fn new(self_: ArrayViewMut<'a, A, D>) -> Self {
297 IterMut {
298 inner: match self_.try_into_slice() {
299 Ok(x) => ElementsRepr::Slice(x.iter_mut()),
300 Err(self_) => ElementsRepr::Counted(self_.into_elements_base()),
301 },
302 }
303 }
304}
305
306#[derive(Clone)]
307pub enum ElementsRepr<S, C> {
308 Slice(S),
309 Counted(C),
310}
311
312pub struct Iter<'a, A, D> {
318 inner: ElementsRepr<SliceIter<'a, A>, ElementsBase<'a, A, D>>,
319}
320
321pub struct ElementsBase<'a, A, D> {
323 inner: Baseiter<A, D>,
324 life: PhantomData<&'a A>,
325}
326
327pub struct IterMut<'a, A, D> {
333 inner: ElementsRepr<SliceIterMut<'a, A>, ElementsBaseMut<'a, A, D>>,
334}
335
336pub struct ElementsBaseMut<'a, A, D> {
340 inner: Baseiter<A, D>,
341 life: PhantomData<&'a mut A>,
342}
343
344impl<'a, A, D: Dimension> ElementsBaseMut<'a, A, D> {
345 pub fn new(v: ArrayViewMut<'a, A, D>) -> Self {
346 ElementsBaseMut {
347 inner: v.into_base_iter(),
348 life: PhantomData,
349 }
350 }
351}
352
353#[derive(Clone)]
357pub struct IndexedIter<'a, A, D>(ElementsBase<'a, A, D>);
358pub struct IndexedIterMut<'a, A, D>(ElementsBaseMut<'a, A, D>);
362
363impl<'a, A, D> IndexedIter<'a, A, D>
364where
365 D: Dimension,
366{
367 pub(crate) fn new(x: ElementsBase<'a, A, D>) -> Self {
368 IndexedIter(x)
369 }
370}
371
372impl<'a, A, D> IndexedIterMut<'a, A, D>
373where
374 D: Dimension,
375{
376 pub(crate) fn new(x: ElementsBaseMut<'a, A, D>) -> Self {
377 IndexedIterMut(x)
378 }
379}
380
381impl<'a, A, D: Dimension> Iterator for Iter<'a, A, D> {
382 type Item = &'a A;
383 #[inline]
384 fn next(&mut self) -> Option<&'a A> {
385 either_mut!(self.inner, iter => iter.next())
386 }
387
388 fn size_hint(&self) -> (usize, Option<usize>) {
389 either!(self.inner, ref iter => iter.size_hint())
390 }
391
392 fn fold<Acc, G>(self, init: Acc, g: G) -> Acc
393 where
394 G: FnMut(Acc, Self::Item) -> Acc,
395 {
396 either!(self.inner, iter => iter.fold(init, g))
397 }
398
399 fn nth(&mut self, n: usize) -> Option<Self::Item> {
400 either_mut!(self.inner, iter => iter.nth(n))
401 }
402
403 fn collect<B>(self) -> B
404 where
405 B: FromIterator<Self::Item>,
406 {
407 either!(self.inner, iter => iter.collect())
408 }
409
410 fn all<F>(&mut self, f: F) -> bool
411 where
412 F: FnMut(Self::Item) -> bool,
413 {
414 either_mut!(self.inner, iter => iter.all(f))
415 }
416
417 fn any<F>(&mut self, f: F) -> bool
418 where
419 F: FnMut(Self::Item) -> bool,
420 {
421 either_mut!(self.inner, iter => iter.any(f))
422 }
423
424 fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
425 where
426 P: FnMut(&Self::Item) -> bool,
427 {
428 either_mut!(self.inner, iter => iter.find(predicate))
429 }
430
431 fn find_map<B, F>(&mut self, f: F) -> Option<B>
432 where
433 F: FnMut(Self::Item) -> Option<B>,
434 {
435 either_mut!(self.inner, iter => iter.find_map(f))
436 }
437
438 fn count(self) -> usize {
439 either!(self.inner, iter => iter.count())
440 }
441
442 fn last(self) -> Option<Self::Item> {
443 either!(self.inner, iter => iter.last())
444 }
445
446 fn position<P>(&mut self, predicate: P) -> Option<usize>
447 where
448 P: FnMut(Self::Item) -> bool,
449 {
450 either_mut!(self.inner, iter => iter.position(predicate))
451 }
452}
453
454impl<'a, A> DoubleEndedIterator for Iter<'a, A, Ix1> {
455 #[inline]
456 fn next_back(&mut self) -> Option<&'a A> {
457 either_mut!(self.inner, iter => iter.next_back())
458 }
459
460 fn nth_back(&mut self, n: usize) -> Option<&'a A> {
461 either_mut!(self.inner, iter => iter.nth_back(n))
462 }
463
464 fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
465 where
466 G: FnMut(Acc, Self::Item) -> Acc,
467 {
468 either!(self.inner, iter => iter.rfold(init, g))
469 }
470}
471
472impl<'a, A, D> ExactSizeIterator for Iter<'a, A, D>
473where
474 D: Dimension,
475{
476 fn len(&self) -> usize {
477 either!(self.inner, ref iter => iter.len())
478 }
479}
480
481impl<'a, A, D: Dimension> Iterator for IndexedIter<'a, A, D> {
482 type Item = (D::Pattern, &'a A);
483 #[inline]
484 fn next(&mut self) -> Option<Self::Item> {
485 let index = match self.0.inner.index {
486 None => return None,
487 Some(ref ix) => ix.clone(),
488 };
489 match self.0.next() {
490 None => None,
491 Some(elem) => Some((index.into_pattern(), elem)),
492 }
493 }
494
495 fn size_hint(&self) -> (usize, Option<usize>) {
496 self.0.size_hint()
497 }
498}
499
500impl<'a, A, D> ExactSizeIterator for IndexedIter<'a, A, D>
501where
502 D: Dimension,
503{
504 fn len(&self) -> usize {
505 self.0.inner.len()
506 }
507}
508
509impl<'a, A, D: Dimension> Iterator for IterMut<'a, A, D> {
510 type Item = &'a mut A;
511 #[inline]
512 fn next(&mut self) -> Option<&'a mut A> {
513 either_mut!(self.inner, iter => iter.next())
514 }
515
516 fn size_hint(&self) -> (usize, Option<usize>) {
517 either!(self.inner, ref iter => iter.size_hint())
518 }
519
520 fn fold<Acc, G>(self, init: Acc, g: G) -> Acc
521 where
522 G: FnMut(Acc, Self::Item) -> Acc,
523 {
524 either!(self.inner, iter => iter.fold(init, g))
525 }
526
527 fn nth(&mut self, n: usize) -> Option<Self::Item> {
528 either_mut!(self.inner, iter => iter.nth(n))
529 }
530
531 fn collect<B>(self) -> B
532 where
533 B: FromIterator<Self::Item>,
534 {
535 either!(self.inner, iter => iter.collect())
536 }
537
538 fn all<F>(&mut self, f: F) -> bool
539 where
540 F: FnMut(Self::Item) -> bool,
541 {
542 either_mut!(self.inner, iter => iter.all(f))
543 }
544
545 fn any<F>(&mut self, f: F) -> bool
546 where
547 F: FnMut(Self::Item) -> bool,
548 {
549 either_mut!(self.inner, iter => iter.any(f))
550 }
551
552 fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
553 where
554 P: FnMut(&Self::Item) -> bool,
555 {
556 either_mut!(self.inner, iter => iter.find(predicate))
557 }
558
559 fn find_map<B, F>(&mut self, f: F) -> Option<B>
560 where
561 F: FnMut(Self::Item) -> Option<B>,
562 {
563 either_mut!(self.inner, iter => iter.find_map(f))
564 }
565
566 fn count(self) -> usize {
567 either!(self.inner, iter => iter.count())
568 }
569
570 fn last(self) -> Option<Self::Item> {
571 either!(self.inner, iter => iter.last())
572 }
573
574 fn position<P>(&mut self, predicate: P) -> Option<usize>
575 where
576 P: FnMut(Self::Item) -> bool,
577 {
578 either_mut!(self.inner, iter => iter.position(predicate))
579 }
580}
581
582impl<'a, A> DoubleEndedIterator for IterMut<'a, A, Ix1> {
583 #[inline]
584 fn next_back(&mut self) -> Option<&'a mut A> {
585 either_mut!(self.inner, iter => iter.next_back())
586 }
587
588 fn nth_back(&mut self, n: usize) -> Option<&'a mut A> {
589 either_mut!(self.inner, iter => iter.nth_back(n))
590 }
591
592 fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc
593 where
594 G: FnMut(Acc, Self::Item) -> Acc,
595 {
596 either!(self.inner, iter => iter.rfold(init, g))
597 }
598}
599
600impl<'a, A, D> ExactSizeIterator for IterMut<'a, A, D>
601where
602 D: Dimension,
603{
604 fn len(&self) -> usize {
605 either!(self.inner, ref iter => iter.len())
606 }
607}
608
609impl<'a, A, D: Dimension> Iterator for ElementsBaseMut<'a, A, D> {
610 type Item = &'a mut A;
611 #[inline]
612 fn next(&mut self) -> Option<&'a mut A> {
613 self.inner.next().map(|p| unsafe { &mut *p })
614 }
615
616 fn size_hint(&self) -> (usize, Option<usize>) {
617 self.inner.size_hint()
618 }
619
620 fn fold<Acc, G>(self, init: Acc, mut g: G) -> Acc
621 where
622 G: FnMut(Acc, Self::Item) -> Acc,
623 {
624 unsafe { self.inner.fold(init, move |acc, ptr| g(acc, &mut *ptr)) }
625 }
626}
627
628impl<'a, A> DoubleEndedIterator for ElementsBaseMut<'a, A, Ix1> {
629 #[inline]
630 fn next_back(&mut self) -> Option<&'a mut A> {
631 self.inner.next_back().map(|p| unsafe { &mut *p })
632 }
633
634 fn rfold<Acc, G>(self, init: Acc, mut g: G) -> Acc
635 where
636 G: FnMut(Acc, Self::Item) -> Acc,
637 {
638 unsafe { self.inner.rfold(init, move |acc, ptr| g(acc, &mut *ptr)) }
639 }
640}
641
642impl<'a, A, D> ExactSizeIterator for ElementsBaseMut<'a, A, D>
643where
644 D: Dimension,
645{
646 fn len(&self) -> usize {
647 self.inner.len()
648 }
649}
650
651impl<'a, A, D: Dimension> Iterator for IndexedIterMut<'a, A, D> {
652 type Item = (D::Pattern, &'a mut A);
653 #[inline]
654 fn next(&mut self) -> Option<Self::Item> {
655 let index = match self.0.inner.index {
656 None => return None,
657 Some(ref ix) => ix.clone(),
658 };
659 match self.0.next() {
660 None => None,
661 Some(elem) => Some((index.into_pattern(), elem)),
662 }
663 }
664
665 fn size_hint(&self) -> (usize, Option<usize>) {
666 self.0.size_hint()
667 }
668}
669
670impl<'a, A, D> ExactSizeIterator for IndexedIterMut<'a, A, D>
671where
672 D: Dimension,
673{
674 fn len(&self) -> usize {
675 self.0.inner.len()
676 }
677}
678
679pub struct LanesIter<'a, A, D> {
684 inner_len: Ix,
685 inner_stride: Ixs,
686 iter: Baseiter<A, D>,
687 life: PhantomData<&'a A>,
688}
689
690clone_bounds!(
691 ['a, A, D: Clone]
692 LanesIter['a, A, D] {
693 @copy {
694 inner_len,
695 inner_stride,
696 life,
697 }
698 iter,
699 }
700);
701
702impl<'a, A, D> Iterator for LanesIter<'a, A, D>
703where
704 D: Dimension,
705{
706 type Item = ArrayView<'a, A, Ix1>;
707 fn next(&mut self) -> Option<Self::Item> {
708 self.iter.next().map(|ptr| unsafe {
709 ArrayView::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
710 })
711 }
712
713 fn size_hint(&self) -> (usize, Option<usize>) {
714 self.iter.size_hint()
715 }
716}
717
718impl<'a, A, D> ExactSizeIterator for LanesIter<'a, A, D>
719where
720 D: Dimension,
721{
722 fn len(&self) -> usize {
723 self.iter.len()
724 }
725}
726
727pub struct LanesIterMut<'a, A, D> {
736 inner_len: Ix,
737 inner_stride: Ixs,
738 iter: Baseiter<A, D>,
739 life: PhantomData<&'a mut A>,
740}
741
742impl<'a, A, D> Iterator for LanesIterMut<'a, A, D>
743where
744 D: Dimension,
745{
746 type Item = ArrayViewMut<'a, A, Ix1>;
747 fn next(&mut self) -> Option<Self::Item> {
748 self.iter.next().map(|ptr| unsafe {
749 ArrayViewMut::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
750 })
751 }
752
753 fn size_hint(&self) -> (usize, Option<usize>) {
754 self.iter.size_hint()
755 }
756}
757
758impl<'a, A, D> ExactSizeIterator for LanesIterMut<'a, A, D>
759where
760 D: Dimension,
761{
762 fn len(&self) -> usize {
763 self.iter.len()
764 }
765}
766
767#[derive(Debug)]
768pub struct AxisIterCore<A, D> {
769 index: Ix,
772 end: Ix,
775 stride: Ixs,
777 inner_dim: D,
779 inner_strides: D,
781 ptr: *mut A,
783}
784
785clone_bounds!(
786 [A, D: Clone]
787 AxisIterCore[A, D] {
788 @copy {
789 index,
790 end,
791 stride,
792 ptr,
793 }
794 inner_dim,
795 inner_strides,
796 }
797);
798
799impl<A, D: Dimension> AxisIterCore<A, D> {
800 fn new<S, Di>(v: ArrayBase<S, Di>, axis: Axis) -> Self
802 where
803 Di: RemoveAxis<Smaller = D>,
804 S: Data<Elem = A>,
805 {
806 AxisIterCore {
807 index: 0,
808 end: v.len_of(axis),
809 stride: v.stride_of(axis),
810 inner_dim: v.dim.remove_axis(axis),
811 inner_strides: v.strides.remove_axis(axis),
812 ptr: v.ptr.as_ptr(),
813 }
814 }
815
816 #[inline]
817 unsafe fn offset(&self, index: usize) -> *mut A {
818 debug_assert!(
819 index < self.end,
820 "index={}, end={}, stride={}",
821 index,
822 self.end,
823 self.stride
824 );
825 self.ptr.offset(index as isize * self.stride)
826 }
827
828 fn split_at(self, index: usize) -> (Self, Self) {
836 assert!(index <= self.len());
837 let mid = self.index + index;
838 let left = AxisIterCore {
839 index: self.index,
840 end: mid,
841 stride: self.stride,
842 inner_dim: self.inner_dim.clone(),
843 inner_strides: self.inner_strides.clone(),
844 ptr: self.ptr,
845 };
846 let right = AxisIterCore {
847 index: mid,
848 end: self.end,
849 stride: self.stride,
850 inner_dim: self.inner_dim,
851 inner_strides: self.inner_strides,
852 ptr: self.ptr,
853 };
854 (left, right)
855 }
856
857 fn next_with_index(&mut self) -> Option<(usize, *mut A)> {
860 let index = self.index;
861 self.next().map(|ptr| (index, ptr))
862 }
863
864 fn next_back_with_index(&mut self) -> Option<(usize, *mut A)> {
867 self.next_back().map(|ptr| (self.end, ptr))
868 }
869}
870
871impl<A, D> Iterator for AxisIterCore<A, D>
872where
873 D: Dimension,
874{
875 type Item = *mut A;
876
877 fn next(&mut self) -> Option<Self::Item> {
878 if self.index >= self.end {
879 None
880 } else {
881 let ptr = unsafe { self.offset(self.index) };
882 self.index += 1;
883 Some(ptr)
884 }
885 }
886
887 fn size_hint(&self) -> (usize, Option<usize>) {
888 let len = self.len();
889 (len, Some(len))
890 }
891}
892
893impl<A, D> DoubleEndedIterator for AxisIterCore<A, D>
894where
895 D: Dimension,
896{
897 fn next_back(&mut self) -> Option<Self::Item> {
898 if self.index >= self.end {
899 None
900 } else {
901 let ptr = unsafe { self.offset(self.end - 1) };
902 self.end -= 1;
903 Some(ptr)
904 }
905 }
906}
907
908impl<A, D> ExactSizeIterator for AxisIterCore<A, D>
909where
910 D: Dimension,
911{
912 fn len(&self) -> usize {
913 self.end - self.index
914 }
915}
916
917#[derive(Debug)]
932pub struct AxisIter<'a, A, D> {
933 iter: AxisIterCore<A, D>,
934 life: PhantomData<&'a A>,
935}
936
937clone_bounds!(
938 ['a, A, D: Clone]
939 AxisIter['a, A, D] {
940 @copy {
941 life,
942 }
943 iter,
944 }
945);
946
947impl<'a, A, D: Dimension> AxisIter<'a, A, D> {
948 pub(crate) fn new<Di>(v: ArrayView<'a, A, Di>, axis: Axis) -> Self
950 where
951 Di: RemoveAxis<Smaller = D>,
952 {
953 AxisIter {
954 iter: AxisIterCore::new(v, axis),
955 life: PhantomData,
956 }
957 }
958
959 pub fn split_at(self, index: usize) -> (Self, Self) {
967 let (left, right) = self.iter.split_at(index);
968 (
969 AxisIter {
970 iter: left,
971 life: self.life,
972 },
973 AxisIter {
974 iter: right,
975 life: self.life,
976 },
977 )
978 }
979}
980
981impl<'a, A, D> Iterator for AxisIter<'a, A, D>
982where
983 D: Dimension,
984{
985 type Item = ArrayView<'a, A, D>;
986
987 fn next(&mut self) -> Option<Self::Item> {
988 self.iter.next().map(|ptr| unsafe { self.as_ref(ptr) })
989 }
990
991 fn size_hint(&self) -> (usize, Option<usize>) {
992 self.iter.size_hint()
993 }
994}
995
996impl<'a, A, D> DoubleEndedIterator for AxisIter<'a, A, D>
997where
998 D: Dimension,
999{
1000 fn next_back(&mut self) -> Option<Self::Item> {
1001 self.iter.next_back().map(|ptr| unsafe { self.as_ref(ptr) })
1002 }
1003}
1004
1005impl<'a, A, D> ExactSizeIterator for AxisIter<'a, A, D>
1006where
1007 D: Dimension,
1008{
1009 fn len(&self) -> usize {
1010 self.iter.len()
1011 }
1012}
1013
1014pub struct AxisIterMut<'a, A, D> {
1029 iter: AxisIterCore<A, D>,
1030 life: PhantomData<&'a mut A>,
1031}
1032
1033impl<'a, A, D: Dimension> AxisIterMut<'a, A, D> {
1034 pub(crate) fn new<Di>(v: ArrayViewMut<'a, A, Di>, axis: Axis) -> Self
1036 where
1037 Di: RemoveAxis<Smaller = D>,
1038 {
1039 AxisIterMut {
1040 iter: AxisIterCore::new(v, axis),
1041 life: PhantomData,
1042 }
1043 }
1044
1045 pub fn split_at(self, index: usize) -> (Self, Self) {
1053 let (left, right) = self.iter.split_at(index);
1054 (
1055 AxisIterMut {
1056 iter: left,
1057 life: self.life,
1058 },
1059 AxisIterMut {
1060 iter: right,
1061 life: self.life,
1062 },
1063 )
1064 }
1065}
1066
1067impl<'a, A, D> Iterator for AxisIterMut<'a, A, D>
1068where
1069 D: Dimension,
1070{
1071 type Item = ArrayViewMut<'a, A, D>;
1072
1073 fn next(&mut self) -> Option<Self::Item> {
1074 self.iter.next().map(|ptr| unsafe { self.as_ref(ptr) })
1075 }
1076
1077 fn size_hint(&self) -> (usize, Option<usize>) {
1078 self.iter.size_hint()
1079 }
1080}
1081
1082impl<'a, A, D> DoubleEndedIterator for AxisIterMut<'a, A, D>
1083where
1084 D: Dimension,
1085{
1086 fn next_back(&mut self) -> Option<Self::Item> {
1087 self.iter.next_back().map(|ptr| unsafe { self.as_ref(ptr) })
1088 }
1089}
1090
1091impl<'a, A, D> ExactSizeIterator for AxisIterMut<'a, A, D>
1092where
1093 D: Dimension,
1094{
1095 fn len(&self) -> usize {
1096 self.iter.len()
1097 }
1098}
1099
1100impl<'a, A, D: Dimension> NdProducer for AxisIter<'a, A, D> {
1101 type Item = <Self as Iterator>::Item;
1102 type Dim = Ix1;
1103 type Ptr = *mut A;
1104 type Stride = isize;
1105
1106 #[doc(hidden)]
1107 fn layout(&self) -> crate::Layout {
1108 crate::Layout::one_dimensional()
1109 }
1110 #[doc(hidden)]
1111 fn raw_dim(&self) -> Self::Dim {
1112 Ix1(self.len())
1113 }
1114 #[doc(hidden)]
1115 fn as_ptr(&self) -> Self::Ptr {
1116 if self.len() > 0 {
1117 unsafe { self.iter.offset(self.iter.index) }
1120 } else {
1121 std::ptr::NonNull::dangling().as_ptr()
1125 }
1126 }
1127
1128 fn contiguous_stride(&self) -> isize {
1129 self.iter.stride
1130 }
1131
1132 #[doc(hidden)]
1133 unsafe fn as_ref(&self, ptr: Self::Ptr) -> Self::Item {
1134 ArrayView::new_(
1135 ptr,
1136 self.iter.inner_dim.clone(),
1137 self.iter.inner_strides.clone(),
1138 )
1139 }
1140 #[doc(hidden)]
1141 unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
1142 self.iter.offset(self.iter.index + i[0])
1143 }
1144
1145 #[doc(hidden)]
1146 fn stride_of(&self, _axis: Axis) -> isize {
1147 self.contiguous_stride()
1148 }
1149
1150 #[doc(hidden)]
1151 fn split_at(self, _axis: Axis, index: usize) -> (Self, Self) {
1152 self.split_at(index)
1153 }
1154 private_impl! {}
1155}
1156
1157impl<'a, A, D: Dimension> NdProducer for AxisIterMut<'a, A, D> {
1158 type Item = <Self as Iterator>::Item;
1159 type Dim = Ix1;
1160 type Ptr = *mut A;
1161 type Stride = isize;
1162
1163 #[doc(hidden)]
1164 fn layout(&self) -> crate::Layout {
1165 crate::Layout::one_dimensional()
1166 }
1167 #[doc(hidden)]
1168 fn raw_dim(&self) -> Self::Dim {
1169 Ix1(self.len())
1170 }
1171 #[doc(hidden)]
1172 fn as_ptr(&self) -> Self::Ptr {
1173 if self.len() > 0 {
1174 unsafe { self.iter.offset(self.iter.index) }
1177 } else {
1178 std::ptr::NonNull::dangling().as_ptr()
1182 }
1183 }
1184
1185 fn contiguous_stride(&self) -> isize {
1186 self.iter.stride
1187 }
1188
1189 #[doc(hidden)]
1190 unsafe fn as_ref(&self, ptr: Self::Ptr) -> Self::Item {
1191 ArrayViewMut::new_(
1192 ptr,
1193 self.iter.inner_dim.clone(),
1194 self.iter.inner_strides.clone(),
1195 )
1196 }
1197 #[doc(hidden)]
1198 unsafe fn uget_ptr(&self, i: &Self::Dim) -> Self::Ptr {
1199 self.iter.offset(self.iter.index + i[0])
1200 }
1201
1202 #[doc(hidden)]
1203 fn stride_of(&self, _axis: Axis) -> isize {
1204 self.contiguous_stride()
1205 }
1206
1207 #[doc(hidden)]
1208 fn split_at(self, _axis: Axis, index: usize) -> (Self, Self) {
1209 self.split_at(index)
1210 }
1211 private_impl! {}
1212}
1213
1214pub struct AxisChunksIter<'a, A, D> {
1225 iter: AxisIterCore<A, D>,
1226 partial_chunk_index: usize,
1231 partial_chunk_dim: D,
1233 life: PhantomData<&'a A>,
1234}
1235
1236clone_bounds!(
1237 ['a, A, D: Clone]
1238 AxisChunksIter['a, A, D] {
1239 @copy {
1240 life,
1241 partial_chunk_index,
1242 }
1243 iter,
1244 partial_chunk_dim,
1245 }
1246);
1247
1248fn chunk_iter_parts<A, D: Dimension>(
1257 v: ArrayView<'_, A, D>,
1258 axis: Axis,
1259 size: usize,
1260) -> (AxisIterCore<A, D>, usize, D) {
1261 assert_ne!(size, 0, "Chunk size must be nonzero.");
1262 let axis_len = v.len_of(axis);
1263 let n_whole_chunks = axis_len / size;
1264 let chunk_remainder = axis_len % size;
1265 let iter_len = if chunk_remainder == 0 {
1266 n_whole_chunks
1267 } else {
1268 n_whole_chunks + 1
1269 };
1270 let stride = if n_whole_chunks == 0 {
1271 0
1273 } else {
1274 v.stride_of(axis) * size as isize
1275 };
1276
1277 let axis = axis.index();
1278 let mut inner_dim = v.dim.clone();
1279 inner_dim[axis] = size;
1280
1281 let mut partial_chunk_dim = v.dim;
1282 partial_chunk_dim[axis] = chunk_remainder;
1283 let partial_chunk_index = n_whole_chunks;
1284
1285 let iter = AxisIterCore {
1286 index: 0,
1287 end: iter_len,
1288 stride,
1289 inner_dim,
1290 inner_strides: v.strides,
1291 ptr: v.ptr.as_ptr(),
1292 };
1293
1294 (iter, partial_chunk_index, partial_chunk_dim)
1295}
1296
1297impl<'a, A, D: Dimension> AxisChunksIter<'a, A, D> {
1298 pub(crate) fn new(v: ArrayView<'a, A, D>, axis: Axis, size: usize) -> Self {
1299 let (iter, partial_chunk_index, partial_chunk_dim) = chunk_iter_parts(v, axis, size);
1300 AxisChunksIter {
1301 iter,
1302 partial_chunk_index,
1303 partial_chunk_dim,
1304 life: PhantomData,
1305 }
1306 }
1307}
1308
1309macro_rules! chunk_iter_impl {
1310 ($iter:ident, $array:ident) => {
1311 impl<'a, A, D> $iter<'a, A, D>
1312 where
1313 D: Dimension,
1314 {
1315 fn get_subview(&self, index: usize, ptr: *mut A) -> $array<'a, A, D> {
1316 if index != self.partial_chunk_index {
1317 unsafe {
1318 $array::new_(
1319 ptr,
1320 self.iter.inner_dim.clone(),
1321 self.iter.inner_strides.clone(),
1322 )
1323 }
1324 } else {
1325 unsafe {
1326 $array::new_(
1327 ptr,
1328 self.partial_chunk_dim.clone(),
1329 self.iter.inner_strides.clone(),
1330 )
1331 }
1332 }
1333 }
1334
1335 pub fn split_at(self, index: usize) -> (Self, Self) {
1343 let (left, right) = self.iter.split_at(index);
1344 (
1345 Self {
1346 iter: left,
1347 partial_chunk_index: self.partial_chunk_index,
1348 partial_chunk_dim: self.partial_chunk_dim.clone(),
1349 life: self.life,
1350 },
1351 Self {
1352 iter: right,
1353 partial_chunk_index: self.partial_chunk_index,
1354 partial_chunk_dim: self.partial_chunk_dim,
1355 life: self.life,
1356 },
1357 )
1358 }
1359 }
1360
1361 impl<'a, A, D> Iterator for $iter<'a, A, D>
1362 where
1363 D: Dimension,
1364 {
1365 type Item = $array<'a, A, D>;
1366
1367 fn next(&mut self) -> Option<Self::Item> {
1368 self.iter
1369 .next_with_index()
1370 .map(|(index, ptr)| self.get_subview(index, ptr))
1371 }
1372
1373 fn size_hint(&self) -> (usize, Option<usize>) {
1374 self.iter.size_hint()
1375 }
1376 }
1377
1378 impl<'a, A, D> DoubleEndedIterator for $iter<'a, A, D>
1379 where
1380 D: Dimension,
1381 {
1382 fn next_back(&mut self) -> Option<Self::Item> {
1383 self.iter
1384 .next_back_with_index()
1385 .map(|(index, ptr)| self.get_subview(index, ptr))
1386 }
1387 }
1388
1389 impl<'a, A, D> ExactSizeIterator for $iter<'a, A, D> where D: Dimension {}
1390 };
1391}
1392
1393pub struct AxisChunksIterMut<'a, A, D> {
1405 iter: AxisIterCore<A, D>,
1406 partial_chunk_index: usize,
1407 partial_chunk_dim: D,
1408 life: PhantomData<&'a mut A>,
1409}
1410
1411impl<'a, A, D: Dimension> AxisChunksIterMut<'a, A, D> {
1412 pub(crate) fn new(v: ArrayViewMut<'a, A, D>, axis: Axis, size: usize) -> Self {
1413 let (iter, partial_chunk_index, partial_chunk_dim) =
1414 chunk_iter_parts(v.into_view(), axis, size);
1415 AxisChunksIterMut {
1416 iter,
1417 partial_chunk_index,
1418 partial_chunk_dim,
1419 life: PhantomData,
1420 }
1421 }
1422}
1423
1424chunk_iter_impl!(AxisChunksIter, ArrayView);
1425chunk_iter_impl!(AxisChunksIterMut, ArrayViewMut);
1426
1427send_sync_read_only!(Iter);
1428send_sync_read_only!(IndexedIter);
1429send_sync_read_only!(LanesIter);
1430send_sync_read_only!(AxisIter);
1431send_sync_read_only!(AxisChunksIter);
1432send_sync_read_only!(ElementsBase);
1433
1434send_sync_read_write!(IterMut);
1435send_sync_read_write!(IndexedIterMut);
1436send_sync_read_write!(LanesIterMut);
1437send_sync_read_write!(AxisIterMut);
1438send_sync_read_write!(AxisChunksIterMut);
1439send_sync_read_write!(ElementsBaseMut);
1440
1441pub unsafe trait TrustedIterator {}
1447
1448use crate::indexes::IndicesIterF;
1449use crate::iter::IndicesIter;
1450#[cfg(feature = "std")]
1451use crate::{geomspace::Geomspace, linspace::Linspace, logspace::Logspace};
1452#[cfg(feature = "std")]
1453unsafe impl<F> TrustedIterator for Linspace<F> {}
1454#[cfg(feature = "std")]
1455unsafe impl<F> TrustedIterator for Geomspace<F> {}
1456#[cfg(feature = "std")]
1457unsafe impl<F> TrustedIterator for Logspace<F> {}
1458unsafe impl<'a, A, D> TrustedIterator for Iter<'a, A, D> {}
1459unsafe impl<'a, A, D> TrustedIterator for IterMut<'a, A, D> {}
1460unsafe impl<I> TrustedIterator for std::iter::Cloned<I> where I: TrustedIterator {}
1461unsafe impl<I, F> TrustedIterator for std::iter::Map<I, F> where I: TrustedIterator {}
1462unsafe impl<'a, A> TrustedIterator for slice::Iter<'a, A> {}
1463unsafe impl<'a, A> TrustedIterator for slice::IterMut<'a, A> {}
1464unsafe impl TrustedIterator for ::std::ops::Range<usize> {}
1465unsafe impl<D> TrustedIterator for IndicesIter<D> where D: Dimension {}
1467unsafe impl<D> TrustedIterator for IndicesIterF<D> where D: Dimension {}
1468
1469pub fn to_vec<I>(iter: I) -> Vec<I::Item>
1471where
1472 I: TrustedIterator + ExactSizeIterator,
1473{
1474 to_vec_mapped(iter, |x| x)
1475}
1476
1477pub fn to_vec_mapped<I, F, B>(iter: I, mut f: F) -> Vec<B>
1479where
1480 I: TrustedIterator + ExactSizeIterator,
1481 F: FnMut(I::Item) -> B,
1482{
1483 let (size, _) = iter.size_hint();
1487 let mut result = Vec::with_capacity(size);
1488 let mut out_ptr = result.as_mut_ptr();
1489 let mut len = 0;
1490 iter.fold((), |(), elt| unsafe {
1491 ptr::write(out_ptr, f(elt));
1492 len += 1;
1493 result.set_len(len);
1494 out_ptr = out_ptr.offset(1);
1495 });
1496 debug_assert_eq!(size, result.len());
1497 result
1498}