UniqFM.hs 14.4 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3 4
{-
(c) The University of Glasgow 2006
(c) The AQUA Project, Glasgow University, 1994-1998

Simon Marlow's avatar
Simon Marlow committed
5

6
UniqFM: Specialised finite maps, for things with @Uniques@.
7

8
Basically, the things need to be in class @Uniquable@, and we use the
9
@getUnique@ method to grab their @Uniques@.
10 11 12

(A similar thing to @UniqSet@, as opposed to @Set@.)

13
The interface is based on @FiniteMap@s, but the implementation uses
ian@well-typed.com's avatar
ian@well-typed.com committed
14
@Data.IntMap@, which is both maintained and faster than the past
15 16 17 18 19 20
implementation (see commit log).

The @UniqFM@ interface maps directly to Data.IntMap, only
``Data.IntMap.union'' is left-biased and ``plusUFM'' right-biased
and ``addToUFM\_C'' and ``Data.IntMap.insertWith'' differ in the order
of arguments of combining function.
Austin Seipp's avatar
Austin Seipp committed
21
-}
22

23 24 25
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
26
{-# OPTIONS_GHC -Wall #-}
27

28
module UniqFM (
ian@well-typed.com's avatar
ian@well-typed.com committed
29 30
        -- * Unique-keyed mappings
        UniqFM,       -- abstract type
31

32
        -- ** Manipulating those mappings
ian@well-typed.com's avatar
ian@well-typed.com committed
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
        emptyUFM,
        unitUFM,
        unitDirectlyUFM,
        listToUFM,
        listToUFM_Directly,
        listToUFM_C,
        addToUFM,addToUFM_C,addToUFM_Acc,
        addListToUFM,addListToUFM_C,
        addToUFM_Directly,
        addListToUFM_Directly,
        adjustUFM, alterUFM,
        adjustUFM_Directly,
        delFromUFM,
        delFromUFM_Directly,
        delListFromUFM,
48
        delListFromUFM_Directly,
ian@well-typed.com's avatar
ian@well-typed.com committed
49 50
        plusUFM,
        plusUFM_C,
51
        plusUFM_CD,
ian@well-typed.com's avatar
ian@well-typed.com committed
52 53 54
        minusUFM,
        intersectUFM,
        intersectUFM_C,
Jan Stolarek's avatar
Jan Stolarek committed
55
        disjointUFM,
niteria's avatar
niteria committed
56 57
        nonDetFoldUFM, foldUFM, nonDetFoldUFM_Directly,
        anyUFM, allUFM,
ian@well-typed.com's avatar
ian@well-typed.com committed
58 59
        mapUFM, mapUFM_Directly,
        elemUFM, elemUFM_Directly,
60
        filterUFM, filterUFM_Directly, partitionUFM,
ian@well-typed.com's avatar
ian@well-typed.com committed
61 62 63 64
        sizeUFM,
        isNullUFM,
        lookupUFM, lookupUFM_Directly,
        lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
niteria's avatar
niteria committed
65
        nonDetEltsUFM, eltsUFM, nonDetKeysUFM,
Joachim Breitner's avatar
Joachim Breitner committed
66
        ufmToSet_Directly,
niteria's avatar
niteria committed
67 68
        nonDetUFMToList, ufmToIntMap,
        pprUniqFM, pprUFM, pprUFMWithKeys, pluralUFM
69 70
    ) where

71
import Unique           ( Uniquable(..), Unique, getKey )
72
import Outputable
73 74

import qualified Data.IntMap as M
Joachim Breitner's avatar
Joachim Breitner committed
75
import qualified Data.IntSet as S
76 77
import Data.Typeable
import Data.Data
78 79 80 81
#if __GLASGOW_HASKELL__ > 710
import Data.Semigroup   ( Semigroup )
import qualified Data.Semigroup as Semigroup
#endif
82

Austin Seipp's avatar
Austin Seipp committed
83 84 85
{-
************************************************************************
*                                                                      *
86
\subsection{The signature of the module}
Austin Seipp's avatar
Austin Seipp committed
87 88 89
*                                                                      *
************************************************************************
-}
90

ian@well-typed.com's avatar
ian@well-typed.com committed
91 92 93
emptyUFM        :: UniqFM elt
isNullUFM       :: UniqFM elt -> Bool
unitUFM         :: Uniquable key => key -> elt -> UniqFM elt
94
unitDirectlyUFM -- got the Unique already
ian@well-typed.com's avatar
ian@well-typed.com committed
95 96
                :: Unique -> elt -> UniqFM elt
listToUFM       :: Uniquable key => [(key,elt)] -> UniqFM elt
97
listToUFM_Directly
ian@well-typed.com's avatar
ian@well-typed.com committed
98 99 100
                :: [(Unique, elt)] -> UniqFM elt
listToUFM_C     :: Uniquable key => (elt -> elt -> elt)
                           -> [(key, elt)]
101
                           -> UniqFM elt
102

ian@well-typed.com's avatar
ian@well-typed.com committed
103 104
addToUFM        :: Uniquable key => UniqFM elt -> key -> elt  -> UniqFM elt
addListToUFM    :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
105
addListToUFM_Directly :: UniqFM elt -> [(Unique,elt)] -> UniqFM elt
106
addToUFM_Directly
ian@well-typed.com's avatar
ian@well-typed.com committed
107
                :: UniqFM elt -> Unique -> elt -> UniqFM elt
108

ian@well-typed.com's avatar
ian@well-typed.com committed
109 110 111 112
addToUFM_C      :: Uniquable key => (elt -> elt -> elt) -- old -> new -> result
                           -> UniqFM elt                -- old
                           -> key -> elt                -- new
                           -> UniqFM elt                -- result
113

ian@well-typed.com's avatar
ian@well-typed.com committed
114 115 116 117 118 119
addToUFM_Acc    :: Uniquable key =>
                              (elt -> elts -> elts)     -- Add to existing
                           -> (elt -> elts)             -- New element
                           -> UniqFM elts               -- old
                           -> key -> elt                -- new
                           -> UniqFM elts               -- result
120

ian@well-typed.com's avatar
ian@well-typed.com committed
121
alterUFM        :: Uniquable key =>
122
                              (Maybe elt -> Maybe elt)  -- How to adjust
ian@well-typed.com's avatar
ian@well-typed.com committed
123 124 125
                           -> UniqFM elt                -- old
                           -> key                       -- new
                           -> UniqFM elt                -- result
126

ian@well-typed.com's avatar
ian@well-typed.com committed
127 128 129
addListToUFM_C  :: Uniquable key => (elt -> elt -> elt)
                           -> UniqFM elt -> [(key,elt)]
                           -> UniqFM elt
130

ian@well-typed.com's avatar
ian@well-typed.com committed
131
adjustUFM       :: Uniquable key => (elt -> elt) -> UniqFM elt -> key -> UniqFM elt
132 133
adjustUFM_Directly :: (elt -> elt) -> UniqFM elt -> Unique -> UniqFM elt

ian@well-typed.com's avatar
ian@well-typed.com committed
134 135
delFromUFM      :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
delListFromUFM  :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
136
delListFromUFM_Directly :: UniqFM elt -> [Unique] -> UniqFM elt
137
delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
138

139
-- Bindings in right argument shadow those in the left
ian@well-typed.com's avatar
ian@well-typed.com committed
140
plusUFM         :: UniqFM elt -> UniqFM elt -> UniqFM elt
141

ian@well-typed.com's avatar
ian@well-typed.com committed
142 143
plusUFM_C       :: (elt -> elt -> elt)
                -> UniqFM elt -> UniqFM elt -> UniqFM elt
144

145 146 147 148 149 150 151 152 153 154 155
-- | `plusUFM_CD f m1 d1 m2 d2` merges the maps using `f` as the
-- combinding function and `d1` resp. `d2` as the default value if
-- there is no entry in `m1` reps. `m2`. The domain is the union of
-- the domains of `m1` and `m2`.
--
-- Representative example:
--
-- @
-- plusUFM_CD f {A: 1, B: 2} 23 {B: 3, C: 4} 42
--    == {A: f 1 42, B: f 2 3, C: f 23 4 }
-- @
156 157 158
plusUFM_CD      :: (elt -> elt -> elt)
                -> UniqFM elt -> elt -> UniqFM elt -> elt -> UniqFM elt

ian@well-typed.com's avatar
ian@well-typed.com committed
159
minusUFM        :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
160

ian@well-typed.com's avatar
ian@well-typed.com committed
161 162 163
intersectUFM    :: UniqFM elt -> UniqFM elt -> UniqFM elt
intersectUFM_C  :: (elt1 -> elt2 -> elt3)
                -> UniqFM elt1 -> UniqFM elt2 -> UniqFM elt3
Jan Stolarek's avatar
Jan Stolarek committed
164
disjointUFM     :: UniqFM elt1 -> UniqFM elt2 -> Bool
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
165

ian@well-typed.com's avatar
ian@well-typed.com committed
166 167
foldUFM         :: (elt -> a -> a) -> a -> UniqFM elt -> a
mapUFM          :: (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
Edward Z. Yang's avatar
Edward Z. Yang committed
168
mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
ian@well-typed.com's avatar
ian@well-typed.com committed
169
filterUFM       :: (elt -> Bool) -> UniqFM elt -> UniqFM elt
170
filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM elt -> UniqFM elt
171
partitionUFM    :: (elt -> Bool) -> UniqFM elt -> (UniqFM elt, UniqFM elt)
172

ian@well-typed.com's avatar
ian@well-typed.com committed
173 174 175
sizeUFM         :: UniqFM elt -> Int
--hashUFM               :: UniqFM elt -> Int
elemUFM         :: Uniquable key => key -> UniqFM elt -> Bool
176
elemUFM_Directly:: Unique -> UniqFM elt -> Bool
177

ian@well-typed.com's avatar
ian@well-typed.com committed
178
lookupUFM       :: Uniquable key => UniqFM elt -> key -> Maybe elt
179
lookupUFM_Directly  -- when you've got the Unique already
ian@well-typed.com's avatar
ian@well-typed.com committed
180
                :: UniqFM elt -> Unique -> Maybe elt
181
lookupWithDefaultUFM
ian@well-typed.com's avatar
ian@well-typed.com committed
182
                :: Uniquable key => UniqFM elt -> elt -> key -> elt
183
lookupWithDefaultUFM_Directly
ian@well-typed.com's avatar
ian@well-typed.com committed
184 185
                :: UniqFM elt -> elt -> Unique -> elt
eltsUFM         :: UniqFM elt -> [elt]
Joachim Breitner's avatar
Joachim Breitner committed
186
ufmToSet_Directly :: UniqFM elt -> S.IntSet
187

Austin Seipp's avatar
Austin Seipp committed
188 189 190
{-
************************************************************************
*                                                                      *
cactus's avatar
cactus committed
191
\subsection{Monoid interface}
Austin Seipp's avatar
Austin Seipp committed
192 193 194
*                                                                      *
************************************************************************
-}
cactus's avatar
cactus committed
195

196 197 198 199 200
#if __GLASGOW_HASKELL__ > 710
instance Semigroup (UniqFM a) where
  (<>) = plusUFM
#endif

cactus's avatar
cactus committed
201 202 203 204
instance Monoid (UniqFM a) where
    mempty = emptyUFM
    mappend = plusUFM

Austin Seipp's avatar
Austin Seipp committed
205 206 207
{-
************************************************************************
*                                                                      *
208
\subsection{Implementation using ``Data.IntMap''}
Austin Seipp's avatar
Austin Seipp committed
209 210 211
*                                                                      *
************************************************************************
-}
212

213

214 215 216 217 218 219
newtype UniqFM ele = UFM (M.IntMap ele)
  deriving (Data, Eq, Functor, Typeable)
  -- We used to derive Traversable and Foldable, but they were nondeterministic
  -- and not obvious at the call site. You can use explicit nonDetEltsUFM
  -- and fold a list if needed.
  -- See Note [Deterministic UniqFM] in UniqDFM to learn about determinism.
220

221 222 223 224 225 226 227 228
emptyUFM = UFM M.empty
isNullUFM (UFM m) = M.null m
unitUFM k v = UFM (M.singleton (getKey $ getUnique k) v)
unitDirectlyUFM u v = UFM (M.singleton (getKey u) v)
listToUFM = foldl (\m (k, v) -> addToUFM m k v) emptyUFM
listToUFM_Directly = foldl (\m (u, v) -> addToUFM_Directly m u v) emptyUFM
listToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v) emptyUFM

229 230
alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m)
addToUFM (UFM m) k v   = UFM (M.insert (getKey $ getUnique k) v m)
231 232 233 234 235 236 237 238 239 240 241
addListToUFM = foldl (\m (k, v) -> addToUFM m k v)
addListToUFM_Directly = foldl (\m (k, v) -> addToUFM_Directly m k v)
addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m)

