-- | This module provides set-specific multimap functionality.
module Data.Multimap.Set (
  SetMultimap,
  map,
  delete,
  member', notMember'
) where

import Prelude hiding (map)

import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set

import Data.Multimap.Generic (Multimap(..), modifyMany)

-- | A multimap with 'Set' values. This multimap implementation will automatically deduplicate
-- values per key. For example:
--
-- > let mm = fromList [('a', 1), ('a', 1)] :: SetMultimap Char Int
-- > size mm == 1 -- True
--
--  See "Data.Multimap.Set" for operations specific to this type.
type SetMultimap = Multimap Set

-- | /O(n * log n)/ Maps over the multimap's values. This function is useful since 'Set' is not a
-- functor.
map :: (Ord k, Ord v1 , Ord v2) => (v1 -> v2) -> SetMultimap k v1 -> SetMultimap k v2
map :: (v1 -> v2) -> SetMultimap k v1 -> SetMultimap k v2
map v1 -> v2
f (Multimap Map k (Set v1)
m) = Map k (Set v2) -> SetMultimap k v2
forall (c :: * -> *) k v. Map k (c v) -> Multimap c k v
Multimap (Map k (Set v2) -> SetMultimap k v2)
-> Map k (Set v2) -> SetMultimap k v2
forall a b. (a -> b) -> a -> b
$ (Set v1 -> Set v2) -> Map k (Set v1) -> Map k (Set v2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((v1 -> v2) -> Set v1 -> Set v2
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map v1 -> v2
f) Map k (Set v1)
m

-- | /O(log m)/ Deletes an entry from the multimap.
delete :: (Ord k, Ord v) => k -> v -> SetMultimap k v -> SetMultimap k v
delete :: k -> v -> SetMultimap k v -> SetMultimap k v
delete k
k v
v = (Set v -> Set v) -> k -> SetMultimap k v -> SetMultimap k v
forall (c :: * -> *) v k.
(Collection c, Monoid (c v), Ord k) =>
(c v -> c v) -> k -> Multimap c k v -> Multimap c k v
modifyMany (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.delete v
v) k
k

-- | /O(log n)/ Checks whether an entry (key plus value) is a member of the multimap.
member' :: (Ord k, Ord v) => k -> v -> SetMultimap k v -> Bool
member' :: k -> v -> SetMultimap k v -> Bool
member' k
k v
v (Multimap Map k (Set v)
m) = Bool -> (Set v -> Bool) -> Maybe (Set v) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
v) (Maybe (Set v) -> Bool) -> Maybe (Set v) -> Bool
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

-- | /O(log n)/ Checks whether an entry (key plus value) is a not member of the multimap.
notMember' :: (Ord k, Ord v) => k -> v -> SetMultimap k v -> Bool
notMember' :: k -> v -> SetMultimap k v -> Bool
notMember' k
k v
v SetMultimap k v
mm = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ k -> v -> SetMultimap k v -> Bool
forall k v. (Ord k, Ord v) => k -> v -> SetMultimap k v -> Bool
member' k
k v
v SetMultimap k v
mm