vendor/rustc-rayon/src/iter/flat_map.rs

   1 use super::plumbing::*;
   2 use super::*;
   3
   4 use std::fmt::{self, Debug};
   5
   6 /// `FlatMap` maps each element to an iterator, then flattens these iterators together.
   7 /// This struct is created by the [`flat_map()`] method on [`ParallelIterator`]
   8 ///
   9 /// [`flat_map()`]: trait.ParallelIterator.html#method.flat_map
  10 /// [`ParallelIterator`]: trait.ParallelIterator.html
  11 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
  12 #[derive(Clone)]
  13 pub struct FlatMap<I: ParallelIterator, F> {
  14     base: I,
  15     map_op: F,
  16 }
  17
  18 impl<I: ParallelIterator + Debug, F> Debug for FlatMap<I, F> {
  19     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  20         f.debug_struct("FlatMap").field("base", &self.base).finish()
  21     }
  22 }
  23
  24 impl<I: ParallelIterator, F> FlatMap<I, F> {
  25     /// Create a new `FlatMap` iterator.
  26     pub(super) fn new(base: I, map_op: F) -> Self {
  27         FlatMap { base, map_op }
  28     }
  29 }
  30
  31 impl<I, F, PI> ParallelIterator for FlatMap<I, F>
  32 where
  33     I: ParallelIterator,
  34     F: Fn(I::Item) -> PI + Sync + Send,
  35     PI: IntoParallelIterator,
  36 {
  37     type Item = PI::Item;
  38
  39     fn drive_unindexed<C>(self, consumer: C) -> C::Result
  40     where
  41         C: UnindexedConsumer<Self::Item>,
  42     {
  43         let consumer = FlatMapConsumer {
  44             base: consumer,
  45             map_op: &self.map_op,
  46         };
  47         self.base.drive_unindexed(consumer)
  48     }
  49 }
  50
  51 /// ////////////////////////////////////////////////////////////////////////
  52 /// Consumer implementation
  53
  54 struct FlatMapConsumer<'f, C, F> {
  55     base: C,
  56     map_op: &'f F,
  57 }
  58
  59 impl<'f, C, F> FlatMapConsumer<'f, C, F> {
  60     fn new(base: C, map_op: &'f F) -> Self {
  61         FlatMapConsumer { base, map_op }
  62     }
  63 }
  64
  65 impl<'f, T, U, C, F> Consumer<T> for FlatMapConsumer<'f, C, F>
  66 where
  67     C: UnindexedConsumer<U::Item>,
  68     F: Fn(T) -> U + Sync,
  69     U: IntoParallelIterator,
  70 {
  71     type Folder = FlatMapFolder<'f, C, F, C::Result>;
  72     type Reducer = C::Reducer;
  73     type Result = C::Result;
  74
  75     fn split_at(self, index: usize) -> (Self, Self, C::Reducer) {
  76         let (left, right, reducer) = self.base.split_at(index);
  77         (
  78             FlatMapConsumer::new(left, self.map_op),
  79             FlatMapConsumer::new(right, self.map_op),
  80             reducer,
  81         )
  82     }
  83
  84     fn into_folder(self) -> Self::Folder {
  85         FlatMapFolder {
  86             base: self.base,
  87             map_op: self.map_op,
  88             previous: None,
  89         }
  90     }
  91
  92     fn full(&self) -> bool {
  93         self.base.full()
  94     }
  95 }
  96
  97 impl<'f, T, U, C, F> UnindexedConsumer<T> for FlatMapConsumer<'f, C, F>
  98 where
  99     C: UnindexedConsumer<U::Item>,
 100     F: Fn(T) -> U + Sync,
 101     U: IntoParallelIterator,
 102 {
 103     fn split_off_left(&self) -> Self {
 104         FlatMapConsumer::new(self.base.split_off_left(), self.map_op)
 105     }
 106
 107     fn to_reducer(&self) -> Self::Reducer {
 108         self.base.to_reducer()
 109     }
 110 }
 111
 112 struct FlatMapFolder<'f, C, F, R> {
 113     base: C,
 114     map_op: &'f F,
 115     previous: Option<R>,
 116 }
 117
 118 impl<'f, T, U, C, F> Folder<T> for FlatMapFolder<'f, C, F, C::Result>
 119 where
 120     C: UnindexedConsumer<U::Item>,
 121     F: Fn(T) -> U + Sync,
 122     U: IntoParallelIterator,
 123 {
 124     type Result = C::Result;
 125
 126     fn consume(self, item: T) -> Self {
 127         let map_op = self.map_op;
 128         let par_iter = map_op(item).into_par_iter();
 129         let result = par_iter.drive_unindexed(self.base.split_off_left());
 130
 131         // We expect that `previous` is `None`, because we drive
 132         // the cost up so high, but just in case.
 133         let previous = match self.previous {
 134             None => Some(result),
 135             Some(previous) => {
 136                 let reducer = self.base.to_reducer();
 137                 Some(reducer.reduce(previous, result))
 138             }
 139         };
 140
 141         FlatMapFolder {
 142             base: self.base,
 143             map_op,
 144             previous,
 145         }
 146     }
 147
 148     fn complete(self) -> Self::Result {
 149         match self.previous {
 150             Some(previous) => previous,
 151             None => self.base.into_folder().complete(),
 152         }
 153     }
 154
 155     fn full(&self) -> bool {
 156         self.base.full()
 157     }
 158 }