{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE ExistentialQuantification  #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE TupleSections              #-}

module Control.Distributed.Process.Extras.Internal.Containers.MultiMap
  ( MultiMap
  , Insertable
  , empty
  , insert
  , member
  , lookup
  , delete
  , filter
  , filterWithKey
  , foldrWithKey
  , toList
  , size
  ) where

import qualified Data.Foldable as Foldable
import Data.Foldable (Foldable)
import Data.Hashable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as Map
import Data.HashSet (HashSet)
import qualified Data.HashSet as Set
import Prelude hiding (lookup, filter, pred)

-- | Class of things that can be inserted in a map or
-- a set (of mapped values), for which instances of
-- @Eq@ and @Hashable@ must be present.
--
class (Eq a, Hashable a) => Insertable a
instance (Eq a, Hashable a) => Insertable a

-- | Opaque type of MultiMaps.
data MultiMap k v = M { forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: !(HashMap k (HashSet v)) }

-- instance Foldable

instance Foldable (MultiMap k) where
  foldr :: forall a b. (a -> b -> b) -> b -> MultiMap k a -> b
foldr a -> b -> b
f = (k -> a -> b -> b) -> b -> MultiMap k a -> b
forall k v a. (k -> v -> a -> a) -> a -> MultiMap k v -> a
foldrWithKey ((a -> b -> b) -> k -> a -> b -> b
forall a b. a -> b -> a
const a -> b -> b
f)

empty :: MultiMap k v
empty :: forall k v. MultiMap k v
empty = HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ HashMap k (HashSet v)
forall k v. HashMap k v
Map.empty

size :: MultiMap k v -> Int
size :: forall k a. MultiMap k a -> Int
size = HashMap k (HashSet v) -> Int
forall k v. HashMap k v -> Int
Map.size (HashMap k (HashSet v) -> Int)
-> (MultiMap k v -> HashMap k (HashSet v)) -> MultiMap k v -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiMap k v -> HashMap k (HashSet v)
forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap

insert :: forall k v. (Insertable k, Insertable v)
       => k -> v -> MultiMap k v -> MultiMap k v
insert :: forall k v.
(Insertable k, Insertable v) =>
k -> v -> MultiMap k v -> MultiMap k v
insert k
k' v
v' M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} =
  case k -> HashMap k (HashSet v) -> Maybe (HashSet v)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup k
k' HashMap k (HashSet v)
hmap of
    Maybe (HashSet v)
Nothing -> HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ k -> HashSet v -> HashMap k (HashSet v) -> HashMap k (HashSet v)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert k
k' (v -> HashSet v
forall a. Hashable a => a -> HashSet a
Set.singleton v
v') HashMap k (HashSet v)
hmap
    Just HashSet v
s  -> HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ k -> HashSet v -> HashMap k (HashSet v) -> HashMap k (HashSet v)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert k
k' (v -> HashSet v -> HashSet v
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert v
v' HashSet v
s) HashMap k (HashSet v)
hmap
{-# INLINE insert #-}

member :: (Insertable k) => k -> MultiMap k a -> Bool
member :: forall k a. Insertable k => k -> MultiMap k a -> Bool
member k
k = k -> HashMap k (HashSet a) -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
Map.member k
k (HashMap k (HashSet a) -> Bool)
-> (MultiMap k a -> HashMap k (HashSet a)) -> MultiMap k a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiMap k a -> HashMap k (HashSet a)
forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap

lookup :: (Insertable k) => k -> MultiMap k v -> Maybe [v]
lookup :: forall k v. Insertable k => k -> MultiMap k v -> Maybe [v]
lookup k
k M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} = Maybe [v]
-> (HashSet v -> Maybe [v]) -> Maybe (HashSet v) -> Maybe [v]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe [v]
forall a. Maybe a
Nothing ([v] -> Maybe [v]
forall a. a -> Maybe a
Just ([v] -> Maybe [v]) -> (HashSet v -> [v]) -> HashSet v -> Maybe [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashSet v -> [v]
forall a. HashSet a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList) (Maybe (HashSet v) -> Maybe [v]) -> Maybe (HashSet v) -> Maybe [v]
forall a b. (a -> b) -> a -> b
$ k -> HashMap k (HashSet v) -> Maybe (HashSet v)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup k
k HashMap k (HashSet v)
hmap
{-# INLINE lookup #-}

delete :: (Insertable k) => k -> MultiMap k v -> Maybe ([v], MultiMap k v)
delete :: forall k v.
Insertable k =>
k -> MultiMap k v -> Maybe ([v], MultiMap k v)
delete k
k m :: MultiMap k v
m@M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} = Maybe ([v], MultiMap k v)
-> ([v] -> Maybe ([v], MultiMap k v))
-> Maybe [v]
-> Maybe ([v], MultiMap k v)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe ([v], MultiMap k v)
forall a. Maybe a
Nothing (([v], MultiMap k v) -> Maybe ([v], MultiMap k v)
forall a. a -> Maybe a
Just (([v], MultiMap k v) -> Maybe ([v], MultiMap k v))
-> ([v] -> ([v], MultiMap k v)) -> [v] -> Maybe ([v], MultiMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (, HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ k -> HashMap k (HashSet v) -> HashMap k (HashSet v)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
Map.delete k
k HashMap k (HashSet v)
hmap)) (Maybe [v] -> Maybe ([v], MultiMap k v))
-> Maybe [v] -> Maybe ([v], MultiMap k v)
forall a b. (a -> b) -> a -> b
$ k -> MultiMap k v -> Maybe [v]
forall k v. Insertable k => k -> MultiMap k v -> Maybe [v]
lookup k
k MultiMap k v
m

filter :: forall k v. (Insertable k)
       => (v -> Bool)
       -> MultiMap k v
       -> MultiMap k v
filter :: forall k v.
Insertable k =>
(v -> Bool) -> MultiMap k v -> MultiMap k v
filter v -> Bool
p M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} = HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (HashMap k (HashSet v) -> k -> HashSet v -> HashMap k (HashSet v))
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
Map.foldlWithKey' ((v -> Bool)
-> HashMap k (HashSet v) -> k -> HashSet v -> HashMap k (HashSet v)
forall {k} {a}.
Hashable k =>
(a -> Bool)
-> HashMap k (HashSet a) -> k -> HashSet a -> HashMap k (HashSet a)
matchOn v -> Bool
p) HashMap k (HashSet v)
hmap HashMap k (HashSet v)
hmap
  where
    matchOn :: (a -> Bool)
