1  use std::{hash::Hash, iter::Iterator};
  2  
  3  use rustc_hash::FxHashMap;
  4  
  5  pub trait IterExt: Iterator {
  6      fn take_while_inclusive<P>(self, predicate: P) -> TakeWhileInclusive<Self, P>
  7      where
  8          Self: Sized,
  9          P: FnMut(&Self::Item) -> bool;
 10  
 11      fn chain_if_empty<U>(self, other: U) -> ChainIfEmpty<Self, U::IntoIter>
 12      where
 13          Self: Sized,
 14          U: IntoIterator<Item = Self::Item>;
 15  
 16      fn transpose(self) -> Transpose<<Self::Item as IntoIterator>::IntoIter>
 17      where
 18          Self: Sized,
 19          Self::Item: IntoIterator;
 20  
 21      fn all_consume<P>(self, predicate: P) -> bool
 22      where
 23          Self: Sized,
 24          P: FnMut(Self::Item) -> bool;
 25  
 26      fn any_consume<P>(self, predicate: P) -> bool
 27      where
 28          Self: Sized,
 29          P: FnMut(Self::Item) -> bool;
 30  
 31      fn index(mut self, elem: impl PartialEq<Self::Item>) -> Option<usize>
 32      where
 33          Self: Sized,
 34      {
 35          self.position(|x| elem == x)
 36      }
 37  
 38      fn counts_fx(self) -> FxHashMap<Self::Item, usize>
 39      where
 40          Self: Sized,
 41          Self::Item: Hash + Eq,
 42      {
 43          self.fold(FxHashMap::default(), |mut counts, x| {
 44              counts.entry(x).and_modify(|x| *x += 1).or_insert(1);
 45              counts
 46          })
 47      }
 48  }
 49  
 50  impl<I> IterExt for I
 51  where
 52      I: Iterator,
 53  {
 54      fn take_while_inclusive<P>(self, predicate: P) -> TakeWhileInclusive<I, P>
 55      where
 56          P: FnMut(&I::Item) -> bool,
 57      {
 58          TakeWhileInclusive::new(self, predicate)
 59      }
 60  
 61      fn chain_if_empty<U>(self, other: U) -> ChainIfEmpty<I, U::IntoIter>
 62      where
 63          U: IntoIterator<Item = Self::Item>,
 64      {
 65          ChainIfEmpty {
 66              iter: self,
 67              other: other.into_iter(),
 68              state: State::Unknown,
 69          }
 70      }
 71  
 72      fn transpose(self) -> Transpose<<Self::Item as std::iter::IntoIterator>::IntoIter>
 73      where
 74          Self::Item: IntoIterator,
 75      {
 76          Transpose {
 77              iterators: self.map(|it| it.into_iter()).collect(),
 78          }
 79      }
 80  
 81      fn all_consume<P>(self, mut predicate: P) -> bool
 82      where
 83          Self: Sized,
 84          P: FnMut(Self::Item) -> bool,
 85      {
 86          #[allow(clippy::unnecessary_fold)]
 87          self.fold(true, |acc, x| acc && predicate(x))
 88      }
 89  
 90      fn any_consume<P>(self, mut predicate: P) -> bool
 91      where
 92          Self: Sized,
 93          P: FnMut(Self::Item) -> bool,
 94      {
 95          #[allow(clippy::unnecessary_fold)]
 96          self.fold(false, |acc, x| acc || predicate(x))
 97      }
 98  }
 99  
100  pub struct TakeWhileInclusive<I, P> {
101      iter: I,
102      predicate: P,
103      done: bool,
104  }
105  
106  impl<I, P> TakeWhileInclusive<I, P> {
107      fn new(iter: I, predicate: P) -> Self {
108          TakeWhileInclusive {
109              iter,
110              predicate,
111              done: false,
112          }
113      }
114  }
115  
116  impl<I, P> Iterator for TakeWhileInclusive<I, P>
117  where
118      I: Iterator,
119      P: FnMut(&I::Item) -> bool,
120  {
121      type Item = I::Item;
122  
123      fn next(&mut self) -> Option<Self::Item> {
124          if self.done {
125              return None;
126          }
127          self.iter.next().inspect(|item| {
128              if !(self.predicate)(item) {
129                  self.done = true;
130              }
131          })
132      }
133  }
134  
135  enum State {
136      Empty,
137      NotEmpty,
138      Unknown,
139  }
140  
141  pub struct ChainIfEmpty<I, U> {
142      iter: I,
143      other: U,
144      state: State,
145  }
146  
147  impl<I, U> Iterator for ChainIfEmpty<I, U>
148  where
149      I: Iterator,
150      U: Iterator<Item = I::Item>,
151  {
152      type Item = I::Item;
153  
154      fn next(&mut self) -> Option<Self::Item> {
155          match self.state {
156              State::Empty => self.other.next(),
157              State::NotEmpty => self.iter.next(),
158              State::Unknown => match self.iter.next() {
159                  Some(x) => {
160                      self.state = State::NotEmpty;
161                      Some(x)
162                  }
163                  None => {
164                      self.state = State::Empty;
165                      self.other.next()
166                  }
167              },
168          }
169      }
170  }
171  
172  pub struct Transpose<I: Iterator> {
173      iterators: Vec<I>,
174  }
175  
176  impl<I: Iterator> Iterator for Transpose<I> {
177      type Item = Vec<I::Item>;
178  
179      fn next(&mut self) -> Option<Self::Item> {
180          self.iterators
181              .iter_mut()
182              .map(|it| it.next())
183              .collect::<Option<Vec<_>>>()
184      }
185  }
186  
187  #[cfg(test)]
188  mod tests_take_while_inclusive {
189      use super::*;
190  
191      macro_rules! test {
192          ($name:ident, $inp:expr, $pred:expr, $exp: expr) => {
193              #[test]
194              fn $name() {
195                  let pred = $pred;
196                  assert_eq!(
197                      $inp.iter()
198                          .take_while_inclusive(|&&x| (pred)(x))
199                          .copied()
200                          .collect::<Vec<_>>(),
201                      $exp
202                  )
203              }
204          };
205      }
206  
207      test!(split, [1, 2, 3, 4], |x| x < 3, vec![1, 2, 3]);
208      test!(always_true, [1, 2, 3, 4], |_| true, vec![1, 2, 3, 4]);
209      test!(always_false, [1, 2, 3, 4], |_| false, vec![1]);
210  }
211  
212  #[cfg(test)]
213  mod tests_chain_if_empty {
214      use super::*;
215  
216      macro_rules! test {
217          ($name:ident, $inp1:expr, $inp2:expr, $exp: expr) => {
218              #[test]
219              fn $name() {
220                  assert_eq!(
221                      $inp1.into_iter().chain_if_empty($inp2).collect::<Vec<_>>(),
222                      $exp
223                  );
224              }
225          };
226      }
227  
228      test!(both_empty, Vec::<i32>::new(), [], vec![]);
229      test!(first_empty, [], [4, 5, 6], vec![4, 5, 6]);
230      test!(second_empty, [1, 2, 3], [], vec![1, 2, 3]);
231      test!(both_nonempty, [1, 2, 3], [4, 5, 6], vec![1, 2, 3]);
232  }