{-# LANGUAGE Trustworthy  #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Copyright:  (c) 2018-2021 Kowainik
SPDX-License-Identifier: MIT
Maintainer:  Kowainik <xrom.xkov@gmail.com>
Stability:   Experimental
Portability: Portable

Contains implementation of polymorphic type classes for data types 'Set' and
'Map'.
-}

module Relude.Extra.Map
    ( StaticMap (..)
    , DynamicMap (..)
    , (!?)
    , notMember
    , lookupDefault

      -- * To pairs
    , toPairs
    , keys
    , elems
    ) where

import GHC.Exts (IsList (Item, toList))

import Relude.Base (Eq, Ord, Type)
import Relude.Bool (Bool, guard, not)
import Relude.Container.Reexport (HashMap, HashSet, Hashable, IntMap, IntSet, Map, Set, fst, snd)
import Relude.Function ((.))
import Relude.Functor.Reexport (($>))
import Relude.List (map)
import Relude.Monad.Reexport (Maybe (..), fromMaybe)
import Relude.Numeric (Int)

import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Map.Strict as M
import qualified Data.Set as S

----------------------------------------------------------------------------
-- Static Map
----------------------------------------------------------------------------

{- | Read-only map or set. Contains polymorphic functions which work for both
sets and maps.

@since 0.1.0
-}
class StaticMap t where
    type Key t :: Type
    type Val t :: Type

    size   :: t -> Int
    lookup :: Key t -> t -> Maybe (Val t)
    member :: Key t -> t -> Bool

{- |

@since 0.1.0
-}
instance Ord k => StaticMap (Map k v) where
    type Key (Map k v) = k
    type Val (Map k v) = v

    size :: Map k v -> Int
