graphql-engine/server/src-lib/Data/HashMap/Strict/InsOrd/Extended.hs

module Data.HashMap.Strict.InsOrd.Extended
  ( module OMap,
    groupTuples,
    groupListWith,
    partition,
  )
where

import Data.HashMap.Strict.InsOrd as OMap
import Data.Hashable (Hashable)
import Data.List qualified as L
import Data.Sequence.NonEmpty qualified as NE
import Prelude

groupTuples ::
  (Eq k, Hashable k, Foldable t) =>
  t (k, v) ->
  OMap.InsOrdHashMap k (NE.NESeq v)
groupTuples =
  L.foldl' groupFlds OMap.empty
  where
    groupFlds m (k, v) =
      OMap.insertWith (flip (<>)) k (NE.singleton v) m

groupListWith ::
  (Eq k, Hashable k, Foldable t, Functor t) =>
  (v -> k) ->
  t v ->
  OMap.InsOrdHashMap k (NE.NESeq v)
groupListWith f l =
  groupTuples $ fmap (\v -> (f v, v)) l

partition :: (Eq k, Hashable k) => (v -> Bool) -> OMap.InsOrdHashMap k v -> (OMap.InsOrdHashMap k v, OMap.InsOrdHashMap k v)
partition predicate =
  OMap.foldlWithKey'
    ( \(left, right) key val ->
        if (predicate val)
          then (OMap.insert key val left, right)
          else (left, OMap.insert key val right)
    )
    (mempty, mempty)