{-# LANGUAGE CPP #-}

module Agda.Utils.Map where

import Control.Monad ((<$!>))
import Data.Functor.Compose
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Map.Internal (Map(..), balanceL, balanceR, singleton)

import Agda.Utils.Impossible

-- * Monadic map operations
---------------------------------------------------------------------------

{-# INLINE insertWithGood #-}
-- | Version of `insertWith` that's willing to be properly inlined.
insertWithGood :: forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithGood :: forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithGood a -> a -> a
f k
k !a
a = k -> a -> Map k a -> Map k a
go k
k a
a where
  go :: k -> a -> Map k a -> Map k a
  go :: k -> a -> Map k a -> Map k a
go !k
kx a
x Map k a
Tip               = k -> a -> Map k a
forall k a. k -> a -> Map k a
singleton k
kx a
x
  go  k
kx a
x (Bin Size
sy k
ky a
y Map k a
l Map k a
r) = case k -> k -> Ordering
forall a. Ord a => a -> a -> Ordering
compare k
kx k
ky of
    Ordering
LT -> k -> a -> Map k a -> Map k a -> Map k a
forall k a. k -> a -> Map k a -> Map k a -> Map k a
balanceL k
ky a
y (k -> a -> Map k a -> Map k a
go k
kx a
x Map k a
l) Map k a
r
    Ordering
GT -> k -> a -> Map k a -> Map k a -> Map k a
forall k a. k -> a -> Map k a -> Map k a -> Map k a
balanceR k
ky a
y Map k a
l (k -> a -> Map k a -> Map k a
go k
kx a
x Map k a
r)
    Ordering
EQ -> let !y' :: a
y' = a -> a -> a
f a
x a
y in Size -> k -> a -> Map k a -> Map k a -> Map k a
forall k a. Size -> k -> a -> Map k a -> Map k a -> Map k a
Bin Size
sy k
kx a
y' Map k a
l Map k a
r

{-# INLINE forGood_ #-}
-- | Version of `forM_` that deigns to be properly lambda-lifted for State, Reader, etc.
forGood_ :: forall k v m. Applicative m => Map k v -> (v -> m ()) -> m ()
forGood_ :: forall k v (m :: * -> *).
Applicative m =>
Map k v -> (v -> m ()) -> m ()
forGood_ Map k v
m v -> m ()
f = Map k v -> m ()
go Map k v
m where
  go :: Map k v -> m ()
go Map k v
Tip             = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  go (Bin Size
_ k
k v
v Map k v
l Map k v
r) = Map k v -> m ()
go Map k v
l m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> v -> m ()
f v
v m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Map k v -> m ()
go Map k v
r

{-# INLINE forWithKey_ #-}
forWithKey_ :: forall k v m. Applicative m => Map k v -> (k -> v -> m ()) -> m ()
forWithKey_ :: forall k v (m :: * -> *).
Applicative m =>
Map k v -> (k -> v -> m ()) -> m ()
forWithKey_ Map k v
m k -> v -> m ()
f = Map k v -> m ()
go Map k v
m where
  go :: Map k v -> m ()
go Map k v
Tip             = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  go (Bin Size
_ k
k v
v Map k v
l Map k v
r) = Map k v -> m ()
go Map k v
l m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> k -> v -> m ()
f k
k v
v m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Map k v -> m ()
go Map k v
r

{-# INLINE adjustM #-}
-- | Update monadically the value at one position (must exist!).
adjustM :: (Functor f, Ord k) => (v -> f v) -> k -> Map k v -> f (Map k v)
adjustM :: forall (f :: * -> *) k v.
(Functor f, Ord k) =>
(v -> f v) -> k -> Map k v -> f (Map k v)
adjustM v -> f v
f = (Maybe v -> f (Maybe v)) -> k -> Map k v -> f (Map k v)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF ((Maybe v -> f (Maybe v)) -> k -> Map k v -> f (Map k v))
-> (Maybe v -> f (Maybe v)) -> k -> Map k v -> f (Map k v)
forall a b. (a -> b) -> a -> b
$ \case
  Maybe v
Nothing -> f (Maybe v)
forall a. HasCallStack => a
__IMPOSSIBLE__
  Just v
v  -> v -> Maybe v
forall a. a -> Maybe a
Just (v -> Maybe v) -> f v -> f (Maybe v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v -> f v
f v
v

{-# INLINE adjustM' #-}
-- | Wrapper for 'adjustM' for convenience.
adjustM' :: (Functor f, Ord k) => (v -> f (a, v)) -> k -> Map k v -> f (a, Map k v)
adjustM' :: forall (f :: * -> *) k v a.
(Functor f, Ord k) =>
(v -> f (a, v)) -> k -> Map k v -> f (a, Map k v)
adjustM' v -> f (a, v)
f k
k = Compose f ((,) a) (Map k v) -> f (a, Map k v)
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (Compose f ((,) a) (Map k v) -> f (a, Map k v))
-> (Map k v -> Compose f ((,) a) (Map k v))
-> Map k v
-> f (a, Map k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v -> Compose f ((,) a) v)
-> k -> Map k v -> Compose f ((,) a) (Map k v)
forall (f :: * -> *) k v.
(Functor f, Ord k) =>
(v -> f v) -> k -> Map k v -> f (Map k v)
adjustM (f (a, v) -> Compose f ((,) a) v
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (f (a, v) -> Compose f ((,) a) v)
-> (v -> f (a, v)) -> v -> Compose f ((,) a) v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> f (a, v)
f) k
k

-- * Non-monadic map operations
---------------------------------------------------------------------------

#if !MIN_VERSION_containers(0,8,0)
{-# INLINE filterKeys #-}
-- | Filter a map based on the keys.
filterKeys :: (k -> Bool) -> Map k a -> Map k a
filterKeys :: forall k a. (k -> Bool) -> Map k a -> Map k a
filterKeys k -> Bool
p = (k -> a -> Bool) -> Map k a -> Map k a
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (Bool -> a -> Bool
forall a b. a -> b -> a
const (Bool -> a -> Bool) -> (k -> Bool) -> k -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Bool
p)
#endif

-- | Check whether a map is a singleton.
isSingleMap :: Map k v -> Maybe (k, v)
isSingleMap :: forall k v. Map k v -> Maybe (k, v)
isSingleMap (Bin Size
1 k
k v
v Map k v
_ Map k v
_) = (k, v) -> Maybe (k, v)
forall a. a -> Maybe a
Just (k
k, v
v)
isSingleMap Map k v
_               = Maybe (k, v)
forall a. Maybe a
Nothing