{-# LANGUAGE Safe         #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Copyright:  (c) 2018-2021 Kowainik
SPDX-License-Identifier: MIT
Maintainer:  Kowainik <xrom.xkov@gmail.com>
Stability:   Experimental
Portability: Portable

Polymorphic grouping functions.

@since 0.1.0
-}

module Relude.Extra.Group
    ( groupBy
    , groupOneBy
    ) where

import Relude
import Relude.Extra.Map

import Data.List.NonEmpty ((<|))


{- | Groups elements using results of the given function as keys.

>>> groupBy even [1..6] :: HashMap Bool (NonEmpty Int)
fromList [(False,5 :| [3,1]),(True,6 :| [4,2])]

@since 0.1.0
-}
groupBy :: forall f t a . (Foldable f, DynamicMap t, Val t ~ NonEmpty a, Monoid t)
        => (a -> Key t) -> f a -> t
groupBy :: (a -> Key t) -> f a -> t
groupBy a -> Key t
f = (a -> t -> t) -> t -> f a -> t
forall (f :: * -> *) a b.
Foldable f =>
(a -> b -> b) -> b -> f a -> b
flipfoldl' a -> t -> t
hmGroup t
forall a. Monoid a => a
mempty
  where
    hmGroup :: a -> t -> t
    hmGroup :: a -> t -> t
hmGroup a
x =
        let val :: Maybe (NonEmpty a) -> NonEmpty a
            val :: Maybe (NonEmpty a) -> NonEmpty a
val Maybe (NonEmpty a)
Nothing   = OneItem (NonEmpty a) -> NonEmpty a
forall x. One x => OneItem x -> x
one a
OneItem (NonEmpty a)
x
            val (Just NonEmpty a
xs) = a
x a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty a
xs
        in (Maybe (Val t) -> Maybe (Val t)) -> Key t -> t -> t
forall t.
DynamicMap t =>
(Maybe (Val t) -> Maybe (Val t)) -> Key t -> t -> t
alter (NonEmpty a -> Maybe (NonEmpty a)
forall a. a -> Maybe a
Just (NonEmpty a -> Maybe (NonEmpty a))
-> (Maybe (NonEmpty a) -> NonEmpty a)
-> Maybe (NonEmpty a)
-> Maybe (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (NonEmpty a) -> NonEmpty a
val) (a -> Key t
f a
x)

{- | Similar to 'groupBy' but keeps only one element as value.

>>> groupOneBy even [1 .. 6] :: HashMap Bool Int
fromList [(False,1),(True,2)]

@since 0.1.0
-}
groupOneBy :: forall f t a . (Foldable f, DynamicMap t, Val t ~ a, Monoid t)
           => (a -> Key t) -> f a -> t
groupOneBy :: (a -> Key t) -> f a -> t
groupOneBy a -> Key t
f = (a -> t -> t) -> t -> f a -> t
forall (f :: * -> *) a b.
Foldable f =>
(a -> b -> b) -> b -> f a -> b
flipfoldl' a -> t -> t
hmGroup t
forall a. Monoid a => a
mempty
  where
    hmGroup :: a -> t -> t
    hmGroup :: a -> t -> t
hmGroup a
val t
m = let key :: Key t
key = a -> Key t
f a
val in
        case Key t -> t -> Maybe (Val t)
forall t. StaticMap t => Key t -> t -> Maybe (Val t)
lookup Key t
key t
m of
            Maybe (Val t)
Nothing -> Key t -> Val t -> t -> t
forall t. DynamicMap t => Key t -> Val t -> t -> t
insert Key t
key a
Val t
val t
m
            Just Val t
_  -> t
m