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

{- |
Module                  : Relude.Extra.Group
Copyright               : (c) 2018-2023 Kowainik
SPDX-License-Identifier : MIT
Maintainer              : Kowainik <xrom.xkov@gmail.com>
Stability               : Stable
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 :: 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 = forall (f :: * -> *) a b.
Foldable f =>
(a -> b -> b) -> b -> f a -> b
flipfoldl' a -> t -> t
hmGroup 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   = forall x. One x => OneItem x -> x
one a
x
            val (Just NonEmpty a
xs) = a
x forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty a
xs
        in forall t.
DynamicMap t =>
(Maybe (Val t) -> Maybe (Val t)) -> Key t -> t -> t
alter (forall a. a -> Maybe a
Just 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 :: forall (f :: * -> *) t a.
(Foldable f, DynamicMap t, Val t ~ a, Monoid t) =>
(a -> Key t) -> f a -> t
groupOneBy a -> Key t
f = forall (f :: * -> *) a b.
Foldable f =>
(a -> b -> b) -> b -> f a -> b
flipfoldl' a -> t -> t
hmGroup 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 forall t. StaticMap t => Key t -> t -> Maybe (Val t)
lookup Key t
key t
m of
            Maybe (Val t)
Nothing -> forall t. DynamicMap t => Key t -> Val t -> t -> t
insert Key t
key a
val t
m
            Just Val t
_  -> t
m