-- | 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 = BiMultiMap Map.empty Map.empty -- | A bidirectional multimap containing a single triplet. singleton :: (Ord a, Ord b, Ord v) => a -> b -> v -> BiMultiMap a b v singleton a b v = insert a b v empty -- | Yields the amount of triplets in the multimap. size :: BiMultiMap a b v -> Int size (BiMultiMap m _) = foldl' (+) 0 $ map Set.size $ Map.elems 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 a b v (BiMultiMap m r) = BiMultiMap (Map.insertWith (\_new old -> Set.insert (b, v) old) a (Set.singleton (b, v)) m) (Map.insertWith (\_new old -> Set.insert (a, v) old) b (Set.singleton (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 a (BiMultiMap m _) = maybe Set.empty id $ Map.lookup a 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 b = lookupBy1st b . 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 a b v (BiMultiMap m r) = let m' = Map.update (nothingWhen Set.null . Set.delete (b, v)) a m r' = Map.update (nothingWhen Set.null . Set.delete (a, v)) b r in BiMultiMap m' 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 a (BiMultiMap m r) = let (mm, m') = Map.updateLookupWithKey (\_ _ -> Nothing) a m r' = case mm of Nothing -> r Just mb -> reverseDelete a (Set.toList mb) r in BiMultiMap m' 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 b = flip . deleteAllBy1st b . 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 p (BiMultiMap m r) = let (m0, m1) = Map.partitionWithKey p m r1 = foldl' (\rr (a, mb) -> reverseDelete a (Set.toList mb) rr) r $ Map.toList m0 in (m0, BiMultiMap m1 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 p b = let (m, b') = partitionWithKeyBy1st p $ flip b in (m, flip b') -- | Exchange the first and the second components of all triplets. flip :: BiMultiMap a b v -> BiMultiMap b a v flip (BiMultiMap m r) = BiMultiMap r 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 a bs r = foldl' (\rr (b, v) -> Map.update (rmb v) b rr) r bs where rmb v = nothingWhen Set.null . Set.delete (a, v) -- | @nothingWhen p a@ is @Just a@ when @a@ satisfies predicate @p@. -- Yields @Nothing@ otherwise. nothingWhen :: (a -> Bool) -> a -> Maybe a nothingWhen p a = if p a then Nothing else Just a