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