/*
Even BETTER grouping construct in Rust.
*/

pub struct GroupBy<A, B, I> {
  iter: I,
  test: Box<dyn Fn(&A) -> B>,
  curr: Option<B>,
  last: Option<A>
}

impl<A, B, I> GroupBy<A, B, I> {
  fn new<T: 'static>(iter: I, test: T) -> GroupBy<A, B, I> where T: Fn(&A) -> B {
    GroupBy { iter, test: Box::new(test), curr: None, last: None }
  }
} 

impl<A, B: Eq, I> Iterator for GroupBy<A, B, I>
  where I: Iterator<Item=A>
{
  type Item = Vec<A>;

  fn next(&mut self) -> Option<Self::Item> {
    let mut group = Vec::new();
    if let Some(a) = self.last.take() {
      group.push(a);
    }
    let test = &self.test;
    loop {
      match self.iter.next() {
        None => break,
        Some(a) => {
          let new = Some(test(&a));
          if self.curr == new || self.curr.is_none() {
            group.push(a);
            self.curr = new;
          } else {
            self.last = Some(a);
            self.curr = new;
            break;
          }
        }
      }
    }
    if !group.is_empty() { Some(group) } else { None }
  }
}

pub trait IteratorExt: Iterator {
  fn group_by<B: Eq, T: 'static>(self, test: T) -> GroupBy<Self::Item, B, Self>
    where Self: Sized, T: Fn(&Self::Item) -> B
  {
    GroupBy::new(self, test)
  }

  fn group(self) -> GroupBy<Self::Item, Self::Item, Self>
    where Self: Sized, Self::Item: Eq + Clone
  {
    GroupBy::new(self, |a: &Self::Item| a.clone())
  }
}