// A bit better than my previous attempts

mod group {
  use core::slice::Iter;
  use std::hash::Hash;
  use std::collections::HashMap;

  pub struct GroupByIter<'a, E: 'a, G: Eq + 'a> {
    iter: Iter<'a, E>,
    test: Box<dyn Fn(&'a E) -> G>,
    curr: Option<G>,
    last: Option<&'a E>
  } 

  impl<'a, E: 'a, G: Eq + 'a> GroupByIter<'a, E, G> {
    fn new<F: 'static>(slice: &'a [E], test: F) -> GroupByIter<'a, E, G> where F: Fn(&'a E) -> G {
      GroupByIter { iter: slice.iter(), test: Box::new(test), curr: None, last: None }
    }
  }

  impl<'a, E: 'a, G: Eq + 'a> Iterator for GroupByIter<'a, E, G> {
    type Item = Vec<&'a E>;

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

  pub trait Groupable {
    type Item;

    fn group_by<'a, F: 'static, G: Eq + 'a>(&'a self, test: F) -> GroupByIter<'a, Self::Item, G> where F: Fn(&'a Self::Item) -> G;

    fn group<'a>(&'a self) -> GroupByIter<'a, Self::Item, &'a Self::Item> where Self::Item: Eq + 'a {
      self.group_by(|e| e)
    }

    fn categorize_by<'a, F: 'static, G: Eq + Hash + 'a>(&'a self, test: F) -> HashMap<G, Vec<&'a Self::Item>> where F: Fn(&'a Self::Item) -> G;

    fn counts_by<'a, F: 'static, C: Eq + Hash + 'a>(&'a self, test: F) -> HashMap<C, usize> where F: Fn(&'a Self::Item) -> C {
      self.categorize_by(test).into_iter().map(|g| (g.0, g.1.len())).collect()
    }

    fn counts<'a>(&'a self) -> HashMap<&'a Self::Item, usize> where Self::Item: Eq + Hash {
      self.counts_by(|e| e)
    }
  }

  impl<E> Groupable for [E] {
    type Item = E;

    fn group_by<'a, F: 'static, G: Eq + 'a>(&'a self, test: F) -> GroupByIter<'a, E, G> where E: 'a, F: Fn(&'a E) -> G {
      GroupByIter::new(self, test)
    }

    fn categorize_by<'a, F: 'static, G: Eq + Hash + 'a>(&'a self, test: F) -> HashMap<G, Vec<&'a Self::Item>> where F: Fn(&'a Self::Item) -> G {
      let mut map: HashMap<G, Vec<&Self::Item>> = HashMap::new();
      for e in self {
        map.entry(test(e))
          .and_modify(|g| g.push(e))
          .or_insert(vec![e]);
      }
      map
    }

    fn counts_by<'a, F: 'static, C: Eq + Hash + 'a>(&'a self, test: F) -> HashMap<C, usize> where F: Fn(&'a Self::Item) -> C {
      let mut map = HashMap::new();
      for e in self {
        let c = test(&e);
        map.entry(c)
          .and_modify(|e| *e += 1)
          .or_insert(1);
      }
      map
    }
  }
}