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