size = Map k v -> Int
forall k a. Map k a -> Int
M.size
    {-# INLINE size #-}
    lookup :: Key (Map k v) -> Map k v -> Maybe (Val (Map k v))
lookup = Key (Map k v) -> Map k v -> Maybe (Val (Map k v))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup
    {-# INLINE lookup #-}
    member :: Key (Map k v) -> Map k v -> Bool
member = Key (Map k v) -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member
    {-# INLINE member #-}

{- |

@since 0.1.0
-}
instance (Eq k, Hashable k) => StaticMap (HashMap k v) where
    type Key (HashMap k v) = k
    type Val (HashMap k v) = v

    size :: HashMap k v -> Int
size = HashMap k v -> Int
forall k v. HashMap k v -> Int
HM.size
    {-# INLINE size #-}
    lookup :: Key (HashMap k v) -> HashMap k v -> Maybe (Val (HashMap k v))
lookup = Key (HashMap k v) -> HashMap k v -> Maybe (Val (HashMap k v))
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup
    {-# INLINE lookup #-}
    member :: Key (HashMap k v) -> HashMap k v -> Bool
member = Key (HashMap k v) -> HashMap k v -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
HM.member
    {-# INLINE member #-}

{- |

@since 0.1.0
-}
instance StaticMap (IntMap v) where
    type Key (IntMap v) = Int
    type Val (IntMap v) = v

    size :: IntMap v -> Int
size = IntMap v -> Int
forall v. IntMap v -> Int
IM.size
    {-# INLINE size #-}
    lookup :: Key (IntMap v) -> IntMap v -> Maybe (Val (IntMap v))
lookup = Key (IntMap v) -> IntMap v -> Maybe (Val (IntMap v))
forall a. Int -> IntMap a -> Maybe a
IM.lookup
    {-# INLINE lookup #-}
    member :: Key (IntMap v) -> IntMap v -> Bool
member = Key (IntMap v) -> IntMap v -> Bool
forall a. Int -> IntMap a -> Bool
IM.member
    {-# INLINE member #-}

{- |

@since 0.1.0
-}
instance Ord a => StaticMap (Set a) where
    type Key (Set a) = a
    type Val (Set a) = a

    size :: Set a -> Int
size = Set a -> Int
forall a. Set a -> Int
S.size
    {-# INLINE size #-}
    member :: Key (Set a) -> Set a -> Bool
member = Key (Set a) -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member
    {-# INLINE member #-}
    lookup :: Key (Set a) -> Set a -> Maybe (Val (Set a))
lookup Key (Set a)
k Set a
m = Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Key (Set a) -> Set a -> Bool
forall t. StaticMap t => Key t -> t -> Bool
member Key (Set a)
k Set a
m) Maybe () -> a -> Maybe a
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> a
Key (Set a)
k
    {-# INLINE lookup #-}

{- |

@since 0.1.0
-}
instance (Eq a, Hashable a) => StaticMap (HashSet a) where
    type Key (HashSet a) = a
    type Val (HashSet a) = a

    size :: HashSet a -> Int
size = HashSet a -> Int
forall a. HashSet a -> Int
HS.size
    {-# INLINE size #-}
    member :: Key (HashSet a) -> HashSet a -> Bool
member = Key (HashSet a) -> HashSet a -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HS.member
    {-# INLINE member #-}
    lookup :: Key (HashSet a) -> HashSet a -> Maybe (Val (HashSet a))
lookup Key (HashSet a)
k HashSet a
m = Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Key (HashSet a) -> HashSet a -> Bool
forall t. StaticMap t => Key t -> t -> Bool
member Key (HashSet a)
k HashSet a
m) Maybe () -> a -> Maybe a
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> a
Key (HashSet a)
k
    {-# INLINE lookup #-}

{- |

@since 0.1.0
-}
instance StaticMap IntSet where
    type Key IntSet = Int
    type Val IntSet = Int

    size :: IntSet -> Int
size = IntSet -> Int
IS.size
    {-# INLINE size #-}
    member :: Key IntSet -> IntSet -> Bool
member = Int -> IntSet -> Bool
Key IntSet -> IntSet -> Bool
IS.member
    {-# INLINE member #-}
    lookup :: Key IntSet -> IntSet -> Maybe (Val IntSet)
lookup Key IntSet
k IntSet
m = Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Key IntSet -> IntSet -> Bool
forall t. StaticMap t => Key t -> t -> Bool
member Key IntSet
k IntSet
m) Maybe () -> Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Int
Key IntSet
k
    {-# INLINE lookup #-}

{- | Operator version of 'lookup' function.

>>> let myHashMap = HashMap.fromList [('a', "xxx"), ('b', "yyy")]
>>> myHashMap !? 'b'
Just "yyy"

>>> myHashMap !? 'd'
Nothing

@since 0.1.0
-}
infixl 9 !?
(!?) :: StaticMap t => t -> Key t -> Maybe (Val t)
!? :: forall t. StaticMap t => t -> Key t -> Maybe (Val t)
(!?) t
m Key t
k = Key t -> t -> Maybe (Val t)
forall t. StaticMap t => Key t -> t -> Maybe (Val t)
lookup Key t
k t
m
{-# INLINE (!?) #-}

{- | Inverse of 'member' function.

>>> let myHashMap = HashMap.fromList [('a', "xxx"), ('b', "yyy")]
>>> notMember 'b' myHashMap
False

>>> notMember 'c' myHashMap
True

@since 0.1.0
-}
notMember :: StaticMap t => Key t -> t -> Bool
notMember :: forall t. StaticMap t => Key t -> t -> Bool
notMember Key t
k = Bool -> Bool
not (Bool -> Bool) -> (t -> Bool) -> t -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key t -> t -> Bool
forall t. StaticMap t => Key t -> t -> Bool
member Key t
k
{-# INLINE notMember #-}

{- | Return the value to which the specified key is mapped, or the default value
if this map contains no mapping for the key.

>>> let myHashMap = HashMap.fromList [('a', "xxx"), ('b', "yyy")]
>>> lookupDefault "zzz" 'b' myHashMap
"yyy"

>>> lookupDefault "zzz" 'c' myHashMap
"zzz"

@since 0.1.0
-}
lookupDefault :: StaticMap t
              => Val t -- ^ Default value to return.
              -> Key t -- ^ Key to search
              -> t     -- ^ Container to search
              -> Val t
lookupDefault :: forall t. StaticMap t => Val t -> Key t -> t -> Val t
lookupDefault Val t
def Key t
k = Val t -> Maybe (Val t) -> Val t
forall a. a -> Maybe a -> a
fromMaybe Val t
def (Maybe (Val t) -> Val t) -> (t -> Maybe (Val t)) -> t -> Val t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key t -> t -> Maybe (Val t)
forall t. StaticMap t => Key t -> t -> Maybe (Val t)
lookup Key t
k
{-# INLINE lookupDefault #-}

----------------------------------------------------------------------------
-- Dynamic Map
----------------------------------------------------------------------------

{- | Modifiable Map.

@since 0.1.0
-}
class StaticMap t => DynamicMap t where
    -- insertions
    insert     :: Key t -> Val t -> t -> t
    insertWith :: (Val t -> Val t -> Val t) -> Key t -> Val t -> t -> t

    -- deletions
    delete :: Key t -> t -> t
    alter :: (Maybe (Val t) -> Maybe (Val t)) -> Key t -> t -> t

{- |

@since 0.1.0
-}
instance Ord k => DynamicMap (Map k v) where
    insert :: Key (Map k v) -> Val (Map k v) -> Map k v -> Map k v
insert     = Key (Map k v) -> Val (Map k v) -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
    {-# INLINE insert #-}
    insertWith :: (Val (Map k v) -> Val (Map k v) -> Val (Map k v))
-> Key (Map k v) -> Val (Map k v) -> Map k v -> Map k v
insertWith = (Val (Map k v) -> Val (Map k v) -> Val (Map k v))
-> Key (Map k v) -> Val (Map k v) -> Map k v -> Map k v
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith
    {-# INLINE insertWith #-}
    delete :: Key (Map k v) -> Map k v -> Map k v
delete     = Key (Map k v) -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
M.delete
    {-# INLINE delete #-}
    alter :: (Maybe (Val (Map k v)) -> Maybe (Val (Map k v)))
-> Key (Map k v) -> Map k v -> Map k v
alter      = (Maybe (Val (Map k v)) -> Maybe (Val (Map k v)))
-> Key (Map k v) -> Map k v -> Map k v
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter
    {-# INLINE alter #-}

{- |

@since 0.1.0
-}
instance (Eq k, Hashable k) => DynamicMap (HashMap k v) where
    insert :: Key (HashMap k v)
-> Val (HashMap k v) -> HashMap k v -> HashMap k v
insert     = Key (HashMap k v)
-> Val (HashMap k v) -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert
    {-# INLINE insert #-}
    insertWith :: (Val (HashMap k v) -> Val (HashMap k v) -> Val (HashMap k v))
-> Key (HashMap k v)
-> Val (HashMap k v)
-> HashMap k v
-> HashMap k v
insertWith = (Val (HashMap k v) -> Val (HashMap k v) -> Val (HashMap k v))
-> Key (HashMap k v)
-> Val (HashMap k v)
-> HashMap k v
-> HashMap k v
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith
    {-# INLINE insertWith #-}
    delete :: Key (HashMap k v) -> HashMap k v -> HashMap k v
delete     = Key (HashMap k v) -> HashMap k v -> HashMap k v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HM.delete
    {-# INLINE delete #-}
    alter :: (Maybe (Val (HashMap k v)) -> Maybe (Val (HashMap k v)))
-> Key (HashMap k v) -> HashMap k v -> HashMap k v
alter      = (Maybe (Val (HashMap k v)) -> Maybe (Val (HashMap k v)))
-> Key (HashMap k v) -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
(Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
HM.alter
    {-# INLINE alter #-}

{- |

@since 0.1.0
-}
instance DynamicMap (IntMap v) where
    insert :: Key (IntMap v) -> Val (IntMap v) -> IntMap v -> IntMap v
insert     = Key (IntMap v) -> Val (IntMap v) -> IntMap v -> IntMap v
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert
    {-# INLINE insert #-}
    insertWith :: (Val (IntMap v) -> Val (IntMap v) -> Val (IntMap v))
-> Key (IntMap v) -> Val (IntMap v) -> IntMap v -> IntMap v
insertWith = (Val (IntMap v) -> Val (IntMap v) -> Val (IntMap v))
-> Key (IntMap v) -> Val (IntMap v) -> IntMap v -> IntMap v
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IM.insertWith
    {-# INLINE insertWith #-}
    delete :: Key (IntMap v) -> IntMap v -> IntMap v
delete     = Key (IntMap v) -> IntMap v -> IntMap v
forall a. Int -> IntMap a -> IntMap a
IM.delete
    {-# INLINE delete #-}
    alter :: (Maybe (Val (IntMap v)) -> Maybe (Val (IntMap v)))
-> Key (IntMap v) -> IntMap v -> IntMap v
alter      = (Maybe (Val (IntMap v)) -> Maybe (Val (IntMap v)))
-> Key (IntMap v) -> IntMap v -> IntMap v
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter
    {-# INLINE alter #-}

----------------------------------------------------------------------------
-- ToPairs
----------------------------------------------------------------------------

-- $setup
-- >>> import qualified Data.HashMap.Strict as HashMap

{- | Converts the structure to the list of the key-value pairs.

>>> toPairs (HashMap.fromList [('a', "xxx"), ('b', "yyy")])
[('a',"xxx"),('b',"yyy")]

@since 0.1.0
-}
toPairs :: (IsList t, Item t ~ (a, b)) => t -> [(a, b)]
toPairs :: forall t a b. (IsList t, Item t ~ (a, b)) => t -> [(a, b)]
toPairs = t -> [(a, b)]
forall l. IsList l => l -> [Item l]
toList

{- | Converts the structure to the list of the keys.

>>> keys (HashMap.fromList [('a', "xxx"), ('b', "yyy")])
"ab"

@since 0.1.0
-}
keys :: (IsList t, Item t ~ (a, b)) => t -> [a]
keys :: forall t a b. (IsList t, Item t ~ (a, b)) => t -> [a]
keys = ((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst ([(a, b)] -> [a]) -> (t -> [(a, b)]) -> t -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> [(a, b)]
forall l. IsList l => l -> [Item l]
toList

{- | Converts the structure to the list of the values.

>>> elems (HashMap.fromList [('a', "xxx"), ('b', "yyy")])
["xxx","yyy"]

@since 0.1.0
-}
elems :: (IsList t, Item t ~ (a, b)) => t -> [b]
elems :: forall t a b. (IsList t, Item t ~ (a, b)) => t -> [b]
elems = ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd ([(a, b)] -> [b]) -> (t -> [(a, b)]) -> t -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> [(a, b)]
forall l. IsList l => l -> [Item l]
toList