-- Arguments of combining function of M.insertWith and addToUFM_C are flipped.
addToUFM_C f (UFM m) k v =
  UFM (M.insertWith (flip f) (getKey $ getUnique k) v m)
addToUFM_Acc exi new (UFM m) k v =
  UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
addListToUFM_C f = foldl (\m (k, v) -> addToUFM_C f m k v)

242 243 244
adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m)
adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m)

245 246 247
delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m)
delListFromUFM = foldl delFromUFM
delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m)
248
delListFromUFM_Directly = foldl delFromUFM_Directly
249 250 251

-- M.union is left-biased, plusUFM should be right-biased.
plusUFM (UFM x) (UFM y) = UFM (M.union y x)
252 253 254
     -- Note (M.union y x), with arguments flipped
     -- M.union is left-biased, plusUFM should be right-biased.

255
plusUFM_C f (UFM x) (UFM y) = UFM (M.unionWith f x y)
256 257 258 259 260 261 262

plusUFM_CD f (UFM xm) dx (UFM ym) dy
    = UFM $ M.mergeWithKey
        (\_ x y -> Just (x `f` y))
        (M.map (\x -> x `f` dy))
        (M.map (\y -> dx `f` y))
        xm ym
263 264 265
minusUFM (UFM x) (UFM y) = UFM (M.difference x y)
intersectUFM (UFM x) (UFM y) = UFM (M.intersection x y)
intersectUFM_C f (UFM x) (UFM y) = UFM (M.intersectionWith f x y)
Jan Stolarek's avatar
Jan Stolarek committed
266
disjointUFM (UFM x) (UFM y) = M.null (M.intersection x y)
267 268

