-- | This is an implementation of bidirectional multimaps.
module Control.Distributed.Process.Internal.BiMultiMap
  ( BiMultiMap
  , empty
  , singleton
  , size
  , insert
  , lookupBy1st
  , lookupBy2nd
  , delete
  , deleteAllBy1st
  , deleteAllBy2nd
  , partitionWithKeyBy1st
  , partitionWithKeyBy2nd
  , flip
  ) where

import Data.List (foldl')
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Prelude hiding (flip, lookup)

-- | A bidirectional multimaps @BiMultiMap a b v@ is a set of triplets of type
-- @(a, b, v)@.
--
-- It is possible to lookup values by using either @a@ or @b@ as keys.
--
data BiMultiMap a b v = BiMultiMap !(Map a (Set (b, v))) !(Map b (Set (a, v)))

-- The bidirectional multimap is implemented with a pair of multimaps.
--
-- Each multimap represents a set of triples, and one invariant is that both
-- multimaps should represent exactly the same set of triples.
--
-- Each of the multimaps, however, uses a different component of the triplets
-- as key. This allows to do efficient deletions by any of the two components.

-- | The empty bidirectional multimap.
empty :: BiMultiMap a b v
empty :: forall a b v. BiMultiMap a b v
empty = Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap Map a (Set (b, v))
forall k a. Map k a
Map.empty Map b (Set (a, v))
forall k a. Map k a
Map.empty

-- | A bidirectional multimap containing a single triplet.
singleton :: (Ord a, Ord b, Ord v) => a -> b -> v -> BiMultiMap a b v
singleton :: forall a b v.
(Ord a, Ord b, Ord v) =>
a -> b -> v -> BiMultiMap a b v
singleton a
a b
b v
v = a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
forall a b v.
(Ord a, Ord b, Ord v) =>
a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
insert a
a b
b v
v BiMultiMap a b v
forall a b v. BiMultiMap a b v
empty

-- | Yields the amount of triplets in the multimap.
size :: BiMultiMap a b v -> Int
size :: forall a b v. BiMultiMap a b v -> Int
size (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
_) = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Set (b, v) -> Int) -> [Set (b, v)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Set (b, v) -> Int
forall a. Set a -> Int
Set.size ([Set (b, v)] -> [Int]) -> [Set (b, v)] -> [Int]
forall a b. (a -> b) -> a -> b
$ Map a (Set (b, v)) -> [Set (b, v)]
forall k a. Map k a -> [a]
Map.elems Map a (Set (b, v))
m

-- | Inserts a triplet in the multimap.
insert :: (Ord a, Ord b, Ord v)
       => a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
insert :: forall a b v.
(Ord a, Ord b, Ord v) =>
a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
insert a
a b
b v
v (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
r) =
    Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap ((Set (b, v) -> Set (b, v) -> Set (b, v))
-> a -> Set (b, v) -> Map a (Set (b, v)) -> Map a (Set (b, v))
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith (\Set (b, v)
_new Set (b, v)
old -> (b, v) -> Set (b, v) -> Set (b, v)
forall a. Ord a => a -> Set a -> Set a
Set.insert (b
b, v
v) Set (b, v)
old)
                               a
a
                               ((b, v) -> Set (b, v)
forall a. a -> Set a
Set.singleton (b
b, v
v))
                               Map a (Set (b, v))
m)
               ((Set (a, v) -> Set (a, v) -> Set (a, v))
-> b -> Set (a, v) -> Map b (Set (a, v)) -> Map b (Set (a, v))
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith (\Set (a, v)
_new Set (a, v)
old -> (a, v) -> Set (a, v) -> Set (a, v)
forall a. Ord a => a -> Set a -> Set a
Set.insert (a
a, v
v) Set (a, v)
old)
                               b
b
                               ((a, v) -> Set (a, v)
forall a. a -> Set a
Set.singleton (a
a, v
v))
                               Map b (Set (a, v))
r)

-- | Looks up all the triplets whose first component is the given value.
lookupBy1st :: Ord a => a -> BiMultiMap a b v -> Set (b, v)
lookupBy1st :: forall a b v. Ord a => a -> BiMultiMap a b v -> Set (b, v)
lookupBy1st a
a (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
_) = Set (b, v)
-> (Set (b, v) -> Set (b, v)) -> Maybe (Set (b, v)) -> Set (b, v)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set (b, v)
forall a. Set a
Set.empty Set (b, v) -> Set (b, v)
forall a. a -> a
id (Maybe (Set (b, v)) -> Set (b, v))
-> Maybe (Set (b, v)) -> Set (b, v)
forall a b. (a -> b) -> a -> b
$ a -> Map a (Set (b, v)) -> Maybe (Set (b, v))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a Map a (Set (b, v))
m

-- | Looks up all the triplets whose second component is the given value.
lookupBy2nd :: Ord b => b -> BiMultiMap a b v -> Set (a, v)
lookupBy2nd :: forall b a v. Ord b => b -> BiMultiMap a b v -> Set (a, v)
lookupBy2nd b
b = b -> BiMultiMap b a v -> Set (a, v)
forall a b v. Ord a => a -> BiMultiMap a b v -> Set (b, v)
lookupBy1st b
b (BiMultiMap b a v -> Set (a, v))
-> (BiMultiMap a b v -> BiMultiMap b a v)
-> BiMultiMap a b v
-> Set (a, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultiMap a b v -> BiMultiMap b a v
forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip

-- | Deletes a triplet. It yields the original multimap if the triplet is
-- not present.
delete :: (Ord a, Ord b, Ord v)
       => a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
delete :: forall a b v.
(Ord a, Ord b, Ord v) =>
a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
delete a
a b
b v
v (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
r) =
  let m' :: Map a (Set (b, v))
m' = (Set (b, v) -> Maybe (Set (b, v)))
-> a -> Map a (Set (b, v)) -> Map a (Set (b, v))
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update ((Set (b, v) -> Bool) -> Set (b, v) -> Maybe (Set (b, v))
forall a. (a -> Bool) -> a -> Maybe a
nothingWhen Set (b, v) -> Bool
forall a. Set a -> Bool
Set.null (Set (b, v) -> Maybe (Set (b, v)))
-> (Set (b, v) -> Set (b, v)) -> Set (b, v) -> Maybe (Set (b, v))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b, v) -> Set (b, v) -> Set (b, v)
forall a. Ord a => a -> Set a -> Set a
Set.delete (b
b, v
v)) a
a Map a (Set (b, v))
m
      r' :: Map b (Set (a, v))
r' = (Set (a, v) -> Maybe (Set (a, v)))
-> b -> Map b (Set (a, v)) -> Map b (Set (a, v))
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update ((Set (a, v) -> Bool) -> Set (a, v) -> Maybe (Set (a, v))
forall a. (a -> Bool) -> a -> Maybe a
nothingWhen Set (a, v) -> Bool
forall a. Set a -> Bool
Set.null (Set (a, v) -> Maybe (Set (a, v)))
-> (Set (a, v) -> Set (a, v)) -> Set (a, v) -> Maybe (Set (a, v))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, v) -> Set (a, v) -> Set (a, v)
forall a. Ord a => a -> Set a -> Set a
Set.delete (a
a, v
v)) b
b Map b (Set (a, v))
r
   in Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap Map a (Set (b, v))
