{-# LANGUAGE DeriveFunctor #-}
module Data.Multimap.Internal ( Multimap(..)
, null, size
, empty, singleton, fromMap
, fromList, fromCollectionsList
, member, notMember
, (!)
, insert, insertAll, deleteAll
, filter, filterWithKey
, union
, toMap, toMapWith, toList
, keys, keysSet, keysMultiset
, lift1, liftF1
) 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.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)
newtype Multimap k c v = Multimap { getMultimap :: Map k (c v)
} deriving (Eq, Functor, Ord, Read, Show)
instance (Ord k, Semigroup (c v)) => Semigroup (Multimap k c v) where
(<>) = union
instance (Ord k, Semigroup (c v)) => Monoid (Multimap k c v) where
mempty = empty
mappend = union
instance (Collection c) => Foldable (Multimap k c) where
foldMap f (Multimap m) = Map.foldMapWithKey go m where
go _ v = foldMap f v
null :: Multimap k c v -> Bool
null = Map.null . getMultimap
size :: (Collection c) => Multimap k c v -> Int
size = Map.foldl (\a c -> a + Col.size c) 0 . getMultimap
empty :: Multimap k c v
empty = Multimap Map.empty
singleton :: (Collection c, Ord k) => k -> v -> Multimap k c v
singleton k v = Multimap $ Map.singleton k (Col.singleton v)
fromMap :: (Collection c) => Map k (c v) -> Multimap k c v
fromMap m = Multimap $ Map.filter (not . Col.null) m
fromList :: (Collection c, Semigroup (c v), Ord k) => [(k,v)] -> Multimap k c v
fromList ts = Multimap $ Map.fromListWith (<>) (fmap go ts) where
go (k,v) = (k, Col.singleton v)
fromCollectionsList :: (Collection c, Semigroup (c v), Ord k) => [(k, c v)] -> Multimap k c v
fromCollectionsList = fromMap . Map.fromList
(!) :: (Collection c, Ord k) => Multimap k c v -> k -> c v
(Multimap m) ! k = Map.findWithDefault Col.empty k m
member :: (Collection c, Ord k) => k -> Multimap k c v -> Bool
member k = Map.member k . getMultimap
notMember :: (Collection c, Ord k) => k -> Multimap k c v -> Bool
notMember k = Map.notMember k . getMultimap
count :: (Collection c, Ord k) => k -> Multimap k c v -> Int
count k = Col.size . (! k)
insertAll :: (Ord k, Semigroup (c v)) => k -> c v -> Multimap k c v -> Multimap k c v
insertAll k c (Multimap m) = Multimap $ Map.insertWith (<>) k c m
insert :: (Ord k, Collection c, Semigroup (c v)) => k -> v -> Multimap k c v -> Multimap k c v
insert k v = insertAll k (Col.singleton v)
deleteAll :: (Ord k) => k -> Multimap k c v -> Multimap k c v
deleteAll k (Multimap m) = Multimap $ Map.delete k m
filter :: (Collection c, Semigroup (c v), Ord k) => (v -> Bool) -> Multimap k c v -> Multimap k c v
filter f = filterWithKey (const f)
filterWithKey :: (Collection c, Semigroup (c v), Ord k) => (k -> v -> Bool) -> Multimap k c v -> Multimap k c v
filterWithKey f = fromList . Prelude.filter (uncurry f) . toList
union :: (Ord k, Semigroup (c v)) => Multimap k c v -> Multimap k c v -> Multimap k c v
union (Multimap m1) (Multimap m2) = Multimap $ Map.unionWith (<>) m1 m2
toMap :: Multimap k c v -> Map k (c v)
toMap = getMultimap
toMapWith :: (c v -> a) -> Multimap k c v -> Map k a
toMapWith f = fmap f . getMultimap
toList :: (Collection c) => Multimap k c v -> [(k,v)]
toList (Multimap m) = concat $ fmap go (Map.toList m) where
go (k,cs) = foldr (\c a -> (k,c) : a) [] cs
keys :: Multimap k c v -> [k]
keys = Map.keys . getMultimap
keysSet :: Multimap k c v -> Set k
keysSet = Map.keysSet . getMultimap
keysMultiset :: (Ord k, Collection c) => Multimap k c v -> Multiset k
keysMultiset = Mset.fromCountsList . Map.toList . Map.map Col.size . getMultimap
lift1 :: (Ord k, Collection c) => (c v -> c v)
-> k
-> Multimap k c v
-> Multimap k c v
lift1 f k (Multimap m) = Multimap $ Map.alter (wrap . f . unwrap) k m where
unwrap = fromMaybe Col.empty
wrap cs = if Col.null cs then Nothing else Just cs
liftF1 :: (Ord k, Collection c, Functor f) => (c v -> f (c v))
-> k
-> Multimap k c v
-> f (Multimap k c v)
liftF1 f k (Multimap m) = Multimap <$> Map.alterF (fmap wrap . f . unwrap) k m where
unwrap = fromMaybe Col.empty
wrap cs = if Col.null cs then Nothing else Just cs