foldUFM k z (UFM m) = M.fold k z m
269 270


271
mapUFM f (UFM m) = UFM (M.map f m)
Edward Z. Yang's avatar
Edward Z. Yang committed
272
mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m)
273 274
filterUFM p (UFM m) = UFM (M.filter p m)
filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnique) m)
275 276
partitionUFM p (UFM m) = case M.partition p m of
                           (left, right) -> (UFM left, UFM right)
277 278 279 280 281 282 283 284 285 286

sizeUFM (UFM m) = M.size m
elemUFM k (UFM m) = M.member (getKey $ getUnique k) m
elemUFM_Directly u (UFM m) = M.member (getKey u) m

lookupUFM (UFM m) k = M.lookup (getKey $ getUnique k) m
lookupUFM_Directly (UFM m) u = M.lookup (getKey u) m
lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m
lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m
eltsUFM (UFM m) = M.elems m
Joachim Breitner's avatar
Joachim Breitner committed
287
ufmToSet_Directly (UFM m) = M.keysSet m
288

289 290 291 292 293 294
anyUFM :: (elt -> Bool) -> UniqFM elt -> Bool
anyUFM p (UFM m) = M.fold ((||) . p) False m

allUFM :: (elt -> Bool) -> UniqFM elt -> Bool
allUFM p (UFM m) = M.fold ((&&) . p) True m