m' Map b (Set (a, v))
r'

-- | Deletes all triplets whose first component is the given value.
deleteAllBy1st :: (Ord a, Ord b, Ord v) => a -> BiMultiMap a b v -> BiMultiMap a b v
deleteAllBy1st :: forall a b v.
(Ord a, Ord b, Ord v) =>
a -> BiMultiMap a b v -> BiMultiMap a b v
deleteAllBy1st a
a (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
r) =
  let (Maybe (Set (b, v))
mm, Map a (Set (b, v))
m') = (a -> Set (b, v) -> Maybe (Set (b, v)))
-> a
-> Map a (Set (b, v))
-> (Maybe (Set (b, v)), Map a (Set (b, v)))
forall k a.
Ord k =>
(k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a)
Map.updateLookupWithKey (\a
_ Set (b, v)
_ -> Maybe (Set (b, v))
forall a. Maybe a
Nothing) a
a Map a (Set (b, v))
m
      r' :: Map b (Set (a, v))
r' = case Maybe (Set (b, v))
mm of
            Maybe (Set (b, v))
Nothing -> Map b (Set (a, v))
r
            Just Set (b, v)
mb -> a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
forall a b v.
(Ord a, Ord b, Ord v) =>
a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
reverseDelete a
a (Set (b, v) -> [(b, v)]
forall a. Set a -> [a]
Set.toList Set (b, v)
mb) Map b (Set (a, v))
r
   in Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap Map a (Set (b, v))
m' Map b (Set (a, v))
r'

-- | Like 'deleteAllBy1st' but deletes by the second component of the triplets.
deleteAllBy2nd :: (Ord a, Ord b, Ord v)
               => b -> BiMultiMap a b v -> BiMultiMap a b v
deleteAllBy2nd :: forall a b v.
(Ord a, Ord b, Ord v) =>
b -> BiMultiMap a b v -> BiMultiMap a b v
deleteAllBy2nd b
b = BiMultiMap b a v -> BiMultiMap a b v
forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip (BiMultiMap b a v -> BiMultiMap a b v)
-> (BiMultiMap a b v -> BiMultiMap b a v)
-> BiMultiMap a b v
-> BiMultiMap a b v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> BiMultiMap b a v -> BiMultiMap b a v
forall a b v.
(Ord a, Ord b, Ord v) =>
a -> BiMultiMap a b v -> BiMultiMap a b v
deleteAllBy1st b
b (BiMultiMap b a v -> BiMultiMap b a v)
-> (BiMultiMap a b v -> BiMultiMap b a v)
-> BiMultiMap a b v
-> BiMultiMap b a v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultiMap a b v -> BiMultiMap b a v
forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip

-- | Yields the triplets satisfying the given predicate, and a multimap
-- with all this triplets removed.
partitionWithKeyBy1st :: (Ord a, Ord b, Ord v)
                      => (a -> Set (b, v) -> Bool) -> BiMultiMap a b v
                      -> (Map a (Set (b, v)), BiMultiMap a b v)
partitionWithKeyBy1st :: forall a b v.
(Ord a, Ord b, Ord v) =>
(a -> Set (b, v) -> Bool)
-> BiMultiMap a b v -> (Map a (Set (b, v)), BiMultiMap a b v)
partitionWithKeyBy1st a -> Set (b, v) -> Bool
p (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
r) =
    let (Map a (Set (b, v))
m0, Map a (Set (b, v))
m1) = (a -> Set (b, v) -> Bool)
-> Map a (Set (b, v)) -> (Map a (Set (b, v)), Map a (Set (b, v)))
forall k a. (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
Map.partitionWithKey a -> Set (b, v) -> Bool
p Map a (Set (b, v))
m
        r1 :: Map b (Set (a, v))
r1 = (Map b (Set (a, v)) -> (a, Set (b, v)) -> Map b (Set (a, v)))
-> Map b (Set (a, v)) -> [(a, Set (b, v))] -> Map b (Set (a, v))
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map b (Set (a, v))
rr (a
a, Set (b, v)
mb) -> a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
forall a b v.
(Ord a, Ord b, Ord v) =>
a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
reverseDelete a
a (Set (b, v) -> [(b, v)]
forall a. Set a -> [a]
Set.toList Set (b, v)
mb) Map b (Set (a, v))
rr) Map b (Set (a, v))
r ([(a, Set (b, v))] -> Map b (Set (a, v)))
-> [(a, Set (b, v))] -> Map b (Set (a, v))
forall a b. (a -> b) -> a -> b
$
               Map a (Set (b, v)) -> [(a, Set (b, v))]
forall k a. Map k a -> [(k, a)]
Map.toList Map a (Set (b, v))
m0
     in (Map a (Set (b, v))
m0, Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap Map a (Set (b, v))
m1 Map b (Set (a, v))
r1)

-- | Like 'partitionWithKeyBy1st' but the predicates takes the second component
-- of the triplets as first argument.
partitionWithKeyBy2nd :: (Ord a, Ord b, Ord v)
                      => (b -> Set (a, v) -> Bool) -> BiMultiMap a b v
                      -> (Map b (Set (a, v)), BiMultiMap a b v)
partitionWithKeyBy2nd :: forall a b v.
(Ord a, Ord b, Ord v) =>
(b -> Set (a, v) -> Bool)
-> BiMultiMap a b v -> (Map b (Set (a, v)), BiMultiMap a b v)
partitionWithKeyBy2nd b -> Set (a, v) -> Bool
p BiMultiMap a b v
b = let (Map b (Set (a, v))
m, BiMultiMap b a v
b') = (b -> Set (a, v) -> Bool)
-> BiMultiMap b a v -> (Map b (Set (a, v)), BiMultiMap b a v)
forall a b v.
(Ord a, Ord b, Ord v) =>
(a -> Set (b, v) -> Bool)
-> BiMultiMap a b v -> (Map a (Set (b, v)), BiMultiMap a b v)
partitionWithKeyBy1st b -> Set (a, v) -> Bool
p (BiMultiMap b a v -> (Map b (Set (a, v)), BiMultiMap b a v))
-> BiMultiMap b a v -> (Map b (Set (a, v)), BiMultiMap b a v)
forall a b. (a -> b) -> a -> b
$ BiMultiMap a b v -> BiMultiMap b a v
forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip BiMultiMap a b v
b
                             in (Map b (Set (a, v))
m, BiMultiMap b a v -> BiMultiMap a b v
forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip BiMultiMap b a v
b')

-- | Exchange the first and the second components of all triplets.
flip :: BiMultiMap a b v -> BiMultiMap b a v
flip :: forall a b v. BiMultiMap a b v -> BiMultiMap b a v
flip (BiMultiMap Map a (Set (b, v))
m Map b (Set (a, v))
r) = Map b (Set (a, v)) -> Map a (Set (b, v)) -> BiMultiMap b a v
forall a b v.
Map a (Set (b, v)) -> Map b (Set (a, v)) -> BiMultiMap a b v
BiMultiMap Map b (Set (a, v))
r Map a (Set (b, v))
m

-- Internal functions

-- | @reverseDelete a bs m@ removes from @m@ all the triplets wich have @a@ as
-- first component and second and third components in @bs@.
--
-- The @m@ map is in reversed form, meaning that the second component of the
-- triplets is used as key.
reverseDelete :: (Ord a, Ord b, Ord v)
              => a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
reverseDelete :: forall a b v.
(Ord a, Ord b, Ord v) =>
a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
reverseDelete a
a [(b, v)]
bs Map b (Set (a, v))
r = (Map b (Set (a, v)) -> (b, v) -> Map b (Set (a, v)))
-> Map b (Set (a, v)) -> [(b, v)] -> Map b (Set (a, v))
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map b (Set (a, v))
rr (b
b, v
v) -> (Set (a, v) -> Maybe (Set (a, v)))
-> b -> Map b (Set (a, v)) -> Map b (Set (a, v))
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (v -> Set (a, v) -> Maybe (Set (a, v))
forall {b}. Ord b => b -> Set (a, b) -> Maybe (Set (a, b))
rmb v
v) b
b Map b (Set (a, v))
rr) Map b (Set (a, v))
r [(b, v)]
bs
  where
    rmb :: b -> Set (a, b) -> Maybe (Set (a, b))
rmb b
v = (Set (a, b) -> Bool) -> Set (a, b) -> Maybe (Set (a, b))
forall a. (a -> Bool) -> a -> Maybe a
nothingWhen Set (a, b) -> Bool
forall a. Set a -> Bool
Set.null (Set (a, b) -> Maybe (Set (a, b)))
-> (Set (a, b) -> Set (a, b)) -> Set (a, b) -> Maybe (Set (a, b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> Set (a, b) -> Set (a, b)
forall a. Ord a => a -> Set a -> Set a
Set.delete (a
a, b
v)

-- | @nothingWhen p a@ is @Just a@ when @a@ satisfies predicate @p@.
-- Yields @Nothing@ otherwise.
nothingWhen :: (a -> Bool) -> a -> Maybe a
nothingWhen :: forall a. (a -> Bool) -> a -> Maybe a
nothingWhen a -> Bool
p a
a = if a -> Bool
p a
a then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just a
a