module Skeletest.Internal.Utils.Map (
  findOrEmpty,
  adjustNested,
) where

import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)
import GHC.Exts (IsList, fromList)

findOrEmpty :: (Ord k, IsList (t a)) => k -> Map k (t a) -> t a
findOrEmpty :: forall {k} k (t :: k -> *) (a :: k).
(Ord k, IsList (t a)) =>
k -> Map k (t a) -> t a
findOrEmpty = t a -> k -> Map k (t a) -> t a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault ([Item (t a)] -> t a
forall l. IsList l => [Item l] -> l
fromList [])

-- | Same as 'adjust', except defaulting to an empty structure if it doesn't
-- exist, and deleting the key if the adjusted value is empty.
adjustNested :: (Ord k, Foldable t, IsList (t a)) => (t a -> t a) -> k -> Map k (t a) -> Map k (t a)
adjustNested :: forall k (t :: * -> *) a.
(Ord k, Foldable t, IsList (t a)) =>
(t a -> t a) -> k -> Map k (t a) -> Map k (t a)
adjustNested t a -> t a
f =
  let prune :: t a -> Maybe (t a)
prune t a
m = if t a -> Bool
forall a. t a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null t a
m then Maybe (t a)
forall a. Maybe a
Nothing else t a -> Maybe (t a)
forall a. a -> Maybe a
Just t a
m
   in (Maybe (t a) -> Maybe (t a)) -> k -> Map k (t a) -> Map k (t a)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (t a -> Maybe (t a)
forall {t :: * -> *} {a}. Foldable t => t a -> Maybe (t a)
prune (t a -> Maybe (t a))
-> (Maybe (t a) -> t a) -> Maybe (t a) -> Maybe (t a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> t a
f (t a -> t a) -> (Maybe (t a) -> t a) -> Maybe (t a) -> t a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> Maybe (t a) -> t a
forall a. a -> Maybe a -> a
fromMaybe ([Item (t a)] -> t a
forall l. IsList l => [Item l] -> l
fromList []))