1{- 2(c) The University of Glasgow 2006 3(c) The AQUA Project, Glasgow University, 1994-1998 4 5 6UniqFM: Specialised finite maps, for things with @Uniques@. 7 8Basically, the things need to be in class @Uniquable@, and we use the 9@getUnique@ method to grab their @Uniques@. 10 11(A similar thing to @UniqSet@, as opposed to @Set@.) 12 13The interface is based on @FiniteMap@s, but the implementation uses 14@Data.IntMap@, which is both maintained and faster than the past 15implementation (see commit log). 16 17The @UniqFM@ interface maps directly to Data.IntMap, only 18``Data.IntMap.union'' is left-biased and ``plusUFM'' right-biased 19and ``addToUFM\_C'' and ``Data.IntMap.insertWith'' differ in the order 20of arguments of combining function. 21-} 22 23{-# LANGUAGE DeriveDataTypeable #-} 24{-# LANGUAGE GeneralizedNewtypeDeriving #-} 25{-# LANGUAGE ScopedTypeVariables #-} 26{-# OPTIONS_GHC -Wall #-} 27 28module GHC.Types.Unique.FM ( 29 -- * Unique-keyed mappings 30 UniqFM, -- abstract type 31 NonDetUniqFM(..), -- wrapper for opting into nondeterminism 32 33 -- ** Manipulating those mappings 34 emptyUFM, 35 unitUFM, 36 unitDirectlyUFM, 37 listToUFM, 38 listToUFM_Directly, 39 listToUFM_C, 40 listToIdentityUFM, 41 addToUFM,addToUFM_C,addToUFM_Acc, 42 addListToUFM,addListToUFM_C, 43 addToUFM_Directly, 44 addListToUFM_Directly, 45 adjustUFM, alterUFM, 46 adjustUFM_Directly, 47 delFromUFM, 48 delFromUFM_Directly, 49 delListFromUFM, 50 delListFromUFM_Directly, 51 plusUFM, 52 plusUFM_C, 53 plusUFM_CD, 54 plusUFM_CD2, 55 plusMaybeUFM_C, 56 plusUFMList, 57 minusUFM, 58 intersectUFM, 59 intersectUFM_C, 60 disjointUFM, 61 equalKeysUFM, 62 nonDetStrictFoldUFM, foldUFM, nonDetStrictFoldUFM_Directly, 63 anyUFM, allUFM, seqEltsUFM, 64 mapUFM, mapUFM_Directly, 65 elemUFM, elemUFM_Directly, 66 filterUFM, filterUFM_Directly, partitionUFM, 67 sizeUFM, 68 isNullUFM, 69 lookupUFM, lookupUFM_Directly, 70 lookupWithDefaultUFM, lookupWithDefaultUFM_Directly, 71 nonDetEltsUFM, eltsUFM, nonDetKeysUFM, 72 ufmToSet_Directly, 73 nonDetUFMToList, ufmToIntMap, unsafeIntMapToUFM, 74 unsafeCastUFMKey, 75 pprUniqFM, pprUFM, pprUFMWithKeys, pluralUFM 76 ) where 77 78import GHC.Prelude 79 80import GHC.Types.Unique ( Uniquable(..), Unique, getKey ) 81import GHC.Utils.Outputable 82 83import qualified Data.IntMap as M 84import qualified Data.IntSet as S 85import Data.Data 86import qualified Data.Semigroup as Semi 87import Data.Functor.Classes (Eq1 (..)) 88 89-- | A finite map from @uniques@ of one type to 90-- elements in another type. 91-- 92-- The key is just here to keep us honest. It's always safe 93-- to use a single type as key. 94-- If two types don't overlap in their uniques it's also safe 95-- to index the same map at multiple key types. But this is 96-- very much discouraged. 97newtype UniqFM key ele = UFM (M.IntMap ele) 98 deriving (Data, Eq, Functor) 99 -- Nondeterministic Foldable and Traversable instances are accessible through 100 -- use of the 'NonDetUniqFM' wrapper. 101 -- See Note [Deterministic UniqFM] in GHC.Types.Unique.DFM to learn about determinism. 102 103emptyUFM :: UniqFM key elt 104emptyUFM = UFM M.empty 105 106isNullUFM :: UniqFM key elt -> Bool 107isNullUFM (UFM m) = M.null m 108 109unitUFM :: Uniquable key => key -> elt -> UniqFM key elt 110unitUFM k v = UFM (M.singleton (getKey $ getUnique k) v) 111 112-- when you've got the Unique already 113unitDirectlyUFM :: Unique -> elt -> UniqFM key elt 114unitDirectlyUFM u v = UFM (M.singleton (getKey u) v) 115 116listToUFM :: Uniquable key => [(key,elt)] -> UniqFM key elt 117listToUFM = foldl' (\m (k, v) -> addToUFM m k v) emptyUFM 118 119listToUFM_Directly :: [(Unique, elt)] -> UniqFM key elt 120listToUFM_Directly = foldl' (\m (u, v) -> addToUFM_Directly m u v) emptyUFM 121 122listToIdentityUFM :: Uniquable key => [key] -> UniqFM key key 123listToIdentityUFM = foldl' (\m x -> addToUFM m x x) emptyUFM 124 125listToUFM_C 126 :: Uniquable key 127 => (elt -> elt -> elt) 128 -> [(key, elt)] 129 -> UniqFM key elt 130listToUFM_C f = foldl' (\m (k, v) -> addToUFM_C f m k v) emptyUFM 131 132addToUFM :: Uniquable key => UniqFM key elt -> key -> elt -> UniqFM key elt 133addToUFM (UFM m) k v = UFM (M.insert (getKey $ getUnique k) v m) 134 135addListToUFM :: Uniquable key => UniqFM key elt -> [(key,elt)] -> UniqFM key elt 136addListToUFM = foldl' (\m (k, v) -> addToUFM m k v) 137 138addListToUFM_Directly :: UniqFM key elt -> [(Unique,elt)] -> UniqFM key elt 139addListToUFM_Directly = foldl' (\m (k, v) -> addToUFM_Directly m k v) 140 141addToUFM_Directly :: UniqFM key elt -> Unique -> elt -> UniqFM key elt 142addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m) 143 144addToUFM_C 145 :: Uniquable key 146 => (elt -> elt -> elt) -- old -> new -> result 147 -> UniqFM key elt -- old 148 -> key -> elt -- new 149 -> UniqFM key elt -- result 150-- Arguments of combining function of M.insertWith and addToUFM_C are flipped. 151addToUFM_C f (UFM m) k v = 152 UFM (M.insertWith (flip f) (getKey $ getUnique k) v m) 153 154addToUFM_Acc 155 :: Uniquable key 156 => (elt -> elts -> elts) -- Add to existing 157 -> (elt -> elts) -- New element 158 -> UniqFM key elts -- old 159 -> key -> elt -- new 160 -> UniqFM key elts -- result 161addToUFM_Acc exi new (UFM m) k v = 162 UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m) 163 164alterUFM 165 :: Uniquable key 166 => (Maybe elt -> Maybe elt) -- How to adjust 167 -> UniqFM key elt -- old 168 -> key -- new 169 -> UniqFM key elt -- result 170alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m) 171 172-- | Add elements to the map, combining existing values with inserted ones using 173-- the given function. 174addListToUFM_C 175 :: Uniquable key 176 => (elt -> elt -> elt) 177 -> UniqFM key elt -> [(key,elt)] 178 -> UniqFM key elt 179addListToUFM_C f = foldl' (\m (k, v) -> addToUFM_C f m k v) 180 181adjustUFM :: Uniquable key => (elt -> elt) -> UniqFM key elt -> key -> UniqFM key elt 182adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m) 183 184adjustUFM_Directly :: (elt -> elt) -> UniqFM key elt -> Unique -> UniqFM key elt 185adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m) 186 187delFromUFM :: Uniquable key => UniqFM key elt -> key -> UniqFM key elt 188delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m) 189 190delListFromUFM :: Uniquable key => UniqFM key elt -> [key] -> UniqFM key elt 191delListFromUFM = foldl' delFromUFM 192 193delListFromUFM_Directly :: UniqFM key elt -> [Unique] -> UniqFM key elt 194delListFromUFM_Directly = foldl' delFromUFM_Directly 195 196delFromUFM_Directly :: UniqFM key elt -> Unique -> UniqFM key elt 197delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m) 198 199-- Bindings in right argument shadow those in the left 200plusUFM :: UniqFM key elt -> UniqFM key elt -> UniqFM key elt 201-- M.union is left-biased, plusUFM should be right-biased. 202plusUFM (UFM x) (UFM y) = UFM (M.union y x) 203 -- Note (M.union y x), with arguments flipped 204 -- M.union is left-biased, plusUFM should be right-biased. 205 206plusUFM_C :: (elt -> elt -> elt) -> UniqFM key elt -> UniqFM key elt -> UniqFM key elt 207plusUFM_C f (UFM x) (UFM y) = UFM (M.unionWith f x y) 208 209-- | `plusUFM_CD f m1 d1 m2 d2` merges the maps using `f` as the 210-- combinding function and `d1` resp. `d2` as the default value if 211-- there is no entry in `m1` reps. `m2`. The domain is the union of 212-- the domains of `m1` and `m2`. 213-- 214-- Representative example: 215-- 216-- @ 217-- plusUFM_CD f {A: 1, B: 2} 23 {B: 3, C: 4} 42 218-- == {A: f 1 42, B: f 2 3, C: f 23 4 } 219-- @ 220plusUFM_CD 221 :: (elta -> eltb -> eltc) 222 -> UniqFM key elta -- map X 223 -> elta -- default for X 224 -> UniqFM key eltb -- map Y 225 -> eltb -- default for Y 226 -> UniqFM key eltc 227plusUFM_CD f (UFM xm) dx (UFM ym) dy 228 = UFM $ M.mergeWithKey 229 (\_ x y -> Just (x `f` y)) 230 (M.map (\x -> x `f` dy)) 231 (M.map (\y -> dx `f` y)) 232 xm ym 233 234-- | `plusUFM_CD2 f m1 m2` merges the maps using `f` as the combining 235-- function. Unlike `plusUFM_CD`, a missing value is not defaulted: it is 236-- instead passed as `Nothing` to `f`. `f` can never have both its arguments 237-- be `Nothing`. 238-- 239-- `plusUFM_CD2 f m1 m2` is the same as `plusUFM_CD f (mapUFM Just m1) Nothing 240-- (mapUFM Just m2) Nothing`. 241plusUFM_CD2 242 :: (Maybe elta -> Maybe eltb -> eltc) 243 -> UniqFM key elta -- map X 244 -> UniqFM key eltb -- map Y 245 -> UniqFM key eltc 246plusUFM_CD2 f (UFM xm) (UFM ym) 247 = UFM $ M.mergeWithKey 248 (\_ x y -> Just (Just x `f` Just y)) 249 (M.map (\x -> Just x `f` Nothing)) 250 (M.map (\y -> Nothing `f` Just y)) 251 xm ym 252 253plusMaybeUFM_C :: (elt -> elt -> Maybe elt) 254 -> UniqFM key elt -> UniqFM key elt -> UniqFM key elt 255plusMaybeUFM_C f (UFM xm) (UFM ym) 256 = UFM $ M.mergeWithKey 257 (\_ x y -> x `f` y) 258 id 259 id 260 xm ym 261 262plusUFMList :: [UniqFM key elt] -> UniqFM key elt 263plusUFMList = foldl' plusUFM emptyUFM 264 265minusUFM :: UniqFM key elt1 -> UniqFM key elt2 -> UniqFM key elt1 266minusUFM (UFM x) (UFM y) = UFM (M.difference x y) 267 268intersectUFM :: UniqFM key elt1 -> UniqFM key elt2 -> UniqFM key elt1 269intersectUFM (UFM x) (UFM y) = UFM (M.intersection x y) 270 271intersectUFM_C 272 :: (elt1 -> elt2 -> elt3) 273 -> UniqFM key elt1 274 -> UniqFM key elt2 275 -> UniqFM key elt3 276intersectUFM_C f (UFM x) (UFM y) = UFM (M.intersectionWith f x y) 277 278disjointUFM :: UniqFM key elt1 -> UniqFM key elt2 -> Bool 279disjointUFM (UFM x) (UFM y) = M.disjoint x y 280 281foldUFM :: (elt -> a -> a) -> a -> UniqFM key elt -> a 282foldUFM k z (UFM m) = M.foldr k z m 283 284mapUFM :: (elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2 285mapUFM f (UFM m) = UFM (M.map f m) 286 287mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2 288mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m) 289 290filterUFM :: (elt -> Bool) -> UniqFM key elt -> UniqFM key elt 291filterUFM p (UFM m) = UFM (M.filter p m) 292 293filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM key elt -> UniqFM key elt 294filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnique) m) 295 296partitionUFM :: (elt -> Bool) -> UniqFM key elt -> (UniqFM key elt, UniqFM key elt) 297partitionUFM p (UFM m) = 298 case M.partition p m of 299 (left, right) -> (UFM left, UFM right) 300 301sizeUFM :: UniqFM key elt -> Int 302sizeUFM (UFM m) = M.size m 303 304elemUFM :: Uniquable key => key -> UniqFM key elt -> Bool 305elemUFM k (UFM m) = M.member (getKey $ getUnique k) m 306 307elemUFM_Directly :: Unique -> UniqFM key elt -> Bool 308elemUFM_Directly u (UFM m) = M.member (getKey u) m 309 310lookupUFM :: Uniquable key => UniqFM key elt -> key -> Maybe elt 311lookupUFM (UFM m) k = M.lookup (getKey $ getUnique k) m 312 313-- when you've got the Unique already 314lookupUFM_Directly :: UniqFM key elt -> Unique -> Maybe elt 315lookupUFM_Directly (UFM m) u = M.lookup (getKey u) m 316 317lookupWithDefaultUFM :: Uniquable key => UniqFM key elt -> elt -> key -> elt 318lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m 319 320lookupWithDefaultUFM_Directly :: UniqFM key elt -> elt -> Unique -> elt 321lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m 322 323eltsUFM :: UniqFM key elt -> [elt] 324eltsUFM (UFM m) = M.elems m 325 326ufmToSet_Directly :: UniqFM key elt -> S.IntSet 327ufmToSet_Directly (UFM m) = M.keysSet m 328 329anyUFM :: (elt -> Bool) -> UniqFM key elt -> Bool 330anyUFM p (UFM m) = M.foldr ((||) . p) False m 331 332allUFM :: (elt -> Bool) -> UniqFM key elt -> Bool 333allUFM p (UFM m) = M.foldr ((&&) . p) True m 334 335seqEltsUFM :: ([elt] -> ()) -> UniqFM key elt -> () 336seqEltsUFM seqList = seqList . nonDetEltsUFM 337 -- It's OK to use nonDetEltsUFM here because the type guarantees that 338 -- the only interesting thing this function can do is to force the 339 -- elements. 340 341-- See Note [Deterministic UniqFM] to learn about nondeterminism. 342-- If you use this please provide a justification why it doesn't introduce 343-- nondeterminism. 344nonDetEltsUFM :: UniqFM key elt -> [elt] 345nonDetEltsUFM (UFM m) = M.elems m 346 347-- See Note [Deterministic UniqFM] to learn about nondeterminism. 348-- If you use this please provide a justification why it doesn't introduce 349-- nondeterminism. 350nonDetKeysUFM :: UniqFM key elt -> [Unique] 351nonDetKeysUFM (UFM m) = map getUnique $ M.keys m 352 353-- See Note [Deterministic UniqFM] to learn about nondeterminism. 354-- If you use this please provide a justification why it doesn't introduce 355-- nondeterminism. 356nonDetStrictFoldUFM :: (elt -> a -> a) -> a -> UniqFM key elt -> a 357nonDetStrictFoldUFM k z (UFM m) = M.foldl' (flip k) z m 358 359-- See Note [Deterministic UniqFM] to learn about nondeterminism. 360-- If you use this please provide a justification why it doesn't introduce 361-- nondeterminism. 362nonDetStrictFoldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM key elt -> a 363nonDetStrictFoldUFM_Directly k z (UFM m) = M.foldlWithKey' (\z' i x -> k (getUnique i) x z') z m 364 365-- See Note [Deterministic UniqFM] to learn about nondeterminism. 366-- If you use this please provide a justification why it doesn't introduce 367-- nondeterminism. 368nonDetUFMToList :: UniqFM key elt -> [(Unique, elt)] 369nonDetUFMToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m 370 371-- | A wrapper around 'UniqFM' with the sole purpose of informing call sites 372-- that the provided 'Foldable' and 'Traversable' instances are 373-- nondeterministic. 374-- If you use this please provide a justification why it doesn't introduce 375-- nondeterminism. 376-- See Note [Deterministic UniqFM] in "GHC.Types.Unique.DFM" to learn about determinism. 377newtype NonDetUniqFM key ele = NonDetUniqFM { getNonDet :: UniqFM key ele } 378 deriving (Functor) 379 380-- | Inherently nondeterministic. 381-- If you use this please provide a justification why it doesn't introduce 382-- nondeterminism. 383-- See Note [Deterministic UniqFM] in "GHC.Types.Unique.DFM" to learn about determinism. 384instance forall key. Foldable (NonDetUniqFM key) where 385 foldr f z (NonDetUniqFM (UFM m)) = foldr f z m 386 387-- | Inherently nondeterministic. 388-- If you use this please provide a justification why it doesn't introduce 389-- nondeterminism. 390-- See Note [Deterministic UniqFM] in "GHC.Types.Unique.DFM" to learn about determinism. 391instance forall key. Traversable (NonDetUniqFM key) where 392 traverse f (NonDetUniqFM (UFM m)) = NonDetUniqFM . UFM <$> traverse f m 393 394ufmToIntMap :: UniqFM key elt -> M.IntMap elt 395ufmToIntMap (UFM m) = m 396 397unsafeIntMapToUFM :: M.IntMap elt -> UniqFM key elt 398unsafeIntMapToUFM = UFM 399 400-- | Cast the key domain of a UniqFM. 401-- 402-- As long as the domains don't overlap in their uniques 403-- this is safe. 404unsafeCastUFMKey :: UniqFM key1 elt -> UniqFM key2 elt 405unsafeCastUFMKey (UFM m) = UFM m 406 407-- Determines whether two 'UniqFM's contain the same keys. 408equalKeysUFM :: UniqFM key a -> UniqFM key b -> Bool 409equalKeysUFM (UFM m1) (UFM m2) = liftEq (\_ _ -> True) m1 m2 410 411-- Instances 412 413instance Semi.Semigroup (UniqFM key a) where 414 (<>) = plusUFM 415 416instance Monoid (UniqFM key a) where 417 mempty = emptyUFM 418 mappend = (Semi.<>) 419 420-- Output-ery 421 422instance Outputable a => Outputable (UniqFM key a) where 423 ppr ufm = pprUniqFM ppr ufm 424 425pprUniqFM :: (a -> SDoc) -> UniqFM key a -> SDoc 426pprUniqFM ppr_elt ufm 427 = brackets $ fsep $ punctuate comma $ 428 [ ppr uq <+> text ":->" <+> ppr_elt elt 429 | (uq, elt) <- nonDetUFMToList ufm ] 430 -- It's OK to use nonDetUFMToList here because we only use it for 431 -- pretty-printing. 432 433-- | Pretty-print a non-deterministic set. 434-- The order of variables is non-deterministic and for pretty-printing that 435-- shouldn't be a problem. 436-- Having this function helps contain the non-determinism created with 437-- nonDetEltsUFM. 438pprUFM :: UniqFM key a -- ^ The things to be pretty printed 439 -> ([a] -> SDoc) -- ^ The pretty printing function to use on the elements 440 -> SDoc -- ^ 'SDoc' where the things have been pretty 441 -- printed 442pprUFM ufm pp = pp (nonDetEltsUFM ufm) 443 444-- | Pretty-print a non-deterministic set. 445-- The order of variables is non-deterministic and for pretty-printing that 446-- shouldn't be a problem. 447-- Having this function helps contain the non-determinism created with 448-- nonDetUFMToList. 449pprUFMWithKeys 450 :: UniqFM key a -- ^ The things to be pretty printed 451 -> ([(Unique, a)] -> SDoc) -- ^ The pretty printing function to use on the elements 452 -> SDoc -- ^ 'SDoc' where the things have been pretty 453 -- printed 454pprUFMWithKeys ufm pp = pp (nonDetUFMToList ufm) 455 456-- | Determines the pluralisation suffix appropriate for the length of a set 457-- in the same way that plural from Outputable does for lists. 458pluralUFM :: UniqFM key a -> SDoc 459pluralUFM ufm 460 | sizeUFM ufm == 1 = empty 461 | otherwise = char 's' 462