compiler/utils/UniqFM.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 % (c) The AQUA Project, Glasgow University, 1994-1998
   4 %
   5
   6 UniqFM: Specialised finite maps, for things with @Uniques@
   7
   8 Based on @FiniteMaps@ (as you would expect).
   9
  10 Basically, the things need to be in class @Uniquable@, and we use the
  11 @getUnique@ method to grab their @Uniques@.
  12
  13 (A similar thing to @UniqSet@, as opposed to @Set@.)
  14
  15 \begin{code}
  16 {-# OPTIONS -Wall -fno-warn-name-shadowing #-}
  17 module UniqFM (
  18         -- * Unique-keyed mappings
  19         UniqFM(..),     -- abstract type
  20                         -- (de-abstracted for MachRegs.trivColorable optimisation BL 2007/09)
  21
  22         -- ** Manipulating those mappings
  23         emptyUFM,
  24         unitUFM,
  25         unitDirectlyUFM,
  26         listToUFM,
  27         listToUFM_Directly,
  28         listToUFM_C,
  29         addToUFM,addToUFM_C,addToUFM_Acc,
  30         addListToUFM,addListToUFM_C,
  31         addToUFM_Directly,
  32         addListToUFM_Directly,
  33         delFromUFM,
  34         delFromUFM_Directly,
  35         delListFromUFM,
  36         plusUFM,
  37         plusUFM_C,
  38         minusUFM,
  39         intersectsUFM,
  40         intersectUFM,
  41         intersectUFM_C,
  42         foldUFM, foldUFM_Directly,
  43         mapUFM,
  44         elemUFM, elemUFM_Directly,
  45         filterUFM, filterUFM_Directly,
  46         sizeUFM,
  47         hashUFM,
  48         isNullUFM,
  49         lookupUFM, lookupUFM_Directly,
  50         lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
  51         eltsUFM, keysUFM,
  52         ufmToList
  53     ) where
  54
  55 #include "HsVersions.h"
  56
  57 import Unique           ( Uniquable(..), Unique, getKeyFastInt, mkUniqueGrimily )
  58 import Maybes           ( maybeToBool )
  59 import FastTypes
  60 import Outputable
  61 \end{code}
  62
  63 %************************************************************************
  64 %*                                                                      *
  65 \subsection{The @UniqFM@ type, and signatures for the functions}
  66 %*                                                                      *
  67 %************************************************************************
  68
  69 We use @FiniteMaps@, with a (@getUnique@-able) @Unique@ as ``key''.
  70
  71 \begin{code}
  72 emptyUFM        :: UniqFM elt
  73 isNullUFM       :: UniqFM elt -> Bool
  74 unitUFM         :: Uniquable key => key -> elt -> UniqFM elt
  75 unitDirectlyUFM -- got the Unique already
  76                 :: Unique -> elt -> UniqFM elt
  77 listToUFM       :: Uniquable key => [(key,elt)] -> UniqFM elt
  78 listToUFM_Directly
  79                 :: [(Unique, elt)] -> UniqFM elt
  80 listToUFM_C     :: Uniquable key => (elt -> elt -> elt)
  81                            -> [(key, elt)]
  82                            -> UniqFM elt
  83
  84 addToUFM        :: Uniquable key => UniqFM elt -> key -> elt  -> UniqFM elt
  85 addListToUFM    :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
  86 addListToUFM_Directly :: UniqFM elt -> [(Unique,elt)] -> UniqFM elt
  87 addToUFM_Directly
  88                 :: UniqFM elt -> Unique -> elt -> UniqFM elt
  89
  90 addToUFM_C      :: Uniquable key => (elt -> elt -> elt) -- old -> new -> result
  91                            -> UniqFM elt                -- old
  92                            -> key -> elt                -- new
  93                            -> UniqFM elt                -- result
  94
  95 addToUFM_Acc    :: Uniquable key =>
  96                               (elt -> elts -> elts)     -- Add to existing
  97                            -> (elt -> elts)             -- New element
  98                            -> UniqFM elts               -- old
  99                            -> key -> elt                -- new
 100                            -> UniqFM elts               -- result
 101
 102 addListToUFM_C  :: Uniquable key => (elt -> elt -> elt)
 103                            -> UniqFM elt -> [(key,elt)]
 104                            -> UniqFM elt
 105
 106 delFromUFM      :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
 107 delListFromUFM  :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
 108 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
 109
 110 plusUFM         :: UniqFM elt -> UniqFM elt -> UniqFM elt
 111
 112 plusUFM_C       :: (elt -> elt -> elt)
 113                 -> UniqFM elt -> UniqFM elt -> UniqFM elt
 114
 115 minusUFM        :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
 116
 117 intersectUFM    :: UniqFM elt -> UniqFM elt -> UniqFM elt
 118 intersectUFM_C  :: (elt1 -> elt2 -> elt3)
 119                 -> UniqFM elt1 -> UniqFM elt2 -> UniqFM elt3
 120 intersectsUFM   :: UniqFM elt1 -> UniqFM elt2 -> Bool
 121
 122 foldUFM         :: (elt -> a -> a) -> a -> UniqFM elt -> a
 123 foldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM elt -> a
 124 mapUFM          :: (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
 125 filterUFM       :: (elt -> Bool) -> UniqFM elt -> UniqFM elt
 126 filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM elt -> UniqFM elt
 127
 128 sizeUFM         :: UniqFM elt -> Int
 129 hashUFM         :: UniqFM elt -> Int
 130 elemUFM         :: Uniquable key => key -> UniqFM elt -> Bool
 131 elemUFM_Directly:: Unique -> UniqFM elt -> Bool
 132
 133 lookupUFM       :: Uniquable key => UniqFM elt -> key -> Maybe elt
 134 lookupUFM_Directly  -- when you've got the Unique already
 135                 :: UniqFM elt -> Unique -> Maybe elt
 136 lookupWithDefaultUFM
 137                 :: Uniquable key => UniqFM elt -> elt -> key -> elt
 138 lookupWithDefaultUFM_Directly
 139                 :: UniqFM elt -> elt -> Unique -> elt
 140
 141 keysUFM         :: UniqFM elt -> [Unique]       -- Get the keys
 142 eltsUFM         :: UniqFM elt -> [elt]
 143 ufmToList       :: UniqFM elt -> [(Unique, elt)]
 144 \end{code}
 145
 146 %************************************************************************
 147 %*                                                                      *
 148 \subsection{The @IdFinMap@ and @TyVarFinMap@ specialisations for Ids/TyVars}
 149 %*                                                                      *
 150 %************************************************************************
 151
 152 \begin{code}
 153 -- Turn off for now, these need to be updated (SDM 4/98)
 154
 155 #if 0
 156 #ifdef __GLASGOW_HASKELL__
 157 -- I don't think HBC was too happy about this (WDP 94/10)
 158
 159 {-# SPECIALIZE
 160     addListToUFM :: UniqFM elt -> [(Name,   elt)] -> UniqFM elt
 161   #-}
 162 {-# SPECIALIZE
 163     addListToUFM_C :: (elt -> elt -> elt) -> UniqFM elt -> [(Name,  elt)] -> UniqFM elt
 164   #-}
 165 {-# SPECIALIZE
 166     addToUFM    :: UniqFM elt -> Unique -> elt  -> UniqFM elt
 167   #-}
 168 {-# SPECIALIZE
 169     listToUFM   :: [(Unique, elt)]     -> UniqFM elt
 170   #-}
 171 {-# SPECIALIZE
 172     lookupUFM   :: UniqFM elt -> Name   -> Maybe elt
 173                  , UniqFM elt -> Unique -> Maybe elt
 174   #-}
 175
 176 #endif /* __GLASGOW_HASKELL__ */
 177 #endif
 178 \end{code}
 179
 180 %************************************************************************
 181 %*                                                                      *
 182 \subsection{Andy Gill's underlying @UniqFM@ machinery}
 183 %*                                                                      *
 184 %************************************************************************
 185
 186 ``Uniq Finite maps'' are the heart and soul of the compiler's
 187 lookup-tables/environments.  Important stuff!  It works well with
 188 Dense and Sparse ranges.
 189 Both @Uq@ Finite maps and @Hash@ Finite Maps
 190 are built ontop of Int Finite Maps.
 191
 192 This code is explained in the paper:
 193 \begin{display}
 194         A Gill, S Peyton Jones, B O'Sullivan, W Partain and Aqua Friends
 195         "A Cheap balancing act that grows on a tree"
 196         Glasgow FP Workshop, Sep 1994, pp??-??
 197 \end{display}
 198
 199 %************************************************************************
 200 %*                                                                      *
 201 \subsubsection{The @UniqFM@ type, and signatures for the functions}
 202 %*                                                                      *
 203 %************************************************************************
 204
 205 First, the DataType itself; which is either a Node, a Leaf, or an Empty.
 206
 207 \begin{code}
 208 -- | @UniqFM a@ is a mapping from Unique to @a@. DO NOT use these constructors
 209 -- directly unless you live in this module!
 210 data UniqFM ele
 211   = EmptyUFM
 212   | LeafUFM !FastInt ele
 213   | NodeUFM !FastInt         -- the switching
 214             !FastInt         -- the delta
 215             (UniqFM ele)
 216             (UniqFM ele)
 217 -- INVARIANT: the children of a NodeUFM are never EmptyUFMs
 218
 219 {-
 220 -- for debugging only :-)
 221 instance Outputable (UniqFM a) where
 222         ppr(NodeUFM a b t1 t2) =
 223                 sep [text "NodeUFM " <+> int IBOX(a) <+> int IBOX(b),
 224                      nest 1 (parens (ppr t1)),
 225                      nest 1 (parens (ppr t2))]
 226         ppr (LeafUFM x a) = text "LeafUFM " <+> int IBOX(x)
 227         ppr (EmptyUFM)    = empty
 228 -}
 229 -- and when not debugging the package itself...
 230 instance Outputable a => Outputable (UniqFM a) where
 231     ppr ufm = ppr (ufmToList ufm)
 232 \end{code}
 233
 234 %************************************************************************
 235 %*                                                                      *
 236 \subsubsection{The @UniqFM@ functions}
 237 %*                                                                      *
 238 %************************************************************************
 239
 240 First the ways of building a UniqFM.
 241
 242 \begin{code}
 243 emptyUFM                     = EmptyUFM
 244 unitUFM      key elt = mkLeafUFM (getKeyFastInt (getUnique key)) elt
 245 unitDirectlyUFM key elt = mkLeafUFM (getKeyFastInt key) elt
 246
 247 listToUFM key_elt_pairs
 248   = addListToUFM_C use_snd EmptyUFM key_elt_pairs
 249
 250 listToUFM_Directly uniq_elt_pairs
 251   = addListToUFM_directly_C use_snd EmptyUFM uniq_elt_pairs
 252
 253 listToUFM_C combiner key_elt_pairs
 254   = addListToUFM_C combiner EmptyUFM key_elt_pairs
 255 \end{code}
 256
 257 Now ways of adding things to UniqFMs.
 258
 259 There is an alternative version of @addListToUFM_C@, that uses @plusUFM@,
 260 but the semantics of this operation demands a linear insertion;
 261 perhaps the version without the combinator function
 262 could be optimised using it.
 263
 264 \begin{code}
 265 addToUFM fm key elt = addToUFM_C use_snd fm key elt
 266
 267 addToUFM_Directly fm u elt = insert_ele use_snd fm (getKeyFastInt u) elt
 268
 269 addToUFM_C combiner fm key elt
 270   = insert_ele combiner fm (getKeyFastInt (getUnique key)) elt
 271
 272 addToUFM_Acc add unit fm key item
 273   = insert_ele combiner fm (getKeyFastInt (getUnique key)) (unit item)
 274   where
 275     combiner old _unit_item = add item old
 276
 277 addListToUFM fm key_elt_pairs = addListToUFM_C use_snd fm key_elt_pairs
 278 addListToUFM_Directly fm uniq_elt_pairs = addListToUFM_directly_C use_snd fm uniq_elt_pairs
 279
 280 addListToUFM_C combiner fm key_elt_pairs
 281  = foldl (\ fm (k, e) -> insert_ele combiner fm (getKeyFastInt (getUnique k)) e)
 282          fm key_elt_pairs
 283
 284 addListToUFM_directly_C :: (elt -> elt -> elt) -> UniqFM elt -> [(Unique,elt)] -> UniqFM elt
 285 addListToUFM_directly_C combiner fm uniq_elt_pairs
 286  = foldl (\ fm (k, e) -> insert_ele combiner fm (getKeyFastInt k) e)
 287          fm uniq_elt_pairs
 288 \end{code}
 289
 290 Now ways of removing things from UniqFM.
 291
 292 \begin{code}
 293 delListFromUFM fm lst = foldl delFromUFM fm lst
 294
 295 delFromUFM          fm key = delete fm (getKeyFastInt (getUnique key))
 296 delFromUFM_Directly fm u   = delete fm (getKeyFastInt u)
 297
 298 delete :: UniqFM a -> FastInt -> UniqFM a
 299 delete EmptyUFM _   = EmptyUFM
 300 delete fm       key = del_ele fm
 301   where
 302     del_ele :: UniqFM a -> UniqFM a
 303
 304     del_ele lf@(LeafUFM j _)
 305       | j ==# key       = EmptyUFM
 306       | otherwise       = lf    -- no delete!
 307
 308     del_ele (NodeUFM j p t1 t2)
 309       | j ># key
 310       = mkSLNodeUFM (NodeUFMData j p) (del_ele t1) t2
 311       | otherwise
 312       = mkLSNodeUFM (NodeUFMData j p) t1 (del_ele t2)
 313
 314     del_ele _ = panic "Found EmptyUFM FM when rec-deleting"
 315 \end{code}
 316
 317 Now ways of adding two UniqFM's together.
 318
 319 \begin{code}
 320 plusUFM tr1 tr2 = plusUFM_C use_snd tr1 tr2
 321
 322 plusUFM_C _ EmptyUFM tr = tr
 323 plusUFM_C _ tr EmptyUFM = tr
 324 plusUFM_C f fm1 fm2     = mix_trees fm1 fm2
 325     where
 326         mix_trees (LeafUFM i a) t2 = insert_ele (flip f) t2 i a
 327         mix_trees t1 (LeafUFM i a) = insert_ele f        t1 i a
 328
 329         mix_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 330           = mix_branches
 331                 (ask_about_common_ancestor
 332                         (NodeUFMData j p)
 333                         (NodeUFMData j' p'))
 334           where
 335                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 336                 --
 337                 --        j             j'                      (C j j')
 338                 --       / \    +      / \      ==>             /       \
 339                 --     t1   t2      t1'   t2'                  j         j'
 340                 --                                            / \       / \
 341                 --                                           t1  t2   t1'  t2'
 342                 -- Fast, Ehh !
 343                 --
 344           mix_branches (NewRoot nd False)
 345                 = mkLLNodeUFM nd left_t right_t
 346           mix_branches (NewRoot nd True)
 347                 = mkLLNodeUFM nd right_t left_t
 348
 349                 -- Now, if j == j':
 350                 --
 351                 --        j             j'                       j
 352                 --       / \    +      / \      ==>             / \
 353                 --     t1   t2      t1'   t2'           t1 + t1'   t2 + t2'
 354                 --
 355           mix_branches (SameRoot)
 356                 = mkSSNodeUFM (NodeUFMData j p)
 357                         (mix_trees t1 t1')
 358                         (mix_trees t2 t2')
 359                 -- Now the 4 different other ways; all like this:
 360                 --
 361                 -- Given j >^ j' (and, say,  j > j')
 362                 --
 363                 --        j             j'                       j
 364                 --       / \    +      / \      ==>             / \
 365                 --     t1   t2      t1'   t2'                 t1   t2 + j'
 366                 --                                                     / \
 367                 --                                                   t1'  t2'
 368           mix_branches (LeftRoot Leftt) --  | trace "LL" True
 369             = mkSLNodeUFM
 370                 (NodeUFMData j p)
 371                 (mix_trees t1 right_t)
 372                 t2
 373
 374           mix_branches (LeftRoot Rightt) --  | trace "LR" True
 375             = mkLSNodeUFM
 376                 (NodeUFMData j p)
 377                 t1
 378                 (mix_trees t2 right_t)
 379
 380           mix_branches (RightRoot Leftt) --  | trace "RL" True
 381             = mkSLNodeUFM
 382                 (NodeUFMData j' p')
 383                 (mix_trees left_t t1')
 384                 t2'
 385
 386           mix_branches (RightRoot Rightt) --  | trace "RR" True
 387             = mkLSNodeUFM
 388                 (NodeUFMData j' p')
 389                 t1'
 390                 (mix_trees left_t t2')
 391
 392         mix_trees _ _ = panic "EmptyUFM found when inserting into plusInt"
 393 \end{code}
 394
 395 And ways of subtracting them. First the base cases,
 396 then the full D&C approach.
 397
 398 \begin{code}
 399 minusUFM EmptyUFM _  = EmptyUFM
 400 minusUFM t1 EmptyUFM = t1
 401 minusUFM fm1 fm2     = minus_trees fm1 fm2
 402     where
 403         --
 404         -- Notice the asymetry of subtraction
 405         --
 406         minus_trees lf@(LeafUFM i _a) t2 =
 407                 case lookUp t2 i of
 408                   Nothing -> lf
 409                   Just _ -> EmptyUFM
 410
 411         minus_trees t1 (LeafUFM i _) = delete t1 i
 412
 413         minus_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 414           = minus_branches
 415                 (ask_about_common_ancestor
 416                         (NodeUFMData j p)
 417                         (NodeUFMData j' p'))
 418           where
 419                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 420                 --
 421                 --        j             j'                 j
 422                 --       / \    +      / \      ==>       / \
 423                 --     t1   t2      t1'   t2'            t1  t2
 424                 --
 425                 --
 426                 -- Fast, Ehh !
 427                 --
 428           minus_branches (NewRoot _ _) = left_t
 429
 430                 -- Now, if j == j':
 431                 --
 432                 --        j             j'                       j
 433                 --       / \    +      / \      ==>             / \
 434                 --     t1   t2      t1'   t2'           t1 + t1'   t2 + t2'
 435                 --
 436           minus_branches (SameRoot)
 437                 = mkSSNodeUFM (NodeUFMData j p)
 438                         (minus_trees t1 t1')
 439                         (minus_trees t2 t2')
 440                 -- Now the 4 different other ways; all like this:
 441                 -- again, with asymatry
 442
 443                 --
 444                 -- The left is above the right
 445                 --
 446           minus_branches (LeftRoot Leftt)
 447             = mkSLNodeUFM
 448                 (NodeUFMData j p)
 449                 (minus_trees t1 right_t)
 450                 t2
 451           minus_branches (LeftRoot Rightt)
 452             = mkLSNodeUFM
 453                 (NodeUFMData j p)
 454                 t1
 455                 (minus_trees t2 right_t)
 456
 457                 --
 458                 -- The right is above the left
 459                 --
 460           minus_branches (RightRoot Leftt)
 461             = minus_trees left_t t1'
 462           minus_branches (RightRoot Rightt)
 463             = minus_trees left_t t2'
 464
 465         minus_trees _ _ = panic "EmptyUFM found when insering into plusInt"
 466 \end{code}
 467
 468 And taking the intersection of two UniqFM's.
 469
 470 \begin{code}
 471 intersectUFM  t1 t2 = intersectUFM_C use_snd t1 t2
 472 intersectsUFM t1 t2 = isNullUFM (intersectUFM_C (\ _ _ -> error "urk") t1 t2)
 473
 474 intersectUFM_C _ EmptyUFM _ = EmptyUFM
 475 intersectUFM_C _ _ EmptyUFM = EmptyUFM
 476 intersectUFM_C f fm1 fm2    = intersect_trees fm1 fm2
 477     where
 478         intersect_trees (LeafUFM i a) t2 =
 479                 case lookUp t2 i of
 480                   Nothing -> EmptyUFM
 481                   Just b -> mkLeafUFM i (f a b)
 482
 483         intersect_trees t1 (LeafUFM i a) =
 484                 case lookUp t1 i of
 485                   Nothing -> EmptyUFM
 486                   Just b -> mkLeafUFM i (f b a)
 487
 488         intersect_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 489           = intersect_branches
 490                 (ask_about_common_ancestor
 491                         (NodeUFMData j p)
 492                         (NodeUFMData j' p'))
 493           where
 494                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 495                 --
 496                 --        j             j'
 497                 --       / \    +      / \      ==>             EmptyUFM
 498                 --     t1   t2      t1'   t2'
 499                 --
 500                 -- Fast, Ehh !
 501                 --
 502           intersect_branches (NewRoot _nd _) = EmptyUFM
 503
 504                 -- Now, if j == j':
 505                 --
 506                 --        j             j'                       j
 507                 --       / \    +      / \      ==>             / \
 508                 --     t1   t2      t1'   t2'           t1 x t1'   t2 x t2'
 509                 --
 510           intersect_branches (SameRoot)
 511                 = mkSSNodeUFM (NodeUFMData j p)
 512                         (intersect_trees t1 t1')
 513                         (intersect_trees t2 t2')
 514                 -- Now the 4 different other ways; all like this:
 515                 --
 516                 -- Given j >^ j' (and, say,  j > j')
 517                 --
 518                 --        j             j'                     t2 + j'
 519                 --       / \    +      / \      ==>                / \
 520                 --     t1   t2      t1'   t2'                    t1'  t2'
 521                 --
 522                 -- This does cut down the search space quite a bit.
 523
 524           intersect_branches (LeftRoot Leftt)
 525             = intersect_trees t1 right_t
 526           intersect_branches (LeftRoot Rightt)
 527             = intersect_trees t2 right_t
 528           intersect_branches (RightRoot Leftt)
 529             = intersect_trees left_t t1'
 530           intersect_branches (RightRoot Rightt)
 531             = intersect_trees left_t t2'
 532
 533         intersect_trees _ _ = panic ("EmptyUFM found when intersecting trees")
 534 \end{code}
 535
 536 Now the usual set of `collection' operators, like map, fold, etc.
 537
 538 \begin{code}
 539 foldUFM f a (NodeUFM _ _ t1 t2) = foldUFM f (foldUFM f a t2) t1
 540 foldUFM f a (LeafUFM _ obj)     = f obj a
 541 foldUFM _ a EmptyUFM            = a
 542 \end{code}
 543
 544 \begin{code}
 545 mapUFM _fn EmptyUFM   = EmptyUFM
 546 mapUFM  fn fm         = map_tree fn fm
 547
 548 filterUFM _fn EmptyUFM = EmptyUFM
 549 filterUFM  fn fm       = filter_tree (\_ e -> fn e) fm
 550
 551 filterUFM_Directly _fn EmptyUFM = EmptyUFM
 552 filterUFM_Directly  fn fm       = filter_tree pred fm
 553         where
 554           pred i e = fn (mkUniqueGrimily (iBox i)) e
 555 \end{code}
 556
 557 Note, this takes a long time, O(n), but
 558 because we dont want to do this very often, we put up with this.
 559 O'rable, but how often do we look at the size of
 560 a finite map?
 561
 562 \begin{code}
 563 sizeUFM EmptyUFM            = 0
 564 sizeUFM (NodeUFM _ _ t1 t2) = sizeUFM t1 + sizeUFM t2
 565 sizeUFM (LeafUFM _ _)       = 1
 566
 567 isNullUFM EmptyUFM = True
 568 isNullUFM _        = False
 569
 570 -- hashing is used in VarSet.uniqAway, and should be fast
 571 -- We use a cheap and cheerful method for now
 572 hashUFM EmptyUFM          = 0
 573 hashUFM (NodeUFM n _ _ _) = iBox n
 574 hashUFM (LeafUFM n _)     = iBox n
 575 \end{code}
 576
 577 looking up in a hurry is the {\em whole point} of this binary tree lark.
 578 Lookup up a binary tree is easy (and fast).
 579
 580 \begin{code}
 581 elemUFM          key fm = maybeToBool (lookupUFM fm key)
 582 elemUFM_Directly key fm = maybeToBool (lookupUFM_Directly fm key)
 583
 584 lookupUFM          fm key = lookUp fm (getKeyFastInt (getUnique key))
 585 lookupUFM_Directly fm key = lookUp fm (getKeyFastInt key)
 586
 587 lookupWithDefaultUFM fm deflt key
 588   = case lookUp fm (getKeyFastInt (getUnique key)) of
 589       Nothing  -> deflt
 590       Just elt -> elt
 591
 592 lookupWithDefaultUFM_Directly fm deflt key
 593   = case lookUp fm (getKeyFastInt key) of
 594       Nothing  -> deflt
 595       Just elt -> elt
 596
 597 lookUp :: UniqFM a -> FastInt -> Maybe a
 598 lookUp EmptyUFM _   = Nothing
 599 lookUp fm i         = lookup_tree fm
 600   where
 601         lookup_tree :: UniqFM a -> Maybe a
 602
 603         lookup_tree (LeafUFM j b)
 604           | j ==# i     = Just b
 605           | otherwise   = Nothing
 606         lookup_tree (NodeUFM j _ t1 t2)
 607           | j ># i      = lookup_tree t1
 608           | otherwise   = lookup_tree t2
 609
 610         lookup_tree EmptyUFM = panic "lookup Failed"
 611 \end{code}
 612
 613 folds are *wonderful* things.
 614
 615 \begin{code}
 616 eltsUFM   fm = foldUFM (:) [] fm
 617 keysUFM   fm = foldUFM_Directly (\u _ l -> u      : l) [] fm
 618 ufmToList fm = foldUFM_Directly (\u e l -> (u, e) : l) [] fm
 619 foldUFM_Directly f = fold_tree (\iu e a -> f (mkUniqueGrimily (iBox iu)) e a)
 620
 621 fold_tree :: (FastInt -> elt -> a -> a) -> a -> UniqFM elt -> a
 622 fold_tree f a (NodeUFM _ _ t1 t2) = fold_tree f (fold_tree f a t2) t1
 623 fold_tree f a (LeafUFM iu obj)    = f iu obj a
 624 fold_tree _ a EmptyUFM            = a
 625 \end{code}
 626
 627 %************************************************************************
 628 %*                                                                      *
 629 \subsubsection{The @UniqFM@ type, and its functions}
 630 %*                                                                      *
 631 %************************************************************************
 632
 633 You should always use these to build the tree.
 634 There are 4 versions of mkNodeUFM, depending on
 635 the strictness of the two sub-tree arguments.
 636 The strictness is used *both* to prune out
 637 empty trees, *and* to improve performance,
 638 stoping needless thunks lying around.
 639 The rule of thumb (from experence with these trees)
 640 is make thunks strict, but data structures lazy.
 641 If in doubt, use mkSSNodeUFM, which has the `strongest'
 642 functionality, but may do a few needless evaluations.
 643
 644 \begin{code}
 645 mkLeafUFM :: FastInt -> a -> UniqFM a
 646 mkLeafUFM i a =
 647   ASSERT (iBox i >= 0) -- Note [Uniques must be positive]
 648   LeafUFM i a
 649
 650 -- The *ONLY* ways of building a NodeUFM.
 651
 652 mkSSNodeUFM, mkSLNodeUFM, mkLSNodeUFM, mkLLNodeUFM ::
 653     NodeUFMData -> UniqFM a -> UniqFM a -> UniqFM a
 654
 655 mkSSNodeUFM (NodeUFMData _ _) EmptyUFM t2 = t2
 656 mkSSNodeUFM (NodeUFMData _ _) t1 EmptyUFM = t1
 657 mkSSNodeUFM (NodeUFMData j p) t1 t2
 658   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 659     NodeUFM j p t1 t2
 660
 661 mkSLNodeUFM (NodeUFMData _ _) EmptyUFM t2 = t2
 662 mkSLNodeUFM (NodeUFMData j p) t1 t2
 663   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 664     NodeUFM j p t1 t2
 665
 666 mkLSNodeUFM (NodeUFMData _ _) t1 EmptyUFM = t1
 667 mkLSNodeUFM (NodeUFMData j p) t1 t2
 668   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 669     NodeUFM j p t1 t2
 670
 671 mkLLNodeUFM (NodeUFMData j p) t1 t2
 672   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 673     NodeUFM j p t1 t2
 674
 675 correctNodeUFM
 676         :: Int
 677         -> Int
 678         -> UniqFM a
 679         -> UniqFM a
 680         -> Bool
 681
 682 correctNodeUFM j p t1 t2
 683   = correct (j-p) (j-1) p t1 && correct j ((j-1)+p) p t2
 684   where
 685     correct low high _ (LeafUFM i _)
 686       = low <= iBox i && iBox i <= high
 687     correct low high above_p (NodeUFM j p _ _)
 688       = low <= iBox j && iBox j <= high && above_p > iBox p
 689     correct _ _ _ EmptyUFM = panic "EmptyUFM stored inside a tree"
 690 \end{code}
 691
 692 Note: doing SAT on this by hand seems to make it worse. Todo: Investigate,
 693 and if necessary do $\lambda$ lifting on our functions that are bound.
 694
 695 \begin{code}
 696 insert_ele
 697         :: (a -> a -> a)        -- old -> new -> result
 698         -> UniqFM a
 699         -> FastInt
 700         -> a
 701         -> UniqFM a
 702
 703 insert_ele _f EmptyUFM i new = mkLeafUFM i new
 704
 705 insert_ele  f (LeafUFM j old) i new
 706   | j ># i =
 707           mkLLNodeUFM (getCommonNodeUFMData
 708                           (indexToRoot i)
 709                           (indexToRoot j))
 710                  (mkLeafUFM i new)
 711                  (mkLeafUFM j old)
 712   | j ==# i  = mkLeafUFM j $ f old new
 713   | otherwise =
 714           mkLLNodeUFM (getCommonNodeUFMData
 715                           (indexToRoot i)
 716                           (indexToRoot j))
 717                  (mkLeafUFM j old)
 718                  (mkLeafUFM i new)
 719
 720 insert_ele f n@(NodeUFM j p t1 t2) i a
 721   | i <# j
 722     = if (i >=# (j -# p))
 723       then mkSLNodeUFM (NodeUFMData j p) (insert_ele f t1 i a) t2
 724       else mkLLNodeUFM (getCommonNodeUFMData
 725                           (indexToRoot i)
 726                           ((NodeUFMData j p)))
 727                   (mkLeafUFM i a)
 728                   n
 729   | otherwise
 730     = if (i <=# ((j -# _ILIT(1)) +# p))
 731       then mkLSNodeUFM (NodeUFMData j p) t1 (insert_ele f t2 i a)
 732       else mkLLNodeUFM (getCommonNodeUFMData
 733                           (indexToRoot i)
 734                           ((NodeUFMData j p)))
 735                   n
 736                   (mkLeafUFM i a)
 737 \end{code}
 738
 739
 740
 741 \begin{code}
 742 map_tree :: (a -> b) -> UniqFM a -> UniqFM b
 743 map_tree f (NodeUFM j p t1 t2)
 744   = mkLLNodeUFM (NodeUFMData j p) (map_tree f t1) (map_tree f t2)
 745         -- NB. lazy! we know the tree is well-formed.
 746 map_tree f (LeafUFM i obj)
 747   = mkLeafUFM i (f obj)
 748 map_tree _ _ = panic "map_tree failed"
 749 \end{code}
 750
 751 \begin{code}
 752 filter_tree :: (FastInt -> a -> Bool) -> UniqFM a -> UniqFM a
 753 filter_tree f (NodeUFM j p t1 t2)
 754   = mkSSNodeUFM (NodeUFMData j p) (filter_tree f t1) (filter_tree f t2)
 755
 756 filter_tree f lf@(LeafUFM i obj)
 757   | f i obj = lf
 758   | otherwise = EmptyUFM
 759 filter_tree _ _ = panic "filter_tree failed"
 760 \end{code}
 761
 762 %************************************************************************
 763 %*                                                                      *
 764 \subsubsection{The @UniqFM@ type, and signatures for the functions}
 765 %*                                                                      *
 766 %************************************************************************
 767
 768 Now some Utilities;
 769
 770 This is the information that is held inside a NodeUFM, packaged up for
 771 consumer use.
 772
 773 \begin{code}
 774 data NodeUFMData
 775   = NodeUFMData FastInt
 776                 FastInt
 777 \end{code}
 778
 779 This is the information used when computing new NodeUFMs.
 780
 781 \begin{code}
 782 data Side = Leftt | Rightt -- NB: avoid 1.3 names "Left" and "Right"
 783 data CommonRoot
 784   = LeftRoot  Side      -- which side is the right down ?
 785   | RightRoot Side      -- which side is the left down ?
 786   | SameRoot            -- they are the same !
 787   | NewRoot NodeUFMData -- here's the new, common, root
 788             Bool        -- do you need to swap left and right ?
 789 \end{code}
 790
 791 This specifies the relationship between NodeUFMData and CalcNodeUFMData.
 792
 793 \begin{code}
 794 indexToRoot :: FastInt -> NodeUFMData
 795
 796 indexToRoot i
 797   = NodeUFMData ((shiftL1 (shiftR1 i)) +# _ILIT(1)) (_ILIT(1))
 798
 799 getCommonNodeUFMData :: NodeUFMData -> NodeUFMData -> NodeUFMData
 800
 801 getCommonNodeUFMData (NodeUFMData i p) (NodeUFMData i2 p2)
 802   | p ==# p2    = getCommonNodeUFMData_ p j j2
 803   | p <# p2     = getCommonNodeUFMData_ p2 (j `quotFastInt` (p2 `quotFastInt` p)) j2
 804   | otherwise   = getCommonNodeUFMData_ p j (j2 `quotFastInt` (p `quotFastInt` p2))
 805   where
 806     !j  = i  `quotFastInt` (shiftL1 p)
 807     !j2 = i2 `quotFastInt` (shiftL1 p2)
 808
 809     getCommonNodeUFMData_ :: FastInt -> FastInt -> FastInt -> NodeUFMData
 810
 811     getCommonNodeUFMData_ p j j_
 812       | j ==# j_
 813       = NodeUFMData (((shiftL1 j) +# _ILIT(1)) *# p) p
 814       | otherwise
 815       = getCommonNodeUFMData_ (shiftL1 p) (shiftR1 j) (shiftR1 j_)
 816
 817 ask_about_common_ancestor :: NodeUFMData -> NodeUFMData -> CommonRoot
 818
 819 ask_about_common_ancestor x@(NodeUFMData j _p) y@(NodeUFMData j2 _p2)
 820   | j ==# j2 = SameRoot
 821   | otherwise
 822   = case getCommonNodeUFMData x y of
 823       nd@(NodeUFMData j3 _p3)
 824         | j3 ==# j  -> LeftRoot (decideSide (j ># j2))
 825         | j3 ==# j2 -> RightRoot (decideSide (j <# j2))
 826         | otherwise   -> NewRoot nd (j ># j2)
 827     where
 828         decideSide :: Bool -> Side
 829         decideSide True  = Leftt
 830         decideSide False = Rightt
 831 \end{code}
 832
 833 This might be better in Util.lhs ?
 834
 835
 836 Now the bit twiddling functions.
 837 \begin{code}
 838 shiftL1 :: FastInt -> FastInt
 839 shiftR1 :: FastInt -> FastInt
 840
 841 {-# INLINE shiftL1 #-}
 842 {-# INLINE shiftR1 #-}
 843
 844 shiftL1 n = n `shiftLFastInt` _ILIT(1)
 845 shiftR1 n = n `shiftR_FastInt` _ILIT(1)
 846 \end{code}
 847
 848 \begin{code}
 849 use_snd :: a -> b -> b
 850 use_snd _ b = b
 851 \end{code}
 852
 853 {- Note [Uniques must be positive]
 854 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 855 The getCommonNodeUFMData function assumes that the nodes use
 856 positive uniques. Specifically, the inner `loop' shifts the
 857 low bits out of two uniques until the shifted uniques are the same.
 858 At the same time, it computes a new delta, by shifting
 859 to the left.
 860
 861 The failure case I (JPD) encountered:
 862 If one of the uniques is negative, the shifting may continue
 863 until all 64 bits have been shifted out, resulting in a new delta
 864 of 0, which is wrong and can trigger later assertion failures.
 865
 866 Where do the negative uniques come from? Both Simom M and
 867 I have run into this problem when hashing a data structure.
 868 In both cases, we have avoided the problem by ensuring that
 869 the hashes remain positive.
 870 -}