library/core/src/iter/adapters/flatten.rs

   1 use crate::fmt;
   2 use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen};
   3 use crate::ops::Try;
   4
   5 /// An iterator that maps each element to an iterator, and yields the elements
   6 /// of the produced iterators.
   7 ///
   8 /// This `struct` is created by [`Iterator::flat_map`]. See its documentation
   9 /// for more.
  10 #[must_use = "iterators are lazy and do nothing unless consumed"]
  11 #[stable(feature = "rust1", since = "1.0.0")]
  12 pub struct FlatMap<I, U: IntoIterator, F> {
  13     inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>,
  14 }
  15
  16 impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> {
  17     pub(in crate::iter) fn new(iter: I, f: F) -> FlatMap<I, U, F> {
  18         FlatMap { inner: FlattenCompat::new(iter.map(f)) }
  19     }
  20 }
  21
  22 #[stable(feature = "rust1", since = "1.0.0")]
  23 impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F>
  24 where
  25     U: Clone + IntoIterator<IntoIter: Clone>,
  26 {
  27     fn clone(&self) -> Self {
  28         FlatMap { inner: self.inner.clone() }
  29     }
  30 }
  31
  32 #[stable(feature = "core_impl_debug", since = "1.9.0")]
  33 impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F>
  34 where
  35     U: IntoIterator<IntoIter: fmt::Debug>,
  36 {
  37     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  38         f.debug_struct("FlatMap").field("inner", &self.inner).finish()
  39     }
  40 }
  41
  42 #[stable(feature = "rust1", since = "1.0.0")]
  43 impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
  44 where
  45     F: FnMut(I::Item) -> U,
  46 {
  47     type Item = U::Item;
  48
  49     #[inline]
  50     fn next(&mut self) -> Option<U::Item> {
  51         self.inner.next()
  52     }
  53
  54     #[inline]
  55     fn size_hint(&self) -> (usize, Option<usize>) {
  56         self.inner.size_hint()
  57     }
  58
  59     #[inline]
  60     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  61     where
  62         Self: Sized,
  63         Fold: FnMut(Acc, Self::Item) -> R,
  64         R: Try<Output = Acc>,
  65     {
  66         self.inner.try_fold(init, fold)
  67     }
  68
  69     #[inline]
  70     fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
  71     where
  72         Fold: FnMut(Acc, Self::Item) -> Acc,
  73     {
  74         self.inner.fold(init, fold)
  75     }
  76 }
  77
  78 #[stable(feature = "rust1", since = "1.0.0")]
  79 impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F>
  80 where
  81     F: FnMut(I::Item) -> U,
  82     U: IntoIterator<IntoIter: DoubleEndedIterator>,
  83 {
  84     #[inline]
  85     fn next_back(&mut self) -> Option<U::Item> {
  86         self.inner.next_back()
  87     }
  88
  89     #[inline]
  90     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  91     where
  92         Self: Sized,
  93         Fold: FnMut(Acc, Self::Item) -> R,
  94         R: Try<Output = Acc>,
  95     {
  96         self.inner.try_rfold(init, fold)
  97     }
  98
  99     #[inline]
 100     fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 101     where
 102         Fold: FnMut(Acc, Self::Item) -> Acc,
 103     {
 104         self.inner.rfold(init, fold)
 105     }
 106 }
 107
 108 #[stable(feature = "fused", since = "1.26.0")]
 109 impl<I, U, F> FusedIterator for FlatMap<I, U, F>
 110 where
 111     I: FusedIterator,
 112     U: IntoIterator,
 113     F: FnMut(I::Item) -> U,
 114 {
 115 }
 116
 117 #[unstable(feature = "trusted_len", issue = "37572")]
 118 unsafe impl<T, I, F, const N: usize> TrustedLen for FlatMap<I, [T; N], F>
 119 where
 120     I: TrustedLen,
 121     F: FnMut(I::Item) -> [T; N],
 122 {
 123 }
 124
 125 #[unstable(feature = "trusted_len", issue = "37572")]
 126 unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a [T; N], F>
 127 where
 128     I: TrustedLen,
 129     F: FnMut(I::Item) -> &'a [T; N],
 130 {
 131 }
 132
 133 #[unstable(feature = "trusted_len", issue = "37572")]
 134 unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a mut [T; N], F>
 135 where
 136     I: TrustedLen,
 137     F: FnMut(I::Item) -> &'a mut [T; N],
 138 {
 139 }
 140
 141 /// An iterator that flattens one level of nesting in an iterator of things
 142 /// that can be turned into iterators.
 143 ///
 144 /// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its
 145 /// documentation for more.
 146 ///
 147 /// [`flatten`]: Iterator::flatten()
 148 #[must_use = "iterators are lazy and do nothing unless consumed"]
 149 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 150 pub struct Flatten<I: Iterator<Item: IntoIterator>> {
 151     inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
 152 }
 153
 154 impl<I: Iterator<Item: IntoIterator>> Flatten<I> {
 155     pub(in super::super) fn new(iter: I) -> Flatten<I> {
 156         Flatten { inner: FlattenCompat::new(iter) }
 157     }
 158 }
 159
 160 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 161 impl<I, U> fmt::Debug for Flatten<I>
 162 where
 163     I: fmt::Debug + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 164     U: fmt::Debug + Iterator,
 165 {
 166     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 167         f.debug_struct("Flatten").field("inner", &self.inner).finish()
 168     }
 169 }
 170
 171 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 172 impl<I, U> Clone for Flatten<I>
 173 where
 174     I: Clone + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 175     U: Clone + Iterator,
 176 {
 177     fn clone(&self) -> Self {
 178         Flatten { inner: self.inner.clone() }
 179     }
 180 }
 181
 182 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 183 impl<I, U> Iterator for Flatten<I>
 184 where
 185     I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 186     U: Iterator,
 187 {
 188     type Item = U::Item;
 189
 190     #[inline]
 191     fn next(&mut self) -> Option<U::Item> {
 192         self.inner.next()
 193     }
 194
 195     #[inline]
 196     fn size_hint(&self) -> (usize, Option<usize>) {
 197         self.inner.size_hint()
 198     }
 199
 200     #[inline]
 201     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 202     where
 203         Self: Sized,
 204         Fold: FnMut(Acc, Self::Item) -> R,
 205         R: Try<Output = Acc>,
 206     {
 207         self.inner.try_fold(init, fold)
 208     }
 209
 210     #[inline]
 211     fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 212     where
 213         Fold: FnMut(Acc, Self::Item) -> Acc,
 214     {
 215         self.inner.fold(init, fold)
 216     }
 217 }
 218
 219 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 220 impl<I, U> DoubleEndedIterator for Flatten<I>
 221 where
 222     I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 223     U: DoubleEndedIterator,
 224 {
 225     #[inline]
 226     fn next_back(&mut self) -> Option<U::Item> {
 227         self.inner.next_back()
 228     }
 229
 230     #[inline]
 231     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 232     where
 233         Self: Sized,
 234         Fold: FnMut(Acc, Self::Item) -> R,
 235         R: Try<Output = Acc>,
 236     {
 237         self.inner.try_rfold(init, fold)
 238     }
 239
 240     #[inline]
 241     fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 242     where
 243         Fold: FnMut(Acc, Self::Item) -> Acc,
 244     {
 245         self.inner.rfold(init, fold)
 246     }
 247 }
 248
 249 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 250 impl<I, U> FusedIterator for Flatten<I>
 251 where
 252     I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 253     U: Iterator,
 254 {
 255 }
 256
 257 #[unstable(feature = "trusted_len", issue = "37572")]
 258 unsafe impl<I> TrustedLen for Flatten<I>
 259 where
 260     I: TrustedLen,
 261     <I as Iterator>::Item: TrustedConstSize,
 262 {
 263 }
 264
 265 /// Real logic of both `Flatten` and `FlatMap` which simply delegate to
 266 /// this type.
 267 #[derive(Clone, Debug)]
 268 struct FlattenCompat<I, U> {
 269     iter: Fuse<I>,
 270     frontiter: Option<U>,
 271     backiter: Option<U>,
 272 }
 273 impl<I, U> FlattenCompat<I, U>
 274 where
 275     I: Iterator,
 276 {
 277     /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
 278     fn new(iter: I) -> FlattenCompat<I, U> {
 279         FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None }
 280     }
 281 }
 282
 283 impl<I, U> Iterator for FlattenCompat<I, U>
 284 where
 285     I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 286     U: Iterator,
 287 {
 288     type Item = U::Item;
 289
 290     #[inline]
 291     fn next(&mut self) -> Option<U::Item> {
 292         loop {
 293             if let Some(ref mut inner) = self.frontiter {
 294                 match inner.next() {
 295                     None => self.frontiter = None,
 296                     elt @ Some(_) => return elt,
 297                 }
 298             }
 299             match self.iter.next() {
 300                 None => match self.backiter.as_mut()?.next() {
 301                     None => {
 302                         self.backiter = None;
 303                         return None;
 304                     }
 305                     elt @ Some(_) => return elt,
 306                 },
 307                 Some(inner) => self.frontiter = Some(inner.into_iter()),
 308             }
 309         }
 310     }
 311
 312     #[inline]
 313     fn size_hint(&self) -> (usize, Option<usize>) {
 314         let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
 315         let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
 316         let lo = flo.saturating_add(blo);
 317
 318         if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIntoIterator>::size() {
 319             let (lower, upper) = self.iter.size_hint();
 320
 321             let lower = lower.saturating_mul(fixed_size).saturating_add(lo);
 322             let upper =
 323                 try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? };
 324
 325             return (lower, upper);
 326         }
 327
 328         match (self.iter.size_hint(), fhi, bhi) {
 329             ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
 330             _ => (lo, None),
 331         }
 332     }
 333
 334     #[inline]
 335     fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
 336     where
 337         Self: Sized,
 338         Fold: FnMut(Acc, Self::Item) -> R,
 339         R: Try<Output = Acc>,
 340     {
 341         #[inline]
 342         fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>(
 343             frontiter: &'a mut Option<T::IntoIter>,
 344             fold: &'a mut impl FnMut(Acc, T::Item) -> R,
 345         ) -> impl FnMut(Acc, T) -> R + 'a {
 346             move |acc, x| {
 347                 let mut mid = x.into_iter();
 348                 let r = mid.try_fold(acc, &mut *fold);
 349                 *frontiter = Some(mid);
 350                 r
 351             }
 352         }
 353
 354         if let Some(ref mut front) = self.frontiter {
 355             init = front.try_fold(init, &mut fold)?;
 356         }
 357         self.frontiter = None;
 358
 359         init = self.iter.try_fold(init, flatten(&mut self.frontiter, &mut fold))?;
 360         self.frontiter = None;
 361
 362         if let Some(ref mut back) = self.backiter {
 363             init = back.try_fold(init, &mut fold)?;
 364         }
 365         self.backiter = None;
 366
 367         try { init }
 368     }
 369
 370     #[inline]
 371     fn fold<Acc, Fold>(self, mut init: Acc, mut fold: Fold) -> Acc
 372     where
 373         Fold: FnMut(Acc, Self::Item) -> Acc,
 374     {
 375         #[inline]
 376         fn flatten<T: IntoIterator, Acc>(
 377             fold: &mut impl FnMut(Acc, T::Item) -> Acc,
 378         ) -> impl FnMut(Acc, T) -> Acc + '_ {
 379             move |acc, x| x.into_iter().fold(acc, &mut *fold)
 380         }
 381
 382         if let Some(front) = self.frontiter {
 383             init = front.fold(init, &mut fold);
 384         }
 385
 386         init = self.iter.fold(init, flatten(&mut fold));
 387
 388         if let Some(back) = self.backiter {
 389             init = back.fold(init, &mut fold);
 390         }
 391
 392         init
 393     }
 394 }
 395
 396 impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
 397 where
 398     I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 399     U: DoubleEndedIterator,
 400 {
 401     #[inline]
 402     fn next_back(&mut self) -> Option<U::Item> {
 403         loop {
 404             if let Some(ref mut inner) = self.backiter {
 405                 match inner.next_back() {
 406                     None => self.backiter = None,
 407                     elt @ Some(_) => return elt,
 408                 }
 409             }
 410             match self.iter.next_back() {
 411                 None => match self.frontiter.as_mut()?.next_back() {
 412                     None => {
 413                         self.frontiter = None;
 414                         return None;
 415                     }
 416                     elt @ Some(_) => return elt,
 417                 },
 418                 next => self.backiter = next.map(IntoIterator::into_iter),
 419             }
 420         }
 421     }
 422
 423     #[inline]
 424     fn try_rfold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
 425     where
 426         Self: Sized,
 427         Fold: FnMut(Acc, Self::Item) -> R,
 428         R: Try<Output = Acc>,
 429     {
 430         #[inline]
 431         fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>(
 432             backiter: &'a mut Option<T::IntoIter>,
 433             fold: &'a mut impl FnMut(Acc, T::Item) -> R,
 434         ) -> impl FnMut(Acc, T) -> R + 'a
 435         where
 436             T::IntoIter: DoubleEndedIterator,
 437         {
 438             move |acc, x| {
 439                 let mut mid = x.into_iter();
 440                 let r = mid.try_rfold(acc, &mut *fold);
 441                 *backiter = Some(mid);
 442                 r
 443             }
 444         }
 445
 446         if let Some(ref mut back) = self.backiter {
 447             init = back.try_rfold(init, &mut fold)?;
 448         }
 449         self.backiter = None;
 450
 451         init = self.iter.try_rfold(init, flatten(&mut self.backiter, &mut fold))?;
 452         self.backiter = None;
 453
 454         if let Some(ref mut front) = self.frontiter {
 455             init = front.try_rfold(init, &mut fold)?;
 456         }
 457         self.frontiter = None;
 458
 459         try { init }
 460     }
 461
 462     #[inline]
 463     fn rfold<Acc, Fold>(self, mut init: Acc, mut fold: Fold) -> Acc
 464     where
 465         Fold: FnMut(Acc, Self::Item) -> Acc,
 466     {
 467         #[inline]
 468         fn flatten<T: IntoIterator, Acc>(
 469             fold: &mut impl FnMut(Acc, T::Item) -> Acc,
 470         ) -> impl FnMut(Acc, T) -> Acc + '_
 471         where
 472             T::IntoIter: DoubleEndedIterator,
 473         {
 474             move |acc, x| x.into_iter().rfold(acc, &mut *fold)
 475         }
 476
 477         if let Some(back) = self.backiter {
 478             init = back.rfold(init, &mut fold);
 479         }
 480
 481         init = self.iter.rfold(init, flatten(&mut fold));
 482
 483         if let Some(front) = self.frontiter {
 484             init = front.rfold(init, &mut fold);
 485         }
 486
 487         init
 488     }
 489 }
 490
 491 trait ConstSizeIntoIterator: IntoIterator {
 492     // FIXME(#31844): convert to an associated const once specialization supports that
 493     fn size() -> Option<usize>;
 494 }
 495
 496 impl<T> ConstSizeIntoIterator for T
 497 where
 498     T: IntoIterator,
 499 {
 500     #[inline]
 501     default fn size() -> Option<usize> {
 502         None
 503     }
 504 }
 505
 506 impl<T, const N: usize> ConstSizeIntoIterator for [T; N] {
 507     #[inline]
 508     fn size() -> Option<usize> {
 509         Some(N)
 510     }
 511 }
 512
 513 impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] {
 514     #[inline]
 515     fn size() -> Option<usize> {
 516         Some(N)
 517     }
 518 }
 519
 520 impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] {
 521     #[inline]
 522     fn size() -> Option<usize> {
 523         Some(N)
 524     }
 525 }
 526
 527 #[doc(hidden)]
 528 #[unstable(feature = "std_internals", issue = "none")]
 529 // FIXME(#20400): Instead of this helper trait there should be multiple impl TrustedLen for Flatten<>
 530 //   blocks with different bounds on Iterator::Item but the compiler erroneously considers them overlapping
 531 pub unsafe trait TrustedConstSize: IntoIterator {}
 532
 533 #[unstable(feature = "std_internals", issue = "none")]
 534 unsafe impl<T, const N: usize> TrustedConstSize for [T; N] {}
 535 #[unstable(feature = "std_internals", issue = "none")]
 536 unsafe impl<T, const N: usize> TrustedConstSize for &'_ [T; N] {}
 537 #[unstable(feature = "std_internals", issue = "none")]
 538 unsafe impl<T, const N: usize> TrustedConstSize for &'_ mut [T; N] {}