-> HashMap k (HashSet a) -> k -> HashSet a -> HashMap k (HashSet a)
matchOn a -> Bool
pred HashMap k (HashSet a)
acc k
key HashSet a
valueSet =
      let vs :: HashSet a
vs = (a -> Bool) -> HashSet a -> HashSet a
forall a. (a -> Bool) -> HashSet a -> HashSet a
Set.filter a -> Bool
pred HashSet a
valueSet in
      if HashSet a -> Bool
forall a. HashSet a -> Bool
Set.null HashSet a
vs then HashMap k (HashSet a)
acc else k -> HashSet a -> HashMap k (HashSet a) -> HashMap k (HashSet a)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert k
key HashSet a
vs HashMap k (HashSet a)
acc
{-# INLINE filter #-}

filterWithKey :: forall k v. (Insertable k)
              => (k -> v -> Bool)
              -> MultiMap k v
              -> MultiMap k v
filterWithKey :: forall k v.
Insertable k =>
(k -> v -> Bool) -> MultiMap k v -> MultiMap k v
filterWithKey k -> v -> Bool
p M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} = HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
M (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (HashMap k (HashSet v) -> k -> HashSet v -> HashMap k (HashSet v))
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
Map.foldlWithKey' ((k -> v -> Bool)
-> HashMap k (HashSet v) -> k -> HashSet v -> HashMap k (HashSet v)
forall {k} {a}.
Hashable k =>
(k -> a -> Bool)
-> HashMap k (HashSet a) -> k -> HashSet a -> HashMap k (HashSet a)
matchOn k -> v -> Bool
p) HashMap k (HashSet v)
hmap HashMap k (HashSet v)
hmap
  where
    matchOn :: (k -> a -> Bool)
-> HashMap k (HashSet a) -> k -> HashSet a -> HashMap k (HashSet a)
matchOn k -> a -> Bool
pred HashMap k (HashSet a)
acc k
key HashSet a
valueSet =
      let vs :: HashSet a
vs = (a -> Bool) -> HashSet a -> HashSet a
forall a. (a -> Bool) -> HashSet a -> HashSet a
Set.filter (k -> a -> Bool
pred k
key) HashSet a
valueSet in
      if HashSet a -> Bool
forall a. HashSet a -> Bool
Set.null HashSet a
vs then HashMap k (HashSet a)
acc else k -> HashSet a -> HashMap k (HashSet a) -> HashMap k (HashSet a)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert k
key HashSet a
vs HashMap k (HashSet a)
acc
{-# INLINE filterWithKey #-}

-- | /O(n)/ Reduce this map by applying a binary operator to all
-- elements, using the given starting value (typically the
-- right-identity of the operator).
foldrWithKey :: (k -> v -> a -> a) -> a -> MultiMap k v -> a
foldrWithKey :: forall k v a. (k -> v -> a -> a) -> a -> MultiMap k v -> a
foldrWithKey k -> v -> a -> a
f a
a M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} =
  let wrap :: k -> v -> a -> a
wrap = \k
k' v
v' a
acc' -> k -> v -> a -> a
f k
k' v
v' a
acc'
  in (k -> HashSet v -> a -> a) -> a -> HashMap k (HashSet v) -> a
forall k v a. (k -> v -> a -> a) -> a -> HashMap k v -> a
Map.foldrWithKey (\k
k HashSet v
v a
acc -> (v -> a -> a) -> a -> HashSet v -> a
forall b a. (b -> a -> a) -> a -> HashSet b -> a
Set.foldr (k -> v -> a -> a
wrap k
k) a
acc HashSet v
v) a
a HashMap k (HashSet v)
hmap
{-# INLINE foldrWithKey #-}

toList :: MultiMap k v -> [(k, v)]
toList :: forall k v. MultiMap k v -> [(k, v)]
toList M{HashMap k (HashSet v)
hmap :: forall k v. MultiMap k v -> HashMap k (HashSet v)
hmap :: HashMap k (HashSet v)
..} = ([(k, v)] -> k -> HashSet v -> [(k, v)])
-> [(k, v)] -> HashMap k (HashSet v) -> [(k, v)]
forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
Map.foldlWithKey' [(k, v)] -> k -> HashSet v -> [(k, v)]
forall {a} {b}. [(a, b)] -> a -> HashSet b -> [(a, b)]
explode [] HashMap k (HashSet v)
hmap
  where
    explode :: [(a, b)] -> a -> HashSet b -> [(a, b)]
explode [(a, b)]
xs a
k HashSet b
vs = ([(a, b)] -> b -> [(a, b)]) -> [(a, b)] -> HashSet b -> [(a, b)]
forall a b. (a -> b -> a) -> a -> HashSet b -> a
Set.foldl' (\[(a, b)]
ys b
v -> ((a
k, b
v)(a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
:[(a, b)]
ys)) [(a, b)]
xs HashSet b
vs
{-# INLINE toList #-}