niteria's avatar
niteria committed
295 296 297
-- See Note [Deterministic UniqFM] to learn about nondeterminism.
-- If you use this please provide a justification why it doesn't introduce
-- nondeterminism.
298 299 300
nonDetEltsUFM :: UniqFM elt -> [elt]
nonDetEltsUFM (UFM m) = M.elems m

niteria's avatar
niteria committed
301 302 303 304 305 306
-- See Note [Deterministic UniqFM] to learn about nondeterminism.
-- If you use this please provide a justification why it doesn't introduce
-- nondeterminism.
nonDetKeysUFM :: UniqFM elt -> [Unique]
nonDetKeysUFM (UFM m) = map getUnique $ M.keys m

niteria's avatar
niteria committed
307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324
-- See Note [Deterministic UniqFM] to learn about nondeterminism.
-- If you use this please provide a justification why it doesn't introduce
-- nondeterminism.
nonDetFoldUFM :: (elt -> a -> a) -> a -> UniqFM elt -> a
nonDetFoldUFM k z (UFM m) = M.fold k z m

-- See Note [Deterministic UniqFM] to learn about nondeterminism.
-- If you use this please provide a justification why it doesn't introduce
-- nondeterminism.
nonDetFoldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM elt -> a
nonDetFoldUFM_Directly k z (UFM m) = M.foldWithKey (k . getUnique) z m

