1
-- | This is an implementation of bidirectional multimaps.
2
module Control.Distributed.Process.Internal.BiMultiMap
3
  ( BiMultiMap
4
  , empty
5
  , singleton
6
  , size
7
  , insert
8
  , lookupBy1st
9
  , lookupBy2nd
10
  , delete
11
  , deleteAllBy1st
12
  , deleteAllBy2nd
13
  , partitionWithKeyBy1st
14
  , partitionWithKeyBy2nd
15
  , flip
16
  ) where
17
18
import Data.List (foldl')
19
import Data.Map.Strict (Map)
20
import qualified Data.Map.Strict as Map
21
import Data.Set (Set)
22
import qualified Data.Set as Set
23
import Prelude hiding (flip, lookup)
24
25
-- | A bidirectional multimaps @BiMultiMap a b v@ is a set of triplets of type
26
-- @(a, b, v)@.
27
--
28
-- It is possible to lookup values by using either @a@ or @b@ as keys.
29
--
30
data BiMultiMap a b v = BiMultiMap !(Map a (Set (b, v))) !(Map b (Set (a, v)))
31
32
-- The bidirectional multimap is implemented with a pair of multimaps.
33
--
34
-- Each multimap represents a set of triples, and one invariant is that both
35
-- multimaps should represent exactly the same set of triples.
36
--
37
-- Each of the multimaps, however, uses a different component of the triplets
38
-- as key. This allows to do efficient deletions by any of the two components.
39
40
-- | The empty bidirectional multimap.
41
empty :: BiMultiMap a b v
42
empty = BiMultiMap Map.empty Map.empty
43
44
-- | A bidirectional multimap containing a single triplet.
45
singleton :: (Ord a, Ord b, Ord v) => a -> b -> v -> BiMultiMap a b v
46
singleton a b v = insert a b v empty
47
48
-- | Yields the amount of triplets in the multimap.
49
size :: BiMultiMap a b v -> Int
50
size (BiMultiMap m _) = foldl' (+) 0 $ map Set.size $ Map.elems m
51
52
-- | Inserts a triplet in the multimap.
53
insert :: (Ord a, Ord b, Ord v)
54
       => a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
55
insert a b v (BiMultiMap m r) =
56
    BiMultiMap (Map.insertWith (\_new old -> Set.insert (b, v) old)
57
                               a
58
                               (Set.singleton (b, v))
59
                               m)
60
               (Map.insertWith (\_new old -> Set.insert (a, v) old)
61
                               b
62
                               (Set.singleton (a, v))
63
                               r)
64
65
-- | Looks up all the triplets whose first component is the given value.
66
lookupBy1st :: Ord a => a -> BiMultiMap a b v -> Set (b, v)
67
lookupBy1st a (BiMultiMap m _) = maybe Set.empty id $ Map.lookup a m
68
69
-- | Looks up all the triplets whose second component is the given value.
70
lookupBy2nd :: Ord b => b -> BiMultiMap a b v -> Set (a, v)
71
lookupBy2nd b = lookupBy1st b . flip
72
73
-- | Deletes a triplet. It yields the original multimap if the triplet is
74
-- not present.
75
delete :: (Ord a, Ord b, Ord v)
76
       => a -> b -> v -> BiMultiMap a b v -> BiMultiMap a b v
77
delete a b v (BiMultiMap m r) =
78
  let m' = Map.update (nothingWhen Set.null . Set.delete (b, v)) a m
79
      r' = Map.update (nothingWhen Set.null . Set.delete (a, v)) b r
80
   in BiMultiMap m' r'
81
82
-- | Deletes all triplets whose first component is the given value.
83
deleteAllBy1st :: (Ord a, Ord b, Ord v) => a -> BiMultiMap a b v -> BiMultiMap a b v
84
deleteAllBy1st a (BiMultiMap m r) =
85
  let (mm, m') = Map.updateLookupWithKey (\_ _ -> Nothing) a m
86
      r' = case mm of
87
            Nothing -> r
88
            Just mb -> reverseDelete a (Set.toList mb) r
89
   in BiMultiMap m' r'
90
91
-- | Like 'deleteAllBy1st' but deletes by the second component of the triplets.
92
deleteAllBy2nd :: (Ord a, Ord b, Ord v)
93
               => b -> BiMultiMap a b v -> BiMultiMap a b v
94
deleteAllBy2nd b = flip . deleteAllBy1st b . flip
95
96
-- | Yields the triplets satisfying the given predicate, and a multimap
97
-- with all this triplets removed.
98
partitionWithKeyBy1st :: (Ord a, Ord b, Ord v)
99
                      => (a -> Set (b, v) -> Bool) -> BiMultiMap a b v
100
                      -> (Map a (Set (b, v)), BiMultiMap a b v)
101
partitionWithKeyBy1st p (BiMultiMap m r) =
102
    let (m0, m1) = Map.partitionWithKey p m
103
        r1 = foldl' (\rr (a, mb) -> reverseDelete a (Set.toList mb) rr) r $
104
               Map.toList m0
105
     in (m0, BiMultiMap m1 r1)
106
107
-- | Like 'partitionWithKeyBy1st' but the predicates takes the second component
108
-- of the triplets as first argument.
109
partitionWithKeyBy2nd :: (Ord a, Ord b, Ord v)
110
                      => (b -> Set (a, v) -> Bool) -> BiMultiMap a b v
111
                      -> (Map b (Set (a, v)), BiMultiMap a b v)
112
partitionWithKeyBy2nd p b = let (m, b') = partitionWithKeyBy1st p $ flip b
113
                             in (m, flip b')
114
115
-- | Exchange the first and the second components of all triplets.
116
flip :: BiMultiMap a b v -> BiMultiMap b a v
117
flip (BiMultiMap m r) = BiMultiMap r m
118
119
-- Internal functions
120
121
-- | @reverseDelete a bs m@ removes from @m@ all the triplets wich have @a@ as
122
-- first component and second and third components in @bs@.
123
--
124
-- The @m@ map is in reversed form, meaning that the second component of the
125
-- triplets is used as key.
126
reverseDelete :: (Ord a, Ord b, Ord v)
127
              => a -> [(b, v)] -> Map b (Set (a, v)) -> Map b (Set (a, v))
128
reverseDelete a bs r = foldl' (\rr (b, v) -> Map.update (rmb v) b rr) r bs
129
  where
130
    rmb v = nothingWhen Set.null . Set.delete (a, v)
131
132
-- | @nothingWhen p a@ is @Just a@ when @a@ satisfies predicate @p@.
133
-- Yields @Nothing@ otherwise.
134
nothingWhen :: (a -> Bool) -> a -> Maybe a
135
nothingWhen p a = if p a then Nothing else Just a