{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Serokell.Util.Group
( groupBy
, groupMapBy
) where
import Universum
import Data.List.NonEmpty ((<|))
import qualified Data.HashMap.Strict as HM
groupBy :: forall t a b . (Container t, Element t ~ a, Eq b, Hashable b)
=> (a -> b) -> t -> HashMap b (NonEmpty a)
groupBy f = flipfoldl' hmGroup mempty
where
hmGroup :: a -> HashMap b (NonEmpty a) -> HashMap b (NonEmpty a)
hmGroup x =
let val :: Maybe (NonEmpty a) -> NonEmpty a
val Nothing = one x
val (Just xs) = x <| xs
in HM.alter (Just . val) (f x)
groupMapBy :: forall t a b . (Container t, Element t ~ a, Eq b, Hashable b)
=> (a -> b) -> t -> HashMap b a
groupMapBy f = flipfoldl' hmGroup mempty
where
hmGroup :: a -> HashMap b a -> HashMap b a
hmGroup val hm = let key = f val in
case HM.lookup key hm of
Nothing -> HM.insert key val hm
Just _ -> hm