-- See Note [Deterministic UniqFM] to learn about nondeterminism.
-- If you use this please provide a justification why it doesn't introduce
-- nondeterminism.
nonDetUFMToList :: UniqFM elt -> [(Unique, elt)]
nonDetUFMToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m

325 326 327
ufmToIntMap :: UniqFM elt -> M.IntMap elt
ufmToIntMap (UFM m) = m

Austin Seipp's avatar
Austin Seipp committed
328 329 330
{-
************************************************************************
*                                                                      *
331
\subsection{Output-ery}
Austin Seipp's avatar
Austin Seipp committed
332 333 334
*                                                                      *
************************************************************************
-}
335

336
instance Outputable a => Outputable (UniqFM a) where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
337 338 339 340 341
    ppr ufm = pprUniqFM ppr ufm

pprUniqFM :: (a -> SDoc) -> UniqFM a -> SDoc
pprUniqFM ppr_elt ufm
  = brackets $ fsep $ punctuate comma $
342
    [ ppr uq <+> text ":->" <+> ppr_elt elt
niteria's avatar
niteria committed
343 344 345
    | (uq, elt) <- nonDetUFMToList ufm ]
  -- It's OK to use nonDetUFMToList here because we only use it for
  -- pretty-printing.
346 347 348 349 350

-- | Pretty-print a non-deterministic set.
-- The order of variables is non-deterministic and for pretty-printing that
-- shouldn't be a problem.
-- Having this function helps contain the non-determinism created with
351 352 353
-- nonDetEltsUFM.
pprUFM :: UniqFM a      -- ^ The things to be pretty printed
       -> ([a] -> SDoc) -- ^ The pretty printing function to use on the elements
354 355
       -> SDoc          -- ^ 'SDoc' where the things have been pretty
                        -- printed
356
pprUFM ufm pp = pp (nonDetEltsUFM ufm)
357

niteria's avatar
niteria committed
358 359 360 361 362 363 364 365 366 367 368 369
-- | Pretty-print a non-deterministic set.
-- The order of variables is non-deterministic and for pretty-printing that
-- shouldn't be a problem.
-- Having this function helps contain the non-determinism created with
-- nonDetUFMToList.
pprUFMWithKeys
       :: UniqFM a                -- ^ The things to be pretty printed
       -> ([(Unique, a)] -> SDoc) -- ^ The pretty printing function to use on the elements
       -> SDoc                    -- ^ 'SDoc' where the things have been pretty
                                  -- printed
pprUFMWithKeys ufm pp = pp (nonDetUFMToList ufm)

370 371 372 373 374 375
-- | Determines the pluralisation suffix appropriate for the length of a set
-- in the same way that plural from Outputable does for lists.
pluralUFM :: UniqFM a -> SDoc
pluralUFM ufm
  | sizeUFM ufm == 1 = empty
  | otherwise = char 's'