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