-- | HashMap to multiple values.
module Data.HashMap.Strict.Multi
  ( -- * Type
    MultiMap,

    -- * Construction and conversions
    singleton,
    fromMap,
    toMap,
    fromList,
    toList,

    -- * Basic interface
    lookup,
    insert,
    keys,
    elems,
  )
where

import Data.Aeson (ToJSON)
import Data.HashMap.Strict qualified as M
import Data.Hashable (Hashable)
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Prelude hiding (lookup)

-------------------------------------------------------------------------------

-- | Map from keys to sets of values
newtype MultiMap k v = MultiMap
  { MultiMap k v -> HashMap k (Set v)
unMultiMap :: M.HashMap k (S.Set v)
  }
  deriving newtype (MultiMap k v -> MultiMap k v -> Bool
(MultiMap k v -> MultiMap k v -> Bool)
-> (MultiMap k v -> MultiMap k v -> Bool) -> Eq (MultiMap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => MultiMap k v -> MultiMap k v -> Bool
/= :: MultiMap k v -> MultiMap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => MultiMap k v -> MultiMap k v -> Bool
== :: MultiMap k v -> MultiMap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => MultiMap k v -> MultiMap k v -> Bool
Eq, Int -> MultiMap k v -> ShowS
[MultiMap k v] -> ShowS
MultiMap k v -> String
(Int -> MultiMap k v -> ShowS)
-> (MultiMap k v -> String)
-> ([MultiMap k v] -> ShowS)
-> Show (MultiMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> MultiMap k v -> ShowS
forall k v. (Show k, Show v) => [MultiMap k v] -> ShowS
forall k v. (Show k, Show v) => MultiMap k v -> String
showList :: [MultiMap k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [MultiMap k v] -> ShowS
show :: MultiMap k v -> String
$cshow :: forall k v. (Show k, Show v) => MultiMap k v -> String
showsPrec :: Int -> MultiMap k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> MultiMap k v -> ShowS
Show, [MultiMap k v] -> Value
[MultiMap k v] -> Encoding
MultiMap k v -> Value
MultiMap k v -> Encoding
(MultiMap k v -> Value)
-> (MultiMap k v -> Encoding)
-> ([MultiMap k v] -> Value)
-> ([MultiMap k v] -> Encoding)
-> ToJSON (MultiMap k v)
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
forall k v. (ToJSON v, ToJSONKey k) => [MultiMap k v] -> Value
forall k v. (ToJSON v, ToJSONKey k) => [MultiMap k v] -> Encoding
forall k v. (ToJSON v, ToJSONKey k) => MultiMap k v -> Value
forall k v. (ToJSON v, ToJSONKey k) => MultiMap k v -> Encoding
toEncodingList :: [MultiMap k v] -> Encoding
$ctoEncodingList :: forall k v. (ToJSON v, ToJSONKey k) => [MultiMap k v] -> Encoding
toJSONList :: [MultiMap k v] -> Value
$ctoJSONList :: forall k v. (ToJSON v, ToJSONKey k) => [MultiMap k v] -> Value
toEncoding :: MultiMap k v -> Encoding
$ctoEncoding :: forall k v. (ToJSON v, ToJSONKey k) => MultiMap k v -> Encoding
toJSON :: MultiMap k v -> Value
$ctoJSON :: forall k v. (ToJSON v, ToJSONKey k) => MultiMap k v -> Value
ToJSON)

instance (Eq k, Hashable k, Ord v) => Semigroup (MultiMap k v) where
  MultiMap HashMap k (Set v)
m0 <> :: MultiMap k v -> MultiMap k v -> MultiMap k v
<> MultiMap HashMap k (Set v)
m1 = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap (HashMap k (Set v) -> MultiMap k v)
-> HashMap k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Set v -> Set v)
-> HashMap k (Set v) -> HashMap k (Set v) -> HashMap k (Set v)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
M.unionWith Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
S.union HashMap k (Set v)
m0 HashMap k (Set v)
m1

instance (Eq k, Hashable k, Ord v) => Monoid (MultiMap k v) where
  mempty :: MultiMap k v
mempty = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap HashMap k (Set v)
forall a. Monoid a => a
mempty

-------------------------------------------------------------------------------

-- | Construct a 'MmultiMap' with a single key, to which only one
-- value is associated.
singleton :: Hashable k => k -> v -> MultiMap k v
singleton :: k -> v -> MultiMap k v
singleton k
k v
v = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap (HashMap k (Set v) -> MultiMap k v)
-> HashMap k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ k -> Set v -> HashMap k (Set v)
forall k v. Hashable k => k -> v -> HashMap k v
M.singleton k
k (v -> Set v
forall a. a -> Set a
S.singleton v
v)

-- | Construct a 'MultiMap' with the supplied mappings.
fromMap :: M.HashMap k (S.Set v) -> MultiMap k v
fromMap :: HashMap k (Set v) -> MultiMap k v
fromMap = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap

-- | Convert a 'MultiMap' to a 'HashMap'.
toMap :: MultiMap k v -> M.HashMap k (S.Set v)
toMap :: MultiMap k v -> HashMap k (Set v)
toMap = MultiMap k v -> HashMap k (Set v)
forall k v. MultiMap k v -> HashMap k (Set v)
unMultiMap

-- | Creates a 'MultiMap' from an association list.
--
-- If the provided list constains duplicate mappings, the resulting
-- 'MultiMap' will store the set of all mapped values for each
-- duplicate key.
fromList :: (Eq k, Hashable k, Ord v) => [(k, v)] -> MultiMap k v
fromList :: [(k, v)] -> MultiMap k v
fromList [(k, v)]
l = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap (HashMap k (Set v) -> MultiMap k v)
-> HashMap k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Set v -> Set v) -> [(k, Set v)] -> HashMap k (Set v)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
M.fromListWith (Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
S.union) ([(k, Set v)] -> HashMap k (Set v))
-> [(k, Set v)] -> HashMap k (Set v)
forall a b. (a -> b) -> a -> b
$ ((k, v) -> (k, Set v)) -> [(k, v)] -> [(k, Set v)]
forall a b. (a -> b) -> [a] -> [b]
map ((v -> Set v) -> (k, v) -> (k, Set v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap v -> Set v
forall a. a -> Set a
S.singleton) [(k, v)]
l

-- | Creates an association list from a 'MultiMap'.
--
-- Each set of values associated with a given key is transformed back
-- into a list.
toList :: MultiMap k v -> [(k, [v])]
toList :: MultiMap k v -> [(k, [v])]
toList (MultiMap HashMap k (Set v)
m) = HashMap k [v] -> [(k, [v])]
forall k v. HashMap k v -> [(k, v)]
M.toList (HashMap k [v] -> [(k, [v])]) -> HashMap k [v] -> [(k, [v])]
forall a b. (a -> b) -> a -> b
$ (Set v -> [v]) -> HashMap k (Set v) -> HashMap k [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set v -> [v]
forall a. Set a -> [a]
S.toList) HashMap k (Set v)
m

-------------------------------------------------------------------------------

-- | Return the value to which the specified key is mapped, or 'Nothing' if
-- this map contains no mapping for the key.
lookup :: (Eq k, Hashable k) => k -> MultiMap k v -> S.Set v
lookup :: k -> MultiMap k v -> Set v
lookup k
k (MultiMap HashMap k (Set v)
m) = Set v -> Maybe (Set v) -> Set v
forall a. a -> Maybe a -> a
fromMaybe Set v
forall a. Set a
S.empty (Maybe (Set v) -> Set v) -> Maybe (Set v) -> Set v
forall a b. (a -> b) -> a -> b
$ k -> HashMap k (Set v) -> Maybe (Set v)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup k
k HashMap k (Set v)
m

-- | Associate the specified value with the specified key in this map.
--
-- If this map previously contained a mapping for the key, the new value is
-- inserted in the set, and does not replace the previous mapping.
insert :: (Eq k, Hashable k, Ord v) => k -> v -> MultiMap k v -> MultiMap k v
insert :: k -> v -> MultiMap k v -> MultiMap k v
insert k
k v
v (MultiMap HashMap k (Set v)
m) = HashMap k (Set v) -> MultiMap k v
forall k v. HashMap k (Set v) -> MultiMap k v
MultiMap (HashMap k (Set v) -> MultiMap k v)
-> HashMap k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Set v -> Set v)
-> k -> Set v -> HashMap k (Set v) -> HashMap k (Set v)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
M.insertWith (Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
S.union) k
k (v -> Set v
forall a. a -> Set a
S.singleton v
v) HashMap k (Set v)
m

-- | Returns a list of this map's keys.
keys :: MultiMap k v -> [k]
keys :: MultiMap k v -> [k]
keys = HashMap k (Set v) -> [k]
forall k v. HashMap k v -> [k]
M.keys (HashMap k (Set v) -> [k])
-> (MultiMap k v -> HashMap k (Set v)) -> MultiMap k v -> [k]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiMap k v -> HashMap k (Set v)
forall k v. MultiMap k v -> HashMap k (Set v)
unMultiMap

-- | Returns a list of this map's set of values.
elems :: MultiMap k v -> [S.Set v]
elems :: MultiMap k v -> [Set v]
elems = HashMap k (Set v) -> [Set v]
forall k v. HashMap k v -> [v]
M.elems (HashMap k (Set v) -> [Set v])
-> (MultiMap k v -> HashMap k (Set v)) -> MultiMap k v -> [Set v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiMap k v -> HashMap k (Set v)
forall k v. MultiMap k v -> HashMap k (Set v)
unMultiMap