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

   1 use crate::fmt;
   2 use crate::ops::Try;
   3
   4 use super::super::{DoubleEndedIterator, Fuse, FusedIterator, Iterator};
   5 use super::Map;
   6
   7 /// An iterator that maps each element to an iterator, and yields the elements
   8 /// of the produced iterators.
   9 ///
  10 /// This `struct` is created by [`Iterator::flat_map`]. See its documentation
  11 /// for more.
  12 #[must_use = "iterators are lazy and do nothing unless consumed"]
  13 #[stable(feature = "rust1", since = "1.0.0")]
  14 pub struct FlatMap<I, U: IntoIterator, F> {
  15     inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>,
  16 }
  17 impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> {
  18     pub(in super::super) fn new(iter: I, f: F) -> FlatMap<I, U, F> {
  19         FlatMap { inner: FlattenCompat::new(iter.map(f)) }
  20     }
  21 }
  22
  23 #[stable(feature = "rust1", since = "1.0.0")]
  24 impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F>
  25 where
  26     U: Clone + IntoIterator<IntoIter: Clone>,
  27 {
  28     fn clone(&self) -> Self {
  29         FlatMap { inner: self.inner.clone() }
  30     }
  31 }
  32
  33 #[stable(feature = "core_impl_debug", since = "1.9.0")]
  34 impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F>
  35 where
  36     U: IntoIterator<IntoIter: fmt::Debug>,
  37 {
  38     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  39         f.debug_struct("FlatMap").field("inner", &self.inner).finish()
  40     }
  41 }
  42
  43 #[stable(feature = "rust1", since = "1.0.0")]
  44 impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
  45 where
  46     F: FnMut(I::Item) -> U,
  47 {
  48     type Item = U::Item;
  49
  50     #[inline]
  51     fn next(&mut self) -> Option<U::Item> {
  52         self.inner.next()
  53     }
  54
  55     #[inline]
  56     fn size_hint(&self) -> (usize, Option<usize>) {
  57         self.inner.size_hint()
  58     }
  59
  60     #[inline]
  61     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  62     where
  63         Self: Sized,
  64         Fold: FnMut(Acc, Self::Item) -> R,
  65         R: Try<Ok = Acc>,
  66     {
  67         self.inner.try_fold(init, fold)
  68     }
  69
  70     #[inline]
  71     fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
  72     where
  73         Fold: FnMut(Acc, Self::Item) -> Acc,
  74     {
  75         self.inner.fold(init, fold)
  76     }
  77 }
  78
  79 #[stable(feature = "rust1", since = "1.0.0")]
  80 impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F>
  81 where
  82     F: FnMut(I::Item) -> U,
  83     U: IntoIterator<IntoIter: DoubleEndedIterator>,
  84 {
  85     #[inline]
  86     fn next_back(&mut self) -> Option<U::Item> {
  87         self.inner.next_back()
  88     }
  89
  90     #[inline]
  91     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  92     where
  93         Self: Sized,
  94         Fold: FnMut(Acc, Self::Item) -> R,
  95         R: Try<Ok = Acc>,
  96     {
  97         self.inner.try_rfold(init, fold)
  98     }
  99
 100     #[inline]
 101     fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 102     where
 103         Fold: FnMut(Acc, Self::Item) -> Acc,
 104     {
 105         self.inner.rfold(init, fold)
 106     }
 107 }
 108
 109 #[stable(feature = "fused", since = "1.26.0")]
 110 impl<I, U, F> FusedIterator for FlatMap<I, U, F>
 111 where
 112     I: FusedIterator,
 113     U: IntoIterator,
 114     F: FnMut(I::Item) -> U,
 115 {
 116 }
 117
 118 /// An iterator that flattens one level of nesting in an iterator of things
 119 /// that can be turned into iterators.
 120 ///
 121 /// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its
 122 /// documentation for more.
 123 ///
 124 /// [`flatten`]: Iterator::flatten
 125 /// [`Iterator`]: trait.Iterator.html
 126 #[must_use = "iterators are lazy and do nothing unless consumed"]
 127 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 128 pub struct Flatten<I: Iterator<Item: IntoIterator>> {
 129     inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
 130 }
 131
 132 impl<I: Iterator<Item: IntoIterator>> Flatten<I> {
 133     pub(in super::super) fn new(iter: I) -> Flatten<I> {
 134         Flatten { inner: FlattenCompat::new(iter) }
 135     }
 136 }
 137
 138 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 139 impl<I, U> fmt::Debug for Flatten<I>
 140 where
 141     I: fmt::Debug + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 142     U: fmt::Debug + Iterator,
 143 {
 144     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 145         f.debug_struct("Flatten").field("inner", &self.inner).finish()
 146     }
 147 }
 148
 149 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 150 impl<I, U> Clone for Flatten<I>
 151 where
 152     I: Clone + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 153     U: Clone + Iterator,
 154 {
 155     fn clone(&self) -> Self {
 156         Flatten { inner: self.inner.clone() }
 157     }
 158 }
 159
 160 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 161 impl<I, U> Iterator for Flatten<I>
 162 where
 163     I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 164     U: Iterator,
 165 {
 166     type Item = U::Item;
 167
 168     #[inline]
 169     fn next(&mut self) -> Option<U::Item> {
 170         self.inner.next()
 171     }
 172
 173     #[inline]
 174     fn size_hint(&self) -> (usize, Option<usize>) {
 175         self.inner.size_hint()
 176     }
 177
 178     #[inline]
 179     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 180     where
 181         Self: Sized,
 182         Fold: FnMut(Acc, Self::Item) -> R,
 183         R: Try<Ok = Acc>,
 184     {
 185         self.inner.try_fold(init, fold)
 186     }
 187
 188     #[inline]
 189     fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 190     where
 191         Fold: FnMut(Acc, Self::Item) -> Acc,
 192     {
 193         self.inner.fold(init, fold)
 194     }
 195 }
 196
 197 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 198 impl<I, U> DoubleEndedIterator for Flatten<I>
 199 where
 200     I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 201     U: DoubleEndedIterator,
 202 {
 203     #[inline]
 204     fn next_back(&mut self) -> Option<U::Item> {
 205         self.inner.next_back()
 206     }
 207
 208     #[inline]
 209     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 210     where
 211         Self: Sized,
 212         Fold: FnMut(Acc, Self::Item) -> R,
 213         R: Try<Ok = Acc>,
 214     {
 215         self.inner.try_rfold(init, fold)
 216     }
 217
 218     #[inline]
 219     fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 220     where
 221         Fold: FnMut(Acc, Self::Item) -> Acc,
 222     {
 223         self.inner.rfold(init, fold)
 224     }
 225 }
 226
 227 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 228 impl<I, U> FusedIterator for Flatten<I>
 229 where
 230     I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 231     U: Iterator,
 232 {
 233 }
 234
 235 /// Real logic of both `Flatten` and `FlatMap` which simply delegate to
 236 /// this type.
 237 #[derive(Clone, Debug)]
 238 struct FlattenCompat<I, U> {
 239     iter: Fuse<I>,
 240     frontiter: Option<U>,
 241     backiter: Option<U>,
 242 }
 243 impl<I, U> FlattenCompat<I, U>
 244 where
 245     I: Iterator,
 246 {
 247     /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
 248     fn new(iter: I) -> FlattenCompat<I, U> {
 249         FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None }
 250     }
 251 }
 252
 253 impl<I, U> Iterator for FlattenCompat<I, U>
 254 where
 255     I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 256     U: Iterator,
 257 {
 258     type Item = U::Item;
 259
 260     #[inline]
 261     fn next(&mut self) -> Option<U::Item> {
 262         loop {
 263             if let Some(ref mut inner) = self.frontiter {
 264                 match inner.next() {
 265                     None => self.frontiter = None,
 266                     elt @ Some(_) => return elt,
 267                 }
 268             }
 269             match self.iter.next() {
 270                 None => return self.backiter.as_mut()?.next(),
 271                 Some(inner) => self.frontiter = Some(inner.into_iter()),
 272             }
 273         }
 274     }
 275
 276     #[inline]
 277     fn size_hint(&self) -> (usize, Option<usize>) {
 278         let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
 279         let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
 280         let lo = flo.saturating_add(blo);
 281         match (self.iter.size_hint(), fhi, bhi) {
 282             ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
 283             _ => (lo, None),
 284         }
 285     }
 286
 287     #[inline]
 288     fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
 289     where
 290         Self: Sized,
 291         Fold: FnMut(Acc, Self::Item) -> R,
 292         R: Try<Ok = Acc>,
 293     {
 294         #[inline]
 295         fn flatten<'a, T: IntoIterator, Acc, R: Try<Ok = Acc>>(
 296             frontiter: &'a mut Option<T::IntoIter>,
 297             fold: &'a mut impl FnMut(Acc, T::Item) -> R,
 298         ) -> impl FnMut(Acc, T) -> R + 'a {
 299             move |acc, x| {
 300                 let mut mid = x.into_iter();
 301                 let r = mid.try_fold(acc, &mut *fold);
 302                 *frontiter = Some(mid);
 303                 r
 304             }
 305         }
 306
 307         if let Some(ref mut front) = self.frontiter {
 308             init = front.try_fold(init, &mut fold)?;
 309         }
 310         self.frontiter = None;
 311
 312         init = self.iter.try_fold(init, flatten(&mut self.frontiter, &mut fold))?;
 313         self.frontiter = None;
 314
 315         if let Some(ref mut back) = self.backiter {
 316             init = back.try_fold(init, &mut fold)?;
 317         }
 318         self.backiter = None;
 319
 320         try { init }
 321     }
 322
 323     #[inline]
 324     fn fold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
 325     where
 326         Fold: FnMut(Acc, Self::Item) -> Acc,
 327     {
 328         #[inline]
 329         fn flatten<U: Iterator, Acc>(
 330             fold: &mut impl FnMut(Acc, U::Item) -> Acc,
 331         ) -> impl FnMut(Acc, U) -> Acc + '_ {
 332             move |acc, iter| iter.fold(acc, &mut *fold)
 333         }
 334
 335         self.frontiter
 336             .into_iter()
 337             .chain(self.iter.map(IntoIterator::into_iter))
 338             .chain(self.backiter)
 339             .fold(init, flatten(fold))
 340     }
 341 }
 342
 343 impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
 344 where
 345     I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
 346     U: DoubleEndedIterator,
 347 {
 348     #[inline]
 349     fn next_back(&mut self) -> Option<U::Item> {
 350         loop {
 351             if let Some(ref mut inner) = self.backiter {
 352                 match inner.next_back() {
 353                     None => self.backiter = None,
 354                     elt @ Some(_) => return elt,
 355                 }
 356             }
 357             match self.iter.next_back() {
 358                 None => return self.frontiter.as_mut()?.next_back(),
 359                 next => self.backiter = next.map(IntoIterator::into_iter),
 360             }
 361         }
 362     }
 363
 364     #[inline]
 365     fn try_rfold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
 366     where
 367         Self: Sized,
 368         Fold: FnMut(Acc, Self::Item) -> R,
 369         R: Try<Ok = Acc>,
 370     {
 371         #[inline]
 372         fn flatten<'a, T: IntoIterator, Acc, R: Try<Ok = Acc>>(
 373             backiter: &'a mut Option<T::IntoIter>,
 374             fold: &'a mut impl FnMut(Acc, T::Item) -> R,
 375         ) -> impl FnMut(Acc, T) -> R + 'a
 376         where
 377             T::IntoIter: DoubleEndedIterator,
 378         {
 379             move |acc, x| {
 380                 let mut mid = x.into_iter();
 381                 let r = mid.try_rfold(acc, &mut *fold);
 382                 *backiter = Some(mid);
 383                 r
 384             }
 385         }
 386
 387         if let Some(ref mut back) = self.backiter {
 388             init = back.try_rfold(init, &mut fold)?;
 389         }
 390         self.backiter = None;
 391
 392         init = self.iter.try_rfold(init, flatten(&mut self.backiter, &mut fold))?;
 393         self.backiter = None;
 394
 395         if let Some(ref mut front) = self.frontiter {
 396             init = front.try_rfold(init, &mut fold)?;
 397         }
 398         self.frontiter = None;
 399
 400         try { init }
 401     }
 402
 403     #[inline]
 404     fn rfold<Acc, Fold>(self, init: Acc, ref mut fold: Fold) -> Acc
 405     where
 406         Fold: FnMut(Acc, Self::Item) -> Acc,
 407     {
 408         #[inline]
 409         fn flatten<U: DoubleEndedIterator, Acc>(
 410             fold: &mut impl FnMut(Acc, U::Item) -> Acc,
 411         ) -> impl FnMut(Acc, U) -> Acc + '_ {
 412             move |acc, iter| iter.rfold(acc, &mut *fold)
 413         }
 414
 415         self.frontiter
 416             .into_iter()
 417             .chain(self.iter.map(IntoIterator::into_iter))
 418             .chain(self.backiter)
 419             .rfold(init, flatten(fold))
 420     }
 421 }