module Data.MultiMap
    ( MultiMap (..)
    , assocs
    , delete
    , deleteMany
    , empty
    , insert
    , keysSet
    , lookup
    , singleton
    ) where

import           Prelude hiding (lookup)

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Maybe (fromMaybe)
import           Data.Set (Set)
import qualified Data.Set as Set

newtype MultiMap k v = MultiMap (Map k (Set v))
    deriving (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)

assocs :: MultiMap k v -> [(k, [v])]
assocs :: MultiMap k v -> [(k, [v])]
assocs (MultiMap Map k (Set v)
m) = Map k [v] -> [(k, [v])]
forall k a. Map k a -> [(k, a)]
Map.assocs (Map k [v] -> [(k, [v])]) -> Map k [v] -> [(k, [v])]
forall a b. (a -> b) -> a -> b
$ Set v -> [v]
forall a. Set a -> [a]
Set.toList (Set v -> [v]) -> Map k (Set v) -> Map k [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (Set v)
m

delete :: (Ord k, Ord v) => k -> v -> MultiMap k v -> MultiMap k v
delete :: k -> v -> MultiMap k v -> MultiMap k v
delete k
k v
v (MultiMap Map k (Set v)
m) = Map k (Set v) -> MultiMap k v
forall k v. Map k (Set v) -> MultiMap k v
MultiMap (Map k (Set v) -> MultiMap k v) -> Map k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Maybe (Set v)) -> k -> Map k (Set v) -> Map k (Set v)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update Set v -> Maybe (Set v)
delete' k
k Map k (Set v)
m
  where
    delete' :: Set v -> Maybe (Set v)
delete' Set v
s = let s' :: Set v
s' = v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.delete v
v Set v
s in if Set v -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set v
s' then Maybe (Set v)
forall a. Maybe a
Nothing else Set v -> Maybe (Set v)
forall a. a -> Maybe a
Just Set v
s'

deleteMany ::
    (Ord k, Ord v) => k -> Set v -> MultiMap k v -> MultiMap k v
deleteMany :: k -> Set v -> MultiMap k v -> MultiMap k v
deleteMany k
k Set v
vs (MultiMap Map k (Set v)
m) = Map k (Set v) -> MultiMap k v
forall k v. Map k (Set v) -> MultiMap k v
MultiMap (Map k (Set v) -> MultiMap k v) -> Map k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Maybe (Set v)) -> k -> Map k (Set v) -> Map k (Set v)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update Set v -> Maybe (Set v)
deleteMany' k
k Map k (Set v)
m
  where
    deleteMany' :: Set v -> Maybe (Set v)
deleteMany' Set v
s =
        let s' :: Set v
s' = Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set v
s Set v
vs in if Set v -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Set v
s' then Maybe (Set v)
forall a. Maybe a
Nothing else Set v -> Maybe (Set v)
forall a. a -> Maybe a
Just Set v
s'

empty :: MultiMap k v
empty :: MultiMap k v
empty = Map k (Set v) -> MultiMap k v
forall k v. Map k (Set v) -> MultiMap k v
MultiMap Map k (Set v)
forall k a. Map k a
Map.empty

insert :: (Ord 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 Map k (Set v)
m) =
    Map k (Set v) -> MultiMap k v
forall k v. Map k (Set v) -> MultiMap k v
MultiMap (Map k (Set v) -> MultiMap k v) -> Map k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (Set v -> Set v -> Set v)
-> k -> Set v -> Map k (Set v) -> Map k (Set v)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Set v -> Set v -> Set v
forall a. Semigroup a => a -> a -> a
(<>) k
k (v -> Set v
forall a. a -> Set a
Set.singleton v
v) Map k (Set v)
m

keysSet :: MultiMap k v -> Set k
keysSet :: MultiMap k v -> Set k
keysSet (MultiMap Map k (Set v)
m) = Map k (Set v) -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k (Set v)
m

-- | If no key in the map then the result is empty.
lookup :: Ord k => k -> MultiMap k v -> Set v
lookup :: k -> MultiMap k v -> Set v
lookup k
k (MultiMap Map k (Set v)
m) = Set v -> Maybe (Set v) -> Set v
forall a. a -> Maybe a -> a
fromMaybe Set v
forall a. Set a
Set.empty (Maybe (Set v) -> Set v) -> Maybe (Set v) -> Set v
forall a b. (a -> b) -> a -> b
$ k -> Map k (Set v) -> Maybe (Set v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (Set v)
m

singleton :: k -> v -> MultiMap k v
singleton :: k -> v -> MultiMap k v
singleton k
k v
v = Map k (Set v) -> MultiMap k v
forall k v. Map k (Set v) -> MultiMap k v
MultiMap (Map k (Set v) -> MultiMap k v) -> Map k (Set v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ k -> Set v -> Map k (Set v)
forall k a. k -> a -> Map k a
Map.singleton k
k (Set v -> Map k (Set v)) -> Set v -> Map k (Set v)
forall a b. (a -> b) -> a -> b
$ v -> Set v
forall a. a -> Set a
Set.singleton v
v