1use crate::iter::adapters::SourceIter;
2use crate::iter::{
3 Cloned, Copied, Empty, Filter, FilterMap, Fuse, FusedIterator, Map, Once, OnceWith,
4 TrustedFused, TrustedLen,
5};
6use crate::num::NonZero;
7use crate::ops::{ControlFlow, Try};
8use crate::{array, fmt, option, result};
9
10/// An iterator that maps each element to an iterator, and yields the elements
11/// of the produced iterators.
12///
13/// This `struct` is created by [`Iterator::flat_map`]. See its documentation
14/// for more.
15#[must_use = "iterators are lazy and do nothing unless consumed"]
16#[stable(feature = "rust1", since = "1.0.0")]
17pub struct FlatMap<I, U: IntoIterator, F> {
18 inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>,
19}
20
21impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> {
22 pub(in crate::iter) fn new(iter: I, f: F) -> FlatMap<I, U, F> {
23 FlatMap { inner: FlattenCompat::new(iter.map(f)) }
24 }
25
26 pub(crate) fn into_parts(self) -> (Option<U::IntoIter>, Option<I>, Option<U::IntoIter>) {
27 (
28 self.inner.frontiter,
29 self.inner.iter.into_inner().map(Map::into_inner),
30 self.inner.backiter,
31 )
32 }
33}
34
35#[stable(feature = "rust1", since = "1.0.0")]
36impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F>
37where
38 U: Clone + IntoIterator<IntoIter: Clone>,
39{
40 fn clone(&self) -> Self {
41 FlatMap { inner: self.inner.clone() }
42 }
43}
44
45#[stable(feature = "core_impl_debug", since = "1.9.0")]
46impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F>
47where
48 U: IntoIterator<IntoIter: fmt::Debug>,
49{
50 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51 f.debug_struct("FlatMap").field(name:"inner", &self.inner).finish()
52 }
53}
54
55#[stable(feature = "rust1", since = "1.0.0")]
56impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
57where
58 F: FnMut(I::Item) -> U,
59{
60 type Item = U::Item;
61
62 #[inline]
63 fn next(&mut self) -> Option<U::Item> {
64 self.inner.next()
65 }
66
67 #[inline]
68 fn size_hint(&self) -> (usize, Option<usize>) {
69 self.inner.size_hint()
70 }
71
72 #[inline]
73 fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
74 where
75 Self: Sized,
76 Fold: FnMut(Acc, Self::Item) -> R,
77 R: Try<Output = Acc>,
78 {
79 self.inner.try_fold(init, fold)
80 }
81
82 #[inline]
83 fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
84 where
85 Fold: FnMut(Acc, Self::Item) -> Acc,
86 {
87 self.inner.fold(init, fold)
88 }
89
90 #[inline]
91 fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
92 self.inner.advance_by(n)
93 }
94
95 #[inline]
96 fn count(self) -> usize {
97 self.inner.count()
98 }
99
100 #[inline]
101 fn last(self) -> Option<Self::Item> {
102 self.inner.last()
103 }
104}
105
106#[stable(feature = "rust1", since = "1.0.0")]
107impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F>
108where
109 F: FnMut(I::Item) -> U,
110 U: IntoIterator<IntoIter: DoubleEndedIterator>,
111{
112 #[inline]
113 fn next_back(&mut self) -> Option<U::Item> {
114 self.inner.next_back()
115 }
116
117 #[inline]
118 fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
119 where
120 Self: Sized,
121 Fold: FnMut(Acc, Self::Item) -> R,
122 R: Try<Output = Acc>,
123 {
124 self.inner.try_rfold(init, fold)
125 }
126
127 #[inline]
128 fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
129 where
130 Fold: FnMut(Acc, Self::Item) -> Acc,
131 {
132 self.inner.rfold(init, fold)
133 }
134
135 #[inline]
136 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
137 self.inner.advance_back_by(n)
138 }
139}
140
141#[stable(feature = "fused", since = "1.26.0")]
142impl<I, U, F> FusedIterator for FlatMap<I, U, F>
143where
144 I: FusedIterator,
145 U: IntoIterator,
146 F: FnMut(I::Item) -> U,
147{
148}
149
150#[unstable(feature = "trusted_len", issue = "37572")]
151unsafe impl<I, U, F> TrustedLen for FlatMap<I, U, F>
152where
153 I: Iterator,
154 U: IntoIterator,
155 F: FnMut(I::Item) -> U,
156 FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>: TrustedLen,
157{
158}
159
160#[unstable(issue = "none", feature = "inplace_iteration")]
161unsafe impl<I, U, F> SourceIter for FlatMap<I, U, F>
162where
163 I: SourceIter + TrustedFused,
164 U: IntoIterator,
165{
166 type Source = I::Source;
167
168 #[inline]
169 unsafe fn as_inner(&mut self) -> &mut I::Source {
170 // SAFETY: unsafe function forwarding to unsafe function with the same requirements
171 unsafe { SourceIter::as_inner(&mut self.inner.iter) }
172 }
173}
174
175/// An iterator that flattens one level of nesting in an iterator of things
176/// that can be turned into iterators.
177///
178/// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its
179/// documentation for more.
180///
181/// [`flatten`]: Iterator::flatten()
182#[must_use = "iterators are lazy and do nothing unless consumed"]
183#[stable(feature = "iterator_flatten", since = "1.29.0")]
184pub struct Flatten<I: Iterator<Item: IntoIterator>> {
185 inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
186}
187
188impl<I: Iterator<Item: IntoIterator>> Flatten<I> {
189 pub(in super::super) fn new(iter: I) -> Flatten<I> {
190 Flatten { inner: FlattenCompat::new(iter) }
191 }
192}
193
194#[stable(feature = "iterator_flatten", since = "1.29.0")]
195impl<I, U> fmt::Debug for Flatten<I>
196where
197 I: fmt::Debug + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
198 U: fmt::Debug + Iterator,
199{
200 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
201 f.debug_struct("Flatten").field(name:"inner", &self.inner).finish()
202 }
203}
204
205#[stable(feature = "iterator_flatten", since = "1.29.0")]
206impl<I, U> Clone for Flatten<I>
207where
208 I: Clone + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
209 U: Clone + Iterator,
210{
211 fn clone(&self) -> Self {
212 Flatten { inner: self.inner.clone() }
213 }
214}
215
216#[stable(feature = "iterator_flatten", since = "1.29.0")]
217impl<I, U> Iterator for Flatten<I>
218where
219 I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
220 U: Iterator,
221{
222 type Item = U::Item;
223
224 #[inline]
225 fn next(&mut self) -> Option<U::Item> {
226 self.inner.next()
227 }
228
229 #[inline]
230 fn size_hint(&self) -> (usize, Option<usize>) {
231 self.inner.size_hint()
232 }
233
234 #[inline]
235 fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
236 where
237 Self: Sized,
238 Fold: FnMut(Acc, Self::Item) -> R,
239 R: Try<Output = Acc>,
240 {
241 self.inner.try_fold(init, fold)
242 }
243
244 #[inline]
245 fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
246 where
247 Fold: FnMut(Acc, Self::Item) -> Acc,
248 {
249 self.inner.fold(init, fold)
250 }
251
252 #[inline]
253 fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
254 self.inner.advance_by(n)
255 }
256
257 #[inline]
258 fn count(self) -> usize {
259 self.inner.count()
260 }
261
262 #[inline]
263 fn last(self) -> Option<Self::Item> {
264 self.inner.last()
265 }
266}
267
268#[stable(feature = "iterator_flatten", since = "1.29.0")]
269impl<I, U> DoubleEndedIterator for Flatten<I>
270where
271 I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
272 U: DoubleEndedIterator,
273{
274 #[inline]
275 fn next_back(&mut self) -> Option<U::Item> {
276 self.inner.next_back()
277 }
278
279 #[inline]
280 fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
281 where
282 Self: Sized,
283 Fold: FnMut(Acc, Self::Item) -> R,
284 R: Try<Output = Acc>,
285 {
286 self.inner.try_rfold(init, fold)
287 }
288
289 #[inline]
290 fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
291 where
292 Fold: FnMut(Acc, Self::Item) -> Acc,
293 {
294 self.inner.rfold(init, fold)
295 }
296
297 #[inline]
298 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
299 self.inner.advance_back_by(n)
300 }
301}
302
303#[stable(feature = "iterator_flatten", since = "1.29.0")]
304impl<I, U> FusedIterator for Flatten<I>
305where
306 I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
307 U: Iterator,
308{
309}
310
311#[unstable(feature = "trusted_len", issue = "37572")]
312unsafe impl<I> TrustedLen for Flatten<I>
313where
314 I: Iterator<Item: IntoIterator>,
315 FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>: TrustedLen,
316{
317}
318
319#[unstable(issue = "none", feature = "inplace_iteration")]
320unsafe impl<I> SourceIter for Flatten<I>
321where
322 I: SourceIter + TrustedFused + Iterator,
323 <I as Iterator>::Item: IntoIterator,
324{
325 type Source = I::Source;
326
327 #[inline]
328 unsafe fn as_inner(&mut self) -> &mut I::Source {
329 // SAFETY: unsafe function forwarding to unsafe function with the same requirements
330 unsafe { SourceIter::as_inner(&mut self.inner.iter) }
331 }
332}
333
334#[stable(feature = "default_iters", since = "1.70.0")]
335impl<I> Default for Flatten<I>
336where
337 I: Default + Iterator<Item: IntoIterator>,
338{
339 /// Creates a `Flatten` iterator from the default value of `I`.
340 ///
341 /// ```
342 /// # use core::slice;
343 /// # use std::iter::Flatten;
344 /// let iter: Flatten<slice::Iter<'_, [u8; 4]>> = Default::default();
345 /// assert_eq!(iter.count(), 0);
346 /// ```
347 fn default() -> Self {
348 Flatten::new(iter:Default::default())
349 }
350}
351
352/// Real logic of both `Flatten` and `FlatMap` which simply delegate to
353/// this type.
354#[derive(Clone, Debug)]
355#[unstable(feature = "trusted_len", issue = "37572")]
356struct FlattenCompat<I, U> {
357 iter: Fuse<I>,
358 frontiter: Option<U>,
359 backiter: Option<U>,
360}
361impl<I, U> FlattenCompat<I, U>
362where
363 I: Iterator,
364{
365 /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
366 fn new(iter: I) -> FlattenCompat<I, U> {
367 FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None }
368 }
369}
370
371impl<I, U> FlattenCompat<I, U>
372where
373 I: Iterator<Item: IntoIterator<IntoIter = U>>,
374{
375 /// Folds the inner iterators into an accumulator by applying an operation.
376 ///
377 /// Folds over the inner iterators, not over their elements. Is used by the `fold`, `count`,
378 /// and `last` methods.
379 #[inline]
380 fn iter_fold<Acc, Fold>(self, mut acc: Acc, mut fold: Fold) -> Acc
381 where
382 Fold: FnMut(Acc, U) -> Acc,
383 {
384 #[inline]
385 fn flatten<T: IntoIterator, Acc>(
386 fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc,
387 ) -> impl FnMut(Acc, T) -> Acc + '_ {
388 move |acc, iter| fold(acc, iter.into_iter())
389 }
390
391 if let Some(iter) = self.frontiter {
392 acc = fold(acc, iter);
393 }
394
395 acc = self.iter.fold(acc, flatten(&mut fold));
396
397 if let Some(iter) = self.backiter {
398 acc = fold(acc, iter);
399 }
400
401 acc
402 }
403
404 /// Folds over the inner iterators as long as the given function returns successfully,
405 /// always storing the most recent inner iterator in `self.frontiter`.
406 ///
407 /// Folds over the inner iterators, not over their elements. Is used by the `try_fold` and
408 /// `advance_by` methods.
409 #[inline]
410 fn iter_try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, mut fold: Fold) -> R
411 where
412 Fold: FnMut(Acc, &mut U) -> R,
413 R: Try<Output = Acc>,
414 {
415 #[inline]
416 fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>(
417 frontiter: &'a mut Option<T::IntoIter>,
418 fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R,
419 ) -> impl FnMut(Acc, T) -> R + 'a {
420 move |acc, iter| fold(acc, frontiter.insert(iter.into_iter()))
421 }
422
423 if let Some(iter) = &mut self.frontiter {
424 acc = fold(acc, iter)?;
425 }
426 self.frontiter = None;
427
428 acc = self.iter.try_fold(acc, flatten(&mut self.frontiter, &mut fold))?;
429 self.frontiter = None;
430
431 if let Some(iter) = &mut self.backiter {
432 acc = fold(acc, iter)?;
433 }
434 self.backiter = None;
435
436 try { acc }
437 }
438}
439
440impl<I, U> FlattenCompat<I, U>
441where
442 I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U>>,
443{
444 /// Folds the inner iterators into an accumulator by applying an operation, starting form the
445 /// back.
446 ///
447 /// Folds over the inner iterators, not over their elements. Is used by the `rfold` method.
448 #[inline]
449 fn iter_rfold<Acc, Fold>(self, mut acc: Acc, mut fold: Fold) -> Acc
450 where
451 Fold: FnMut(Acc, U) -> Acc,
452 {
453 #[inline]
454 fn flatten<T: IntoIterator, Acc>(
455 fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc,
456 ) -> impl FnMut(Acc, T) -> Acc + '_ {
457 move |acc, iter| fold(acc, iter.into_iter())
458 }
459
460 if let Some(iter) = self.backiter {
461 acc = fold(acc, iter);
462 }
463
464 acc = self.iter.rfold(acc, flatten(&mut fold));
465
466 if let Some(iter) = self.frontiter {
467 acc = fold(acc, iter);
468 }
469
470 acc
471 }
472
473 /// Folds over the inner iterators in reverse order as long as the given function returns
474 /// successfully, always storing the most recent inner iterator in `self.backiter`.
475 ///
476 /// Folds over the inner iterators, not over their elements. Is used by the `try_rfold` and
477 /// `advance_back_by` methods.
478 #[inline]
479 fn iter_try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, mut fold: Fold) -> R
480 where
481 Fold: FnMut(Acc, &mut U) -> R,
482 R: Try<Output = Acc>,
483 {
484 #[inline]
485 fn flatten<'a, T: IntoIterator, Acc, R: Try>(
486 backiter: &'a mut Option<T::IntoIter>,
487 fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R,
488 ) -> impl FnMut(Acc, T) -> R + 'a {
489 move |acc, iter| fold(acc, backiter.insert(iter.into_iter()))
490 }
491
492 if let Some(iter) = &mut self.backiter {
493 acc = fold(acc, iter)?;
494 }
495 self.backiter = None;
496
497 acc = self.iter.try_rfold(acc, flatten(&mut self.backiter, &mut fold))?;
498 self.backiter = None;
499
500 if let Some(iter) = &mut self.frontiter {
501 acc = fold(acc, iter)?;
502 }
503 self.frontiter = None;
504
505 try { acc }
506 }
507}
508
509// See also the `OneShot` specialization below.
510impl<I, U> Iterator for FlattenCompat<I, U>
511where
512 I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
513 U: Iterator,
514{
515 type Item = U::Item;
516
517 #[inline]
518 default fn next(&mut self) -> Option<U::Item> {
519 loop {
520 if let elt @ Some(_) = and_then_or_clear(&mut self.frontiter, Iterator::next) {
521 return elt;
522 }
523 match self.iter.next() {
524 None => return and_then_or_clear(&mut self.backiter, Iterator::next),
525 Some(inner) => self.frontiter = Some(inner.into_iter()),
526 }
527 }
528 }
529
530 #[inline]
531 default fn size_hint(&self) -> (usize, Option<usize>) {
532 let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
533 let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
534 let lo = flo.saturating_add(blo);
535
536 if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIntoIterator>::size() {
537 let (lower, upper) = self.iter.size_hint();
538
539 let lower = lower.saturating_mul(fixed_size).saturating_add(lo);
540 let upper =
541 try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? };
542
543 return (lower, upper);
544 }
545
546 match (self.iter.size_hint(), fhi, bhi) {
547 ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
548 _ => (lo, None),
549 }
550 }
551
552 #[inline]
553 default fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
554 where
555 Self: Sized,
556 Fold: FnMut(Acc, Self::Item) -> R,
557 R: Try<Output = Acc>,
558 {
559 #[inline]
560 fn flatten<U: Iterator, Acc, R: Try<Output = Acc>>(
561 mut fold: impl FnMut(Acc, U::Item) -> R,
562 ) -> impl FnMut(Acc, &mut U) -> R {
563 move |acc, iter| iter.try_fold(acc, &mut fold)
564 }
565
566 self.iter_try_fold(init, flatten(fold))
567 }
568
569 #[inline]
570 default fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
571 where
572 Fold: FnMut(Acc, Self::Item) -> Acc,
573 {
574 #[inline]
575 fn flatten<U: Iterator, Acc>(
576 mut fold: impl FnMut(Acc, U::Item) -> Acc,
577 ) -> impl FnMut(Acc, U) -> Acc {
578 move |acc, iter| iter.fold(acc, &mut fold)
579 }
580
581 self.iter_fold(init, flatten(fold))
582 }
583
584 #[inline]
585 #[rustc_inherit_overflow_checks]
586 default fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
587 #[inline]
588 #[rustc_inherit_overflow_checks]
589 fn advance<U: Iterator>(n: usize, iter: &mut U) -> ControlFlow<(), usize> {
590 match iter.advance_by(n) {
591 Ok(()) => ControlFlow::Break(()),
592 Err(remaining) => ControlFlow::Continue(remaining.get()),
593 }
594 }
595
596 match self.iter_try_fold(n, advance) {
597 ControlFlow::Continue(remaining) => NonZero::new(remaining).map_or(Ok(()), Err),
598 _ => Ok(()),
599 }
600 }
601
602 #[inline]
603 default fn count(self) -> usize {
604 #[inline]
605 #[rustc_inherit_overflow_checks]
606 fn count<U: Iterator>(acc: usize, iter: U) -> usize {
607 acc + iter.count()
608 }
609
610 self.iter_fold(0, count)
611 }
612
613 #[inline]
614 default fn last(self) -> Option<Self::Item> {
615 #[inline]
616 fn last<U: Iterator>(last: Option<U::Item>, iter: U) -> Option<U::Item> {
617 iter.last().or(last)
618 }
619
620 self.iter_fold(None, last)
621 }
622}
623
624// See also the `OneShot` specialization below.
625impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
626where
627 I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
628 U: DoubleEndedIterator,
629{
630 #[inline]
631 default fn next_back(&mut self) -> Option<U::Item> {
632 loop {
633 if let elt @ Some(_) = and_then_or_clear(&mut self.backiter, |b| b.next_back()) {
634 return elt;
635 }
636 match self.iter.next_back() {
637 None => return and_then_or_clear(&mut self.frontiter, |f| f.next_back()),
638 Some(inner) => self.backiter = Some(inner.into_iter()),
639 }
640 }
641 }
642
643 #[inline]
644 default fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
645 where
646 Self: Sized,
647 Fold: FnMut(Acc, Self::Item) -> R,
648 R: Try<Output = Acc>,
649 {
650 #[inline]
651 fn flatten<U: DoubleEndedIterator, Acc, R: Try<Output = Acc>>(
652 mut fold: impl FnMut(Acc, U::Item) -> R,
653 ) -> impl FnMut(Acc, &mut U) -> R {
654 move |acc, iter| iter.try_rfold(acc, &mut fold)
655 }
656
657 self.iter_try_rfold(init, flatten(fold))
658 }
659
660 #[inline]
661 default fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
662 where
663 Fold: FnMut(Acc, Self::Item) -> Acc,
664 {
665 #[inline]
666 fn flatten<U: DoubleEndedIterator, Acc>(
667 mut fold: impl FnMut(Acc, U::Item) -> Acc,
668 ) -> impl FnMut(Acc, U) -> Acc {
669 move |acc, iter| iter.rfold(acc, &mut fold)
670 }
671
672 self.iter_rfold(init, flatten(fold))
673 }
674
675 #[inline]
676 #[rustc_inherit_overflow_checks]
677 default fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
678 #[inline]
679 #[rustc_inherit_overflow_checks]
680 fn advance<U: DoubleEndedIterator>(n: usize, iter: &mut U) -> ControlFlow<(), usize> {
681 match iter.advance_back_by(n) {
682 Ok(()) => ControlFlow::Break(()),
683 Err(remaining) => ControlFlow::Continue(remaining.get()),
684 }
685 }
686
687 match self.iter_try_rfold(n, advance) {
688 ControlFlow::Continue(remaining) => NonZero::new(remaining).map_or(Ok(()), Err),
689 _ => Ok(()),
690 }
691 }
692}
693
694unsafe impl<const N: usize, I, T> TrustedLen
695 for FlattenCompat<I, <[T; N] as IntoIterator>::IntoIter>
696where
697 I: TrustedLen<Item = [T; N]>,
698{
699}
700
701unsafe impl<'a, const N: usize, I, T> TrustedLen
702 for FlattenCompat<I, <&'a [T; N] as IntoIterator>::IntoIter>
703where
704 I: TrustedLen<Item = &'a [T; N]>,
705{
706}
707
708unsafe impl<'a, const N: usize, I, T> TrustedLen
709 for FlattenCompat<I, <&'a mut [T; N] as IntoIterator>::IntoIter>
710where
711 I: TrustedLen<Item = &'a mut [T; N]>,
712{
713}
714
715trait ConstSizeIntoIterator: IntoIterator {
716 // FIXME(#31844): convert to an associated const once specialization supports that
717 fn size() -> Option<usize>;
718}
719
720impl<T> ConstSizeIntoIterator for T
721where
722 T: IntoIterator,
723{
724 #[inline]
725 default fn size() -> Option<usize> {
726 None
727 }
728}
729
730impl<T, const N: usize> ConstSizeIntoIterator for [T; N] {
731 #[inline]
732 fn size() -> Option<usize> {
733 Some(N)
734 }
735}
736
737impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] {
738 #[inline]
739 fn size() -> Option<usize> {
740 Some(N)
741 }
742}
743
744impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] {
745 #[inline]
746 fn size() -> Option<usize> {
747 Some(N)
748 }
749}
750
751#[inline]
752fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
753 let x: Option = f(opt.as_mut()?);
754 if x.is_none() {
755 *opt = None;
756 }
757 x
758}
759
760/// Specialization trait for iterator types that never return more than one item.
761///
762/// Note that we still have to deal with the possibility that the iterator was
763/// already exhausted before it came into our control.
764#[rustc_specialization_trait]
765trait OneShot {}
766
767// These all have exactly one item, if not already consumed.
768impl<T> OneShot for Once<T> {}
769impl<F> OneShot for OnceWith<F> {}
770impl<T> OneShot for array::IntoIter<T, 1> {}
771impl<T> OneShot for option::IntoIter<T> {}
772impl<T> OneShot for option::Iter<'_, T> {}
773impl<T> OneShot for option::IterMut<'_, T> {}
774impl<T> OneShot for result::IntoIter<T> {}
775impl<T> OneShot for result::Iter<'_, T> {}
776impl<T> OneShot for result::IterMut<'_, T> {}
777
778// These are always empty, which is fine to optimize too.
779impl<T> OneShot for Empty<T> {}
780impl<T> OneShot for array::IntoIter<T, 0> {}
781
782// These adaptors never increase the number of items.
783// (There are more possible, but for now this matches BoundedSize above.)
784impl<I: OneShot> OneShot for Cloned<I> {}
785impl<I: OneShot> OneShot for Copied<I> {}
786impl<I: OneShot, P> OneShot for Filter<I, P> {}
787impl<I: OneShot, P> OneShot for FilterMap<I, P> {}
788impl<I: OneShot, F> OneShot for Map<I, F> {}
789
790// Blanket impls pass this property through as well
791// (but we can't do `Box<I>` unless we expose this trait to alloc)
792impl<I: OneShot> OneShot for &mut I {}
793
794#[inline]
795fn into_item<I>(inner: I) -> Option<I::Item>
796where
797 I: IntoIterator<IntoIter: OneShot>,
798{
799 inner.into_iter().next()
800}
801
802#[inline]
803fn flatten_one<I: IntoIterator<IntoIter: OneShot>, Acc>(
804 mut fold: impl FnMut(Acc, I::Item) -> Acc,
805) -> impl FnMut(Acc, I) -> Acc {
806 move |acc: Acc, inner: I| match inner.into_iter().next() {
807 Some(item: impl OneShot) => fold(acc, item),
808 None => acc,
809 }
810}
811
812#[inline]
813fn try_flatten_one<I: IntoIterator<IntoIter: OneShot>, Acc, R: Try<Output = Acc>>(
814 mut fold: impl FnMut(Acc, I::Item) -> R,
815) -> impl FnMut(Acc, I) -> R {
816 move |acc: Acc, inner: I| match inner.into_iter().next() {
817 Some(item: impl OneShot) => fold(acc, item),
818 None => try { acc },
819 }
820}
821
822#[inline]
823fn advance_by_one<I>(n: NonZero<usize>, inner: I) -> Option<NonZero<usize>>
824where
825 I: IntoIterator<IntoIter: OneShot>,
826{
827 match inner.into_iter().next() {
828 Some(_) => NonZero::new(n.get() - 1),
829 None => Some(n),
830 }
831}
832
833// Specialization: When the inner iterator `U` never returns more than one item, the `frontiter` and
834// `backiter` states are a waste, because they'll always have already consumed their item. So in
835// this impl, we completely ignore them and just focus on `self.iter`, and we only call the inner
836// `U::next()` one time.
837//
838// It's mostly fine if we accidentally mix this with the more generic impls, e.g. by forgetting to
839// specialize one of the methods. If the other impl did set the front or back, we wouldn't see it
840// here, but it would be empty anyway; and if the other impl looked for a front or back that we
841// didn't bother setting, it would just see `None` (or a previous empty) and move on.
842//
843// An exception to that is `advance_by(0)` and `advance_back_by(0)`, where the generic impls may set
844// `frontiter` or `backiter` without consuming the item, so we **must** override those.
845impl<I, U> Iterator for FlattenCompat<I, U>
846where
847 I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
848 U: Iterator + OneShot,
849{
850 #[inline]
851 fn next(&mut self) -> Option<U::Item> {
852 while let Some(inner) = self.iter.next() {
853 if let item @ Some(_) = inner.into_iter().next() {
854 return item;
855 }
856 }
857 None
858 }
859
860 #[inline]
861 fn size_hint(&self) -> (usize, Option<usize>) {
862 let (lower, upper) = self.iter.size_hint();
863 match <I::Item as ConstSizeIntoIterator>::size() {
864 Some(0) => (0, Some(0)),
865 Some(1) => (lower, upper),
866 _ => (0, upper),
867 }
868 }
869
870 #[inline]
871 fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
872 where
873 Self: Sized,
874 Fold: FnMut(Acc, Self::Item) -> R,
875 R: Try<Output = Acc>,
876 {
877 self.iter.try_fold(init, try_flatten_one(fold))
878 }
879
880 #[inline]
881 fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
882 where
883 Fold: FnMut(Acc, Self::Item) -> Acc,
884 {
885 self.iter.fold(init, flatten_one(fold))
886 }
887
888 #[inline]
889 fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
890 if let Some(n) = NonZero::new(n) {
891 self.iter.try_fold(n, advance_by_one).map_or(Ok(()), Err)
892 } else {
893 // Just advance the outer iterator
894 self.iter.advance_by(0)
895 }
896 }
897
898 #[inline]
899 fn count(self) -> usize {
900 self.iter.filter_map(into_item).count()
901 }
902
903 #[inline]
904 fn last(self) -> Option<Self::Item> {
905 self.iter.filter_map(into_item).last()
906 }
907}
908
909// Note: We don't actually care about `U: DoubleEndedIterator`, since forward and backward are the
910// same for a one-shot iterator, but we have to keep that to match the default specialization.
911impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
912where
913 I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
914 U: DoubleEndedIterator + OneShot,
915{
916 #[inline]
917 fn next_back(&mut self) -> Option<U::Item> {
918 while let Some(inner) = self.iter.next_back() {
919 if let item @ Some(_) = inner.into_iter().next() {
920 return item;
921 }
922 }
923 None
924 }
925
926 #[inline]
927 fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
928 where
929 Self: Sized,
930 Fold: FnMut(Acc, Self::Item) -> R,
931 R: Try<Output = Acc>,
932 {
933 self.iter.try_rfold(init, try_flatten_one(fold))
934 }
935
936 #[inline]
937 fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
938 where
939 Fold: FnMut(Acc, Self::Item) -> Acc,
940 {
941 self.iter.rfold(init, flatten_one(fold))
942 }
943
944 #[inline]
945 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
946 if let Some(n) = NonZero::new(n) {
947 self.iter.try_rfold(n, advance_by_one).map_or(Ok(()), Err)
948 } else {
949 // Just advance the outer iterator
950 self.iter.advance_back_by(0)
951 }
952 }
953}
954