module IntMultimap
(
IntMultimap,
empty,
singleton,
toList,
fromList,
map,
foldlWithKey',
null,
size,
member,
delete,
insert,
elems,
keys,
split,
splitLookup
)
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 Data.Maybe
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 a -> [(Int, a)]
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, Eq a) => Int -> a -> IntMultimap a -> IntMultimap a
insert key value (IntMultimap intMap) =
IntMultimap $ A.insertWith (f value) key (B.singleton value) intMap
where
f v new old = B.insert v old
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
foldlWithKey' :: (Int -> a -> b -> a) -> a -> IntMultimap b -> a
foldlWithKey' f v (IntMultimap intMap) = A.foldlWithKey' (\a k set -> B.foldl' (f k) a set) v intMap
split :: Int -> IntMultimap a -> (IntMultimap a, IntMultimap a)
split key (IntMultimap intMap) = (IntMultimap oldMap, IntMultimap newMap)
where
(oldMap, newMap) = A.split key intMap
splitLookup :: Int -> IntMultimap a -> (IntMultimap a, Maybe (B.HashSet a), IntMultimap a)
splitLookup key (IntMultimap intMap) = (IntMultimap oldMap, elemHashSet, IntMultimap newMap)
where
(oldMap, elemHashSet, newMap) = A.splitLookup key intMap