module Data.HashMap.Strict.InsOrd.Extended
  ( module OMap,
    catMaybes,
    partition,
    alterF,
  )
where

import Data.HashMap.Strict.InsOrd as OMap
import Data.Hashable (Hashable)
import Prelude

catMaybes :: InsOrdHashMap k (Maybe v) -> InsOrdHashMap k v
catMaybes :: InsOrdHashMap k (Maybe v) -> InsOrdHashMap k v
catMaybes = (Maybe v -> Maybe v)
-> InsOrdHashMap k (Maybe v) -> InsOrdHashMap k v
forall v1 v2 k.
(v1 -> Maybe v2) -> InsOrdHashMap k v1 -> InsOrdHashMap k v2
OMap.mapMaybe Maybe v -> Maybe v
forall a. a -> a
id

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

-- | Alter a hashmap using a function that can fail, in which case the entire operation fails.
-- (Maybe a version with the key also being passed to the function could be useful.)
alterF ::
  (Functor f, Eq k, Hashable k) =>
  (Maybe v -> f (Maybe v)) ->
  k ->
  InsOrdHashMap k v ->
  f (InsOrdHashMap k v)
alterF :: (Maybe v -> f (Maybe v))
-> k -> InsOrdHashMap k v -> f (InsOrdHashMap k v)
alterF Maybe v -> f (Maybe v)
f k
k InsOrdHashMap k v
m = Maybe v -> InsOrdHashMap k v
alter' (Maybe v -> InsOrdHashMap k v)
-> f (Maybe v) -> f (InsOrdHashMap k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe v -> f (Maybe v)
f (k -> InsOrdHashMap k v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> InsOrdHashMap k v -> Maybe v
OMap.lookup k
k InsOrdHashMap k v
m)
  where
    alter' :: Maybe v -> InsOrdHashMap k v
alter' = \case
      Maybe v
Nothing -> k -> InsOrdHashMap k v -> InsOrdHashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> InsOrdHashMap k v -> InsOrdHashMap k v
OMap.delete k
k InsOrdHashMap k v
m
      Just v
v -> k -> v -> InsOrdHashMap k v -> InsOrdHashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> InsOrdHashMap k v -> InsOrdHashMap k v
OMap.insert k
k v
v InsOrdHashMap k v
m