{-# LANGUAGE CPP #-} {-# LANGUAGE DeriveDataTypeable #-} {-# LANGUAGE DeriveFunctor #-} {-# LANGUAGE TypeFamilies #-} -- | This internal module implements functionality shared by all multimaps. module Data.Multimap.Generic ( Multimap(..), Group, null, size, distinctSize, empty, singleton, #if __GLASGOW_HASKELL__ >= 708 fromList, inverse, #endif fromListWith, fromGroupList, fromMap, member, notMember, count, find, (!), prepend, prependMany, append, appendMany, deleteMany, inverseWith, filter, filterGroups, mapGroups, toList, toGroupList, toMap, keys, keysSet, keysMultiset, maxViewWith, minViewWith, modifyMany, modifyManyF ) where import Data.Multimap.Collection (Collection) import qualified Data.Multimap.Collection as Col import Data.Multiset (Multiset) import qualified Data.Multiset as Mset import Prelude hiding (filter, foldr, null) import qualified Prelude as Prelude import Data.Binary (Binary(..)) import Data.Data (Data, Typeable) import Data.Foldable (foldl', foldr) import Data.Map.Strict (Map) import qualified Data.Map.Strict as Map import Data.Maybe (fromMaybe) import Data.Semigroup (Semigroup, (<>)) import Data.Set (Set) import Data.Tuple (swap) import qualified GHC.Exts -- | A map where the same key can be present multiple times. newtype Multimap c k v = Multimap { _toMap :: Map k (c v) } deriving ( Eq, Ord, Read, Show, Functor, {-| @since 0.2.1.1 -} Data, {-| @since 0.2.1.1 -} Typeable ) -- | A group of values. type Group k cv = (k, cv) instance (Ord k, Semigroup (c v)) => Semigroup (Multimap c k v) where Multimap m1 <> Multimap m2 = Multimap $ Map.unionWith (<>) m1 m2 instance (Ord k, Monoid (c v)) => Monoid (Multimap c k v) where mempty = empty instance Foldable c => Foldable (Multimap c k) where foldr f r0 (Multimap m) = Map.foldr (\c r -> foldr f r c) r0 m instance Traversable c => Traversable (Multimap c k) where sequenceA (Multimap m) = Multimap <$> sequenceA (fmap sequenceA m) -- | @since 0.2.1.0 instance (Binary k, Binary (c v)) => Binary (Multimap c k v) where put (Multimap m) = put m get = Multimap <$> get #if __GLASGOW_HASKELL__ >= 708 instance (Collection c, GHC.Exts.IsList (c v), GHC.Exts.Item (c v) ~ v, Ord k) => GHC.Exts.IsList (Multimap c k v) where type Item (Multimap c k v) = (k, v) fromList = fromList toList = toList -- | /O(n * log n)/ Builds a multimap from a list of key, value tuples. The values are in the same -- order as in the original list. fromList :: (Collection c, GHC.Exts.IsList (c v), GHC.Exts.Item (c v) ~ v, Ord k) => [(k, v)] -> Multimap c k v fromList = fromListWith GHC.Exts.fromList -- | /O(n)/ Inverts keys and values inside a multimap. inverse :: (Collection c, GHC.Exts.IsList (c k), GHC.Exts.Item (c k) ~ k, Ord k, Ord v) => Multimap c k v -> Multimap c v k inverse = inverseWith GHC.Exts.fromList #endif -- | /O(1)/ Checks whether the multimap is empty. null :: Multimap c k v -> Bool null = Map.null . _toMap -- | /O(m * C)/ Returns the size of the multimap. size :: Collection c => Multimap c k v -> Int size (Multimap m) = Map.foldl' (\n c -> n + Col.size c) 0 m -- | /O(1)/ Returns the number of distinct keys in the multimap. distinctSize :: Multimap c k v -> Int distinctSize = Map.size . _toMap -- | /O(1)/ Creates an empty multimap. empty :: Multimap c k v empty = Multimap Map.empty -- | /O(1)/ Creates a multimap with a single entry. singleton :: Collection c => k -> v -> Multimap c k v singleton k v = Multimap $ Map.singleton k (Col.singleton v) -- | /O(m)/ Builds a multimap from already grouped collections. fromGroupList :: (Collection c, Monoid (c v), Ord k) => [Group k (c v)] -> Multimap c k v fromGroupList = fromMap . Map.fromListWith (<>) -- | /O(n * log n)/ Transforms a list of entries into a multimap, combining the values for each key -- into the chosen collection. The values are in the same order as in the original list. fromListWith :: Ord k => ([v] -> c v) -> [(k, v)] -> Multimap c k v fromListWith f ts = Multimap $ Map.map (f . reverse) $ m where Multimap m = foldl' (\r (k, v) -> modifyMany (v:) k r) empty ts -- | /O(m * C)/ Transforms a map of collections into a multimap. fromMap :: Collection c => Map k (c v) -> Multimap c k v fromMap = Multimap . Map.filter (not . Col.null) -- | /O(log m)/ Returns the collection of values associated with a key. find :: (Monoid (c v), Ord k) => k -> Multimap c k v -> c v find k (Multimap m) = Map.findWithDefault mempty k m -- | /O(log m)/ Infix version of 'find'. (!) :: (Monoid (c v), Ord k) => Multimap c k v -> k -> c v (!) = flip find -- | /O(log m * C)/ Returns the number of times a key is present in a multimap. The complexity of -- the operation depends on the complexity of the underlying collection's 'Data.Collection.size' -- operation. count :: (Collection c, Ord k) => k -> Multimap c k v -> Int count k (Multimap m) = maybe 0 Col.size $ Map.lookup k m -- | /O(log m)/ Checks whether a key is present at least once in a multimap. member :: Ord k => k -> Multimap c k v -> Bool member k = Map.member k . _toMap -- | /O(log m)/ Checks whether a key is absent from a multimap. notMember :: Ord k => k -> Multimap c k v -> Bool notMember k = Map.notMember k . _toMap -- | Modifies a key's collection using an arbitrary function. More specifically, this function lifts -- an operation over a collection of values into a multimap operation. -- -- Sample use to filter even values from a 'Data.Multimap.SetMultimap': -- -- @ -- let ms = fromList [(\'a\', 1), (\'a\', 2)] :: SetMultimap Char Int -- modifyMany (Set.filter even) \'a\' ms == fromList [(\'a\', 1)] -- @ modifyMany :: (Collection c, Monoid (c v), Ord k) => (c v -> c v) -> k -> Multimap c k v -> Multimap c k v modifyMany f k (Multimap m) = Multimap $ Map.alter (wrap . f . unwrap) k m where unwrap = fromMaybe mempty wrap c = if Col.null c then Nothing else Just c -- | Modifies a key's collection using an arbitrary function. This is the applicative version of -- 'modifyMany'. modifyManyF :: (Collection c, Monoid (c v), Ord k, Functor f) => (c v -> f (c v)) -> k -> Multimap c k v -> f (Multimap c k v) modifyManyF f k (Multimap m) = Multimap <$> Map.alterF (fmap wrap . f . unwrap) k m where unwrap = fromMaybe mempty wrap c = if Col.null c then Nothing else Just c -- | /O(log m * C)/ Prepends a value to a key's collection. prepend :: (Collection c, Monoid (c v), Ord k) => k -> v -> Multimap c k v -> Multimap c k v prepend k v = prependMany k (Col.singleton v) -- | /O(log m * C)/ Prepends a collection of values to a key's collection. prependMany :: (Collection c, Monoid (c v), Ord k) => k -> c v -> Multimap c k v -> Multimap c k v prependMany k c = modifyMany (c <>) k -- | /O(log m * C)/ Appends a value to a key's collection. append :: (Collection c, Monoid (c v), Ord k) => k -> v -> Multimap c k v -> Multimap c k v append k v = appendMany k (Col.singleton v) -- | /O(log m * C)/ Appends a collection of values to a key's collection. appendMany :: (Collection c, Monoid (c v), Ord k) => k -> c v -> Multimap c k v -> Multimap c k v appendMany k c = modifyMany (<> c) k -- | /O(log m)/ Removes all entries for the given key. deleteMany :: Ord k => k -> Multimap c k v -> Multimap c k v deleteMany k = Multimap . Map.delete k . _toMap -- | /O(n)/ Inverts keys and values inside a multimap, potentially changing the collection type. inverseWith :: (Collection c1, Ord k, Ord v) => ([k] -> c2 k) -> Multimap c1 k v -> Multimap c2 v k inverseWith f = fromListWith f . fmap swap . toList -- | /O(n)/ Filters multimap entries by value. filter :: (Collection c, Monoid (c v), Ord k) => (v -> Bool) -> Multimap c k v -> Multimap c k v filter f = fromGroupList . fmap (fmap (Col.filter f)) . toGroupList -- | /O(m)/ Filters multimap groups. This enables filtering by key and collection. filterGroups :: (Collection c, Monoid (c v), Ord k) => (Group k (c v) -> Bool) -> Multimap c k v -> Multimap c k v filterGroups f = fromGroupList . Prelude.filter f . toGroupList -- | Maps over the multimap's groups. This method can be used to convert between specific multimaps, -- for example: -- -- > let m1 = fromList [('a', 1), ('a', 1)] :: ListMultimap Char Int -- > let m2 = mapGroups (fmap Set.fromList) m1 :: SetMultimap Char Int mapGroups :: (Collection c2, Monoid (c2 v2), Ord k2) => (Group k1 (c1 v1) -> Group k2 (c2 v2)) -> Multimap c1 k1 v1 -> Multimap c2 k2 v2 mapGroups f = fromGroupList . fmap f . toGroupList -- | /O(n)/ Converts a multimap into its list of entries. Note that this is different from -- 'Data.Foldable.toList' which returns the multimap's values (similar to "Data.Map"). toList :: Collection c => Multimap c k v -> [(k, v)] toList (Multimap m) = concat $ fmap go $ Map.toList m where go (k,c) = foldr (\v a -> (k,v) : a) [] c -- | /O(m)/ Converts a multimap into its list of collections. toGroupList :: Multimap c k v -> [Group k (c v)] toGroupList = Map.toList . _toMap -- | /O(1)/ Converts a multimap into a map of collections. toMap :: Multimap c k v -> Map k (c v) toMap = _toMap viewWith :: (Collection c, Ord k) => (Map k (c v) -> Maybe ((k, c v), Map k (c v))) -> (c v -> Maybe (v, c v)) -> Multimap c k v -> Maybe ((k, v), Multimap c k v) viewWith mapView f (Multimap m) = case mapView m of Nothing -> Nothing Just ((k, c), m') -> case f c of Nothing -> Nothing Just (v, c') -> Just ((k, v), Multimap (if Col.null c' then m' else Map.insert k c' m')) -- | Returns the multimap's maximum key and value. The input function is guaranteed to be called -- with a non-empty collection. -- -- @since 0.2.1.2 maxViewWith :: (Collection c, Ord k) => (c v -> Maybe (v, c v)) -> Multimap c k v -> Maybe ((k, v), Multimap c k v) maxViewWith = viewWith Map.maxViewWithKey -- | Returns the multimap's minimum key and value. The input function is guaranteed to be called -- with a non-empty collection. -- -- @since 0.2.1.2 minViewWith :: (Collection c, Ord k) => (c v -> Maybe (v, c v)) -> Multimap c k v -> Maybe ((k, v), Multimap c k v) minViewWith = viewWith Map.minViewWithKey -- | /O(m)/ Returns a list of the multimap's keys. Each key will be repeated as many times as it is -- present in the multimap. keys :: Collection c => Multimap c k v -> [k] keys (Multimap m) = Map.foldrWithKey go [] m where go k c r = replicate (Col.size c) k <> r -- | /O(m)/ Returns a set of the multimap's (distinct) keys. keysSet :: Multimap c k v -> Set k keysSet = Map.keysSet . _toMap -- | /O(m * C)/ Returns a multiset of the map's keys with matching multiplicities. keysMultiset :: (Collection c, Ord k) => Multimap c k v -> Multiset k keysMultiset = Mset.fromCountMap . Map.map Col.size . _toMap