Commit 09485bba authored by Edward Z. Yang's avatar Edward Z. Yang
Browse files

UniqMap implementation.


Signed-off-by: default avatarEdward Z. Yang <ezyang@cs.stanford.edu>
parent 713ff920
......@@ -508,6 +508,7 @@ Library
UniqDFM
UniqDSet
UniqFM
UniqMap
UniqSet
Util
Vectorise.Builtins.Base
......
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# OPTIONS_GHC -Wall #-}
-- Like 'UniqFM', these are maps for keys which are Uniquable.
-- Unlike 'UniqFM', these maps also remember their keys, which
-- makes them a much better drop in replacement for 'Data.Map.Map'.
--
-- Key preservation is right-biased.
module UniqMap (
UniqMap,
emptyUniqMap,
isNullUniqMap,
unitUniqMap,
listToUniqMap,
listToUniqMap_C,
addToUniqMap,
addListToUniqMap,
addToUniqMap_C,
addToUniqMap_Acc,
alterUniqMap,
addListToUniqMap_C,
adjustUniqMap,
delFromUniqMap,
delListFromUniqMap,
plusUniqMap,
plusUniqMap_C,
plusMaybeUniqMap_C,
plusUniqMapList,
minusUniqMap,
intersectUniqMap,
disjointUniqMap,
mapUniqMap,
filterUniqMap,
partitionUniqMap,
sizeUniqMap,
elemUniqMap,
lookupUniqMap,
lookupWithDefaultUniqMap,
anyUniqMap,
allUniqMap,
-- Non-deterministic functions omitted
) where
import UniqFM
import Unique
import Outputable
#if __GLASGOW_HASKELL__ > 710
import Data.Semigroup ( Semigroup(..) )
#endif
import Data.Coerce
import Data.Maybe
import Data.Typeable
import Data.Data
-- | Maps indexed by 'Uniquable' keys
newtype UniqMap k a = UniqMap (UniqFM (k, a))
deriving (Data, Eq, Functor, Typeable)
type role UniqMap nominal representational
#if __GLASGOW_HASKELL__ > 710
instance Semigroup (UniqMap k a) where
(<>) = plusUniqMap
#endif
instance Monoid (UniqMap k a) where
mempty = emptyUniqMap
mappend = plusUniqMap
instance (Outputable k, Outputable a) => Outputable (UniqMap k a) where
ppr (UniqMap m) =
brackets $ fsep $ punctuate comma $
[ ppr k <+> text "->" <+> ppr v
| (k, v) <- eltsUFM m ]
liftC :: (a -> a -> a) -> (k, a) -> (k, a) -> (k, a)
liftC f (_, v) (k', v') = (k', f v v')
emptyUniqMap :: UniqMap k a
emptyUniqMap = UniqMap emptyUFM
isNullUniqMap :: UniqMap k a -> Bool
isNullUniqMap (UniqMap m) = isNullUFM m
unitUniqMap :: Uniquable k => k -> a -> UniqMap k a
unitUniqMap k v = UniqMap (unitUFM k (k, v))
listToUniqMap :: Uniquable k => [(k,a)] -> UniqMap k a
listToUniqMap kvs = UniqMap (listToUFM [ (k,(k,v)) | (k,v) <- kvs])
listToUniqMap_C :: Uniquable k => (a -> a -> a) -> [(k,a)] -> UniqMap k a
listToUniqMap_C f kvs = UniqMap $
listToUFM_C (liftC f) [ (k,(k,v)) | (k,v) <- kvs]
addToUniqMap :: Uniquable k => UniqMap k a -> k -> a -> UniqMap k a
addToUniqMap (UniqMap m) k v = UniqMap $ addToUFM m k (k, v)
addListToUniqMap :: Uniquable k => UniqMap k a -> [(k,a)] -> UniqMap k a
addListToUniqMap (UniqMap m) kvs = UniqMap $
addListToUFM m [(k,(k,v)) | (k,v) <- kvs]
addToUniqMap_C :: Uniquable k
=> (a -> a -> a)
-> UniqMap k a
-> k
-> a
-> UniqMap k a
addToUniqMap_C f (UniqMap m) k v = UniqMap $
addToUFM_C (liftC f) m k (k, v)
addToUniqMap_Acc :: Uniquable k
=> (b -> a -> a)
-> (b -> a)
-> UniqMap k a
-> k
-> b
-> UniqMap k a
addToUniqMap_Acc exi new (UniqMap m) k0 v0 = UniqMap $
addToUFM_Acc (\b (k, v) -> (k, exi b v))
(\b -> (k0, new b))
m k0 v0
alterUniqMap :: Uniquable k
=> (Maybe a -> Maybe a)
-> UniqMap k a
-> k
-> UniqMap k a
alterUniqMap f (UniqMap m) k = UniqMap $
alterUFM (fmap (k,) . f . fmap snd) m k
addListToUniqMap_C
:: Uniquable k
=> (a -> a -> a)
-> UniqMap k a
-> [(k, a)]
-> UniqMap k a
addListToUniqMap_C f (UniqMap m) kvs = UniqMap $
addListToUFM_C (liftC f) m
[(k,(k,v)) | (k,v) <- kvs]
adjustUniqMap
:: Uniquable k
=> (a -> a)
-> UniqMap k a
-> k
-> UniqMap k a
adjustUniqMap f (UniqMap m) k = UniqMap $
adjustUFM (\(_,v) -> (k,f v)) m k
delFromUniqMap :: Uniquable k => UniqMap k a -> k -> UniqMap k a
delFromUniqMap (UniqMap m) k = UniqMap $ delFromUFM m k
delListFromUniqMap :: Uniquable k => UniqMap k a -> [k] -> UniqMap k a
delListFromUniqMap (UniqMap m) ks = UniqMap $ delListFromUFM m ks
plusUniqMap :: UniqMap k a -> UniqMap k a -> UniqMap k a
plusUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ plusUFM m1 m2
plusUniqMap_C :: (a -> a -> a) -> UniqMap k a -> UniqMap k a -> UniqMap k a
plusUniqMap_C f (UniqMap m1) (UniqMap m2) = UniqMap $
plusUFM_C (liftC f) m1 m2
plusMaybeUniqMap_C :: (a -> a -> Maybe a) -> UniqMap k a -> UniqMap k a -> UniqMap k a
plusMaybeUniqMap_C f (UniqMap m1) (UniqMap m2) = UniqMap $
plusMaybeUFM_C (\(_, v) (k', v') -> fmap (k',) (f v v')) m1 m2
plusUniqMapList :: [UniqMap k a] -> UniqMap k a
plusUniqMapList xs = UniqMap $ plusUFMList (coerce xs)
minusUniqMap :: UniqMap k a -> UniqMap k b -> UniqMap k a
minusUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ minusUFM m1 m2
intersectUniqMap :: UniqMap k a -> UniqMap k b -> UniqMap k a
intersectUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ intersectUFM m1 m2
disjointUniqMap :: UniqMap k a -> UniqMap k b -> Bool
disjointUniqMap (UniqMap m1) (UniqMap m2) = disjointUFM m1 m2
mapUniqMap :: (a -> b) -> UniqMap k a -> UniqMap k b
mapUniqMap f (UniqMap m) = UniqMap $ mapUFM (fmap f) m -- (,) k instance
filterUniqMap :: (a -> Bool) -> UniqMap k a -> UniqMap k a
filterUniqMap f (UniqMap m) = UniqMap $ filterUFM (f . snd) m
partitionUniqMap :: (a -> Bool) -> UniqMap k a -> (UniqMap k a, UniqMap k a)
partitionUniqMap f (UniqMap m) =
coerce $ partitionUFM (f . snd) m
sizeUniqMap :: UniqMap k a -> Int
sizeUniqMap (UniqMap m) = sizeUFM m
elemUniqMap :: Uniquable k => k -> UniqMap k a -> Bool
elemUniqMap k (UniqMap m) = elemUFM k m
lookupUniqMap :: Uniquable k => UniqMap k a -> k -> Maybe a
lookupUniqMap (UniqMap m) k = fmap snd (lookupUFM m k)
lookupWithDefaultUniqMap :: Uniquable k => UniqMap k a -> a -> k -> a
lookupWithDefaultUniqMap (UniqMap m) a k = fromMaybe a (fmap snd (lookupUFM m k))
anyUniqMap :: (a -> Bool) -> UniqMap k a -> Bool
anyUniqMap f (UniqMap m) = anyUFM (f . snd) m
allUniqMap :: (a -> Bool) -> UniqMap k a -> Bool
allUniqMap f (UniqMap m) = allUFM (f . snd) m
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment