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