module UniqDFM (
UniqDFM,
emptyUDFM,
unitUDFM,
addToUDFM,
addToUDFM_C,
addListToUDFM,
delFromUDFM,
delListFromUDFM,
adjustUDFM,
alterUDFM,
mapUDFM,
plusUDFM,
plusUDFM_C,
lookupUDFM, lookupUDFM_Directly,
elemUDFM,
foldUDFM,
eltsUDFM,
filterUDFM, filterUDFM_Directly,
isNullUDFM,
sizeUDFM,
intersectUDFM, udfmIntersectUFM,
intersectsUDFM,
disjointUDFM, disjointUdfmUfm,
minusUDFM,
listToUDFM,
udfmMinusUFM,
partitionUDFM,
anyUDFM, allUDFM,
pprUDFM,
udfmToList,
udfmToUfm,
nonDetFoldUDFM,
alwaysUnsafeUfmToUdfm,
) where
import GhcPrelude
import Unique ( Uniquable(..), Unique, getKey )
import Outputable
import qualified Data.IntMap as M
import Data.Data
import Data.List (sortBy)
import Data.Function (on)
import qualified Data.Semigroup as Semi
import UniqFM (UniqFM, listToUFM_Directly, nonDetUFMToList, ufmToIntMap)
data TaggedVal val =
TaggedVal
val
!Int
deriving Data
taggedFst :: TaggedVal val -> val
taggedFst (TaggedVal v _) = v
taggedSnd :: TaggedVal val -> Int
taggedSnd (TaggedVal _ i) = i
instance Eq val => Eq (TaggedVal val) where
(TaggedVal v1 _) == (TaggedVal v2 _) = v1 == v2
instance Functor TaggedVal where
fmap f (TaggedVal val i) = TaggedVal (f val) i
data UniqDFM ele =
UDFM
!(M.IntMap (TaggedVal ele))
!Int
deriving (Data, Functor)
emptyUDFM :: UniqDFM elt
emptyUDFM = UDFM M.empty 0
unitUDFM :: Uniquable key => key -> elt -> UniqDFM elt
unitUDFM k v = UDFM (M.singleton (getKey $ getUnique k) (TaggedVal v 0)) 1
addToUDFM :: Uniquable key => UniqDFM elt -> key -> elt -> UniqDFM elt
addToUDFM m k v = addToUDFM_Directly m (getUnique k) v
addToUDFM_Directly :: UniqDFM elt -> Unique -> elt -> UniqDFM elt
addToUDFM_Directly (UDFM m i) u v
= UDFM (M.insertWith tf (getKey u) (TaggedVal v i) m) (i + 1)
where
tf (TaggedVal new_v _) (TaggedVal _ old_i) = TaggedVal new_v old_i
addToUDFM_Directly_C
:: (elt -> elt -> elt)
-> UniqDFM elt
-> Unique -> elt
-> UniqDFM elt
addToUDFM_Directly_C f (UDFM m i) u v
= UDFM (M.insertWith tf (getKey u) (TaggedVal v i) m) (i + 1)
where
tf (TaggedVal new_v _) (TaggedVal old_v old_i)
= TaggedVal (f old_v new_v) old_i
addToUDFM_C
:: Uniquable key => (elt -> elt -> elt)
-> UniqDFM elt
-> key -> elt
-> UniqDFM elt
addToUDFM_C f m k v = addToUDFM_Directly_C f m (getUnique k) v
addListToUDFM :: Uniquable key => UniqDFM elt -> [(key,elt)] -> UniqDFM elt
addListToUDFM = foldl (\m (k, v) -> addToUDFM m k v)
addListToUDFM_Directly :: UniqDFM elt -> [(Unique,elt)] -> UniqDFM elt
addListToUDFM_Directly = foldl (\m (k, v) -> addToUDFM_Directly m k v)
addListToUDFM_Directly_C
:: (elt -> elt -> elt) -> UniqDFM elt -> [(Unique,elt)] -> UniqDFM elt
addListToUDFM_Directly_C f = foldl (\m (k, v) -> addToUDFM_Directly_C f m k v)
delFromUDFM :: Uniquable key => UniqDFM elt -> key -> UniqDFM elt
delFromUDFM (UDFM m i) k = UDFM (M.delete (getKey $ getUnique k) m) i
plusUDFM_C :: (elt -> elt -> elt) -> UniqDFM elt -> UniqDFM elt -> UniqDFM elt
plusUDFM_C f udfml@(UDFM _ i) udfmr@(UDFM _ j)
| i > j = insertUDFMIntoLeft_C f udfml udfmr
| otherwise = insertUDFMIntoLeft_C f udfmr udfml
plusUDFM :: UniqDFM elt -> UniqDFM elt -> UniqDFM elt
plusUDFM udfml@(UDFM _ i) udfmr@(UDFM _ j)
| i > j = insertUDFMIntoLeft udfml udfmr
| otherwise = insertUDFMIntoLeft udfmr udfml
insertUDFMIntoLeft :: UniqDFM elt -> UniqDFM elt -> UniqDFM elt
insertUDFMIntoLeft udfml udfmr = addListToUDFM_Directly udfml $ udfmToList udfmr
insertUDFMIntoLeft_C
:: (elt -> elt -> elt) -> UniqDFM elt -> UniqDFM elt -> UniqDFM elt
insertUDFMIntoLeft_C f udfml udfmr =
addListToUDFM_Directly_C f udfml $ udfmToList udfmr
lookupUDFM :: Uniquable key => UniqDFM elt -> key -> Maybe elt
lookupUDFM (UDFM m _i) k = taggedFst `fmap` M.lookup (getKey $ getUnique k) m
lookupUDFM_Directly :: UniqDFM elt -> Unique -> Maybe elt
lookupUDFM_Directly (UDFM m _i) k = taggedFst `fmap` M.lookup (getKey k) m
elemUDFM :: Uniquable key => key -> UniqDFM elt -> Bool
elemUDFM k (UDFM m _i) = M.member (getKey $ getUnique k) m
foldUDFM :: (elt -> a -> a) -> a -> UniqDFM elt -> a
foldUDFM k z m = foldr k z (eltsUDFM m)
nonDetFoldUDFM :: (elt -> a -> a) -> a -> UniqDFM elt -> a
nonDetFoldUDFM k z (UDFM m _i) = foldr k z $ map taggedFst $ M.elems m
eltsUDFM :: UniqDFM elt -> [elt]
eltsUDFM (UDFM m _i) =
map taggedFst $ sortBy (compare `on` taggedSnd) $ M.elems m
filterUDFM :: (elt -> Bool) -> UniqDFM elt -> UniqDFM elt
filterUDFM p (UDFM m i) = UDFM (M.filter (\(TaggedVal v _) -> p v) m) i
filterUDFM_Directly :: (Unique -> elt -> Bool) -> UniqDFM elt -> UniqDFM elt
filterUDFM_Directly p (UDFM m i) = UDFM (M.filterWithKey p' m) i
where
p' k (TaggedVal v _) = p (getUnique k) v
udfmToList :: UniqDFM elt -> [(Unique, elt)]
udfmToList (UDFM m _i) =
[ (getUnique k, taggedFst v)
| (k, v) <- sortBy (compare `on` (taggedSnd . snd)) $ M.toList m ]
isNullUDFM :: UniqDFM elt -> Bool
isNullUDFM (UDFM m _) = M.null m
sizeUDFM :: UniqDFM elt -> Int
sizeUDFM (UDFM m _i) = M.size m
intersectUDFM :: UniqDFM elt -> UniqDFM elt -> UniqDFM elt
intersectUDFM (UDFM x i) (UDFM y _j) = UDFM (M.intersection x y) i
udfmIntersectUFM :: UniqDFM elt1 -> UniqFM elt2 -> UniqDFM elt1
udfmIntersectUFM (UDFM x i) y = UDFM (M.intersection x (ufmToIntMap y)) i
intersectsUDFM :: UniqDFM elt -> UniqDFM elt -> Bool
intersectsUDFM x y = isNullUDFM (x `intersectUDFM` y)
disjointUDFM :: UniqDFM elt -> UniqDFM elt -> Bool
disjointUDFM (UDFM x _i) (UDFM y _j) = M.null (M.intersection x y)
disjointUdfmUfm :: UniqDFM elt -> UniqFM elt2 -> Bool
disjointUdfmUfm (UDFM x _i) y = M.null (M.intersection x (ufmToIntMap y))
minusUDFM :: UniqDFM elt1 -> UniqDFM elt2 -> UniqDFM elt1
minusUDFM (UDFM x i) (UDFM y _j) = UDFM (M.difference x y) i
udfmMinusUFM :: UniqDFM elt1 -> UniqFM elt2 -> UniqDFM elt1
udfmMinusUFM (UDFM x i) y = UDFM (M.difference x (ufmToIntMap y)) i
partitionUDFM :: (elt -> Bool) -> UniqDFM elt -> (UniqDFM elt, UniqDFM elt)
partitionUDFM p (UDFM m i) =
case M.partition (p . taggedFst) m of
(left, right) -> (UDFM left i, UDFM right i)
delListFromUDFM :: Uniquable key => UniqDFM elt -> [key] -> UniqDFM elt
delListFromUDFM = foldl delFromUDFM
udfmToUfm :: UniqDFM elt -> UniqFM elt
udfmToUfm (UDFM m _i) =
listToUFM_Directly [(getUnique k, taggedFst tv) | (k, tv) <- M.toList m]
listToUDFM :: Uniquable key => [(key,elt)] -> UniqDFM elt
listToUDFM = foldl (\m (k, v) -> addToUDFM m k v) emptyUDFM
listToUDFM_Directly :: [(Unique, elt)] -> UniqDFM elt
listToUDFM_Directly = foldl (\m (u, v) -> addToUDFM_Directly m u v) emptyUDFM
adjustUDFM :: Uniquable key => (elt -> elt) -> UniqDFM elt -> key -> UniqDFM elt
adjustUDFM f (UDFM m i) k = UDFM (M.adjust (fmap f) (getKey $ getUnique k) m) i
alterUDFM
:: Uniquable key
=> (Maybe elt -> Maybe elt)
-> UniqDFM elt
-> key
-> UniqDFM elt
alterUDFM f (UDFM m i) k =
UDFM (M.alter alterf (getKey $ getUnique k) m) (i + 1)
where
alterf Nothing = inject $ f Nothing
alterf (Just (TaggedVal v _)) = inject $ f (Just v)
inject Nothing = Nothing
inject (Just v) = Just $ TaggedVal v i
mapUDFM :: (elt1 -> elt2) -> UniqDFM elt1 -> UniqDFM elt2
mapUDFM f (UDFM m i) = UDFM (M.map (fmap f) m) i
anyUDFM :: (elt -> Bool) -> UniqDFM elt -> Bool
anyUDFM p (UDFM m _i) = M.foldr ((||) . p . taggedFst) False m
allUDFM :: (elt -> Bool) -> UniqDFM elt -> Bool
allUDFM p (UDFM m _i) = M.foldr ((&&) . p . taggedFst) True m
instance Semi.Semigroup (UniqDFM a) where
(<>) = plusUDFM
instance Monoid (UniqDFM a) where
mempty = emptyUDFM
mappend = (Semi.<>)
alwaysUnsafeUfmToUdfm :: UniqFM elt -> UniqDFM elt
alwaysUnsafeUfmToUdfm = listToUDFM_Directly . nonDetUFMToList
instance Outputable a => Outputable (UniqDFM a) where
ppr ufm = pprUniqDFM ppr ufm
pprUniqDFM :: (a -> SDoc) -> UniqDFM a -> SDoc
pprUniqDFM ppr_elt ufm
= brackets $ fsep $ punctuate comma $
[ ppr uq <+> text ":->" <+> ppr_elt elt
| (uq, elt) <- udfmToList ufm ]
pprUDFM :: UniqDFM a
-> ([a] -> SDoc)
-> SDoc
pprUDFM ufm pp = pp (eltsUDFM ufm)