module IntMultimap
(
IntMultiMap,
empty,
singleton,
toList,
fromList,
map,
null,
size,
member,
delete,
insert,
elems,
keys,
split
)
where
import qualified Data.IntMap.Strict as A
import qualified Data.HashSet as B
import qualified Data.Foldable as C
import qualified GHC.Exts as G
import GHC.Generics
import Prelude hiding (map, null)
import Data.Int
import Data.Hashable
import Data.Functor
import Data.List(length)
import Control.Monad
newtype IntMultiMap a =
IntMultiMap (A.IntMap (B.HashSet a))
deriving(Foldable, Eq, Show, Generic)
map :: (Eq b, Hashable b) => (a -> b) -> IntMultiMap a -> IntMultiMap b
map f (IntMultiMap intMap) = IntMultiMap $
fmap (\hashSet -> B.map f hashSet) intMap
instance (Eq a, Hashable a) => G.IsList (IntMultiMap a) where
type Item (IntMultiMap a) = (Int, a)
toList = toList
fromList = fromList
toList :: IntMultiMap b -> [(Int, b)]
toList (IntMultiMap multiMap) = do
(key, hashSet) <- A.toList multiMap
fmap ((,) key) $ B.toList hashSet
fromList :: (Eq a, Hashable a) =>
[(Int, a)] -> IntMultiMap a
fromList = IntMultiMap . A.fromListWith B.union . fmap (fmap B.singleton)
empty :: IntMultiMap a
empty = IntMultiMap A.empty
singleton :: (Hashable a) => Int -> a -> IntMultiMap a
singleton k v = IntMultiMap $ A.singleton k $ B.singleton v
null :: IntMultiMap a -> Bool
null (IntMultiMap intMap) = A.null intMap
size :: IntMultiMap a -> Int
size = length . toList
member :: Int -> IntMultiMap a -> Bool
member key (IntMultiMap intMap) = A.member key intMap
insert :: (Hashable a, Ord a) => Int -> a -> IntMultiMap a -> IntMultiMap a
insert key value (IntMultiMap intMap) =
IntMultiMap $ A.update (\hash -> Just $ B.insert value hash) key intMap
delete :: (Hashable a, Eq a) => Int -> a -> IntMultiMap a -> IntMultiMap a
delete key value (IntMultiMap intMap) =
IntMultiMap $ A.update f key intMap
where
f hashSet =
mfilter (not . B.null) . Just $ B.delete value hashSet
elems :: IntMultiMap a -> [a]
elems = foldr (:) []
keys :: IntMultiMap a -> [Int]
keys (IntMultiMap intMap) = A.keys intMap
split :: Int -> IntMultiMap a -> (IntMultiMap a, IntMultiMap a)
split key (IntMultiMap intMap) = (IntMultiMap oldMap, IntMultiMap newMap)
where
(oldMap, newMap) = A.split key intMap