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 -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/Commentary/CodingStyle#Warnings
  21 -- for details
  22
  23 module UniqFM (
  24         UniqFM(..),     -- abstract type
  25                         -- (de-abstracted for MachRegs.trivColorable optimisation BL 2007/09)
  26
  27         emptyUFM,
  28         unitUFM,
  29         unitDirectlyUFM,
  30         listToUFM,
  31         listToUFM_Directly,
  32         addToUFM,addToUFM_C,addToUFM_Acc,
  33         addListToUFM,addListToUFM_C,
  34         addToUFM_Directly,
  35         addListToUFM_Directly,
  36         delFromUFM,
  37         delFromUFM_Directly,
  38         delListFromUFM,
  39         plusUFM,
  40         plusUFM_C,
  41         minusUFM,
  42         intersectsUFM,
  43         intersectUFM,
  44         intersectUFM_C,
  45         foldUFM,
  46         mapUFM,
  47         elemUFM, elemUFM_Directly,
  48         filterUFM, filterUFM_Directly,
  49         sizeUFM,
  50         hashUFM,
  51         isNullUFM,
  52         lookupUFM, lookupUFM_Directly,
  53         lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
  54         eltsUFM, keysUFM,
  55         ufmToList
  56     ) where
  57
  58 #include "HsVersions.h"
  59
  60 import Unique           ( Uniquable(..), Unique, getKey#, mkUniqueGrimily )
  61 import Maybes           ( maybeToBool )
  62 import FastTypes
  63 import Outputable
  64
  65 import GHC.Exts         -- Lots of Int# operations
  66 \end{code}
  67
  68 %************************************************************************
  69 %*                                                                      *
  70 \subsection{The @UniqFM@ type, and signatures for the functions}
  71 %*                                                                      *
  72 %************************************************************************
  73
  74 We use @FiniteMaps@, with a (@getUnique@-able) @Unique@ as ``key''.
  75
  76 \begin{code}
  77 emptyUFM        :: UniqFM elt
  78 isNullUFM       :: UniqFM elt -> Bool
  79 unitUFM         :: Uniquable key => key -> elt -> UniqFM elt
  80 unitDirectlyUFM -- got the Unique already
  81                 :: Unique -> elt -> UniqFM elt
  82 listToUFM       :: Uniquable key => [(key,elt)] -> UniqFM elt
  83 listToUFM_Directly
  84                 :: [(Unique, elt)] -> UniqFM elt
  85
  86 addToUFM        :: Uniquable key => UniqFM elt -> key -> elt  -> UniqFM elt
  87 addListToUFM    :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
  88 addToUFM_Directly
  89                 :: UniqFM elt -> Unique -> elt -> UniqFM elt
  90
  91 addToUFM_C      :: Uniquable key => (elt -> elt -> elt) -- old -> new -> result
  92                            -> UniqFM elt                -- old
  93                            -> key -> elt                -- new
  94                            -> UniqFM elt                -- result
  95
  96 addToUFM_Acc    :: Uniquable key =>
  97                               (elt -> elts -> elts)     -- Add to existing
  98                            -> (elt -> elts)             -- New element
  99                            -> UniqFM elts               -- old
 100                            -> key -> elt                -- new
 101                            -> UniqFM elts               -- result
 102
 103 addListToUFM_C  :: Uniquable key => (elt -> elt -> elt)
 104                            -> UniqFM elt -> [(key,elt)]
 105                            -> UniqFM elt
 106
 107 delFromUFM      :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
 108 delListFromUFM  :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
 109 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
 110
 111 plusUFM         :: UniqFM elt -> UniqFM elt -> UniqFM elt
 112
 113 plusUFM_C       :: (elt -> elt -> elt)
 114                 -> UniqFM elt -> UniqFM elt -> UniqFM elt
 115
 116 minusUFM        :: UniqFM elt1 -> UniqFM elt2 -> UniqFM elt1
 117
 118 intersectUFM    :: UniqFM elt -> UniqFM elt -> UniqFM elt
 119 intersectUFM_C  :: (elt1 -> elt2 -> elt3)
 120                 -> UniqFM elt1 -> UniqFM elt2 -> UniqFM elt3
 121 intersectsUFM   :: UniqFM elt1 -> UniqFM elt2 -> Bool
 122
 123 foldUFM         :: (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 @UniqFM a@ is a mapping from Unique to a.
 206
 207 First, the DataType itself; which is either a Node, a Leaf, or an Empty.
 208
 209 \begin{code}
 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 (getKey# (getUnique key)) elt
 245 unitDirectlyUFM key elt = mkLeafUFM (getKey# 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 \end{code}
 253
 254 Now ways of adding things to UniqFMs.
 255
 256 There is an alternative version of @addListToUFM_C@, that uses @plusUFM@,
 257 but the semantics of this operation demands a linear insertion;
 258 perhaps the version without the combinator function
 259 could be optimised using it.
 260
 261 \begin{code}
 262 addToUFM fm key elt = addToUFM_C use_snd fm key elt
 263
 264 addToUFM_Directly fm u elt = insert_ele use_snd fm (getKey# u) elt
 265
 266 addToUFM_C combiner fm key elt
 267   = insert_ele combiner fm (getKey# (getUnique key)) elt
 268
 269 addToUFM_Acc add unit fm key item
 270   = insert_ele combiner fm (getKey# (getUnique key)) (unit item)
 271   where
 272     combiner old _unit_item = add item old
 273
 274 addListToUFM fm key_elt_pairs = addListToUFM_C use_snd fm key_elt_pairs
 275 addListToUFM_Directly fm uniq_elt_pairs = addListToUFM_directly_C use_snd fm uniq_elt_pairs
 276
 277 addListToUFM_C combiner fm key_elt_pairs
 278  = foldl (\ fm (k, e) -> insert_ele combiner fm (getKey# (getUnique k)) e)
 279          fm key_elt_pairs
 280
 281 addListToUFM_directly_C combiner fm uniq_elt_pairs
 282  = foldl (\ fm (k, e) -> insert_ele combiner fm (getKey# k) e)
 283          fm uniq_elt_pairs
 284 \end{code}
 285
 286 Now ways of removing things from UniqFM.
 287
 288 \begin{code}
 289 delListFromUFM fm lst = foldl delFromUFM fm lst
 290
 291 delFromUFM          fm key = delete fm (getKey# (getUnique key))
 292 delFromUFM_Directly fm u   = delete fm (getKey# u)
 293
 294 delete EmptyUFM _   = EmptyUFM
 295 delete fm       key = del_ele fm
 296   where
 297     del_ele :: UniqFM a -> UniqFM a
 298
 299     del_ele lf@(LeafUFM j _)
 300       | j ==# key       = EmptyUFM
 301       | otherwise       = lf    -- no delete!
 302
 303     del_ele nd@(NodeUFM j p t1 t2)
 304       | j ># key
 305       = mkSLNodeUFM (NodeUFMData j p) (del_ele t1) t2
 306       | otherwise
 307       = mkLSNodeUFM (NodeUFMData j p) t1 (del_ele t2)
 308
 309     del_ele _ = panic "Found EmptyUFM FM when rec-deleting"
 310 \end{code}
 311
 312 Now ways of adding two UniqFM's together.
 313
 314 \begin{code}
 315 plusUFM tr1 tr2 = plusUFM_C use_snd tr1 tr2
 316
 317 plusUFM_C f EmptyUFM tr = tr
 318 plusUFM_C f tr EmptyUFM = tr
 319 plusUFM_C f fm1 fm2     = mix_trees fm1 fm2
 320     where
 321         mix_trees (LeafUFM i a) t2 = insert_ele (flip f) t2 i a
 322         mix_trees t1 (LeafUFM i a) = insert_ele f        t1 i a
 323
 324         mix_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 325           = mix_branches
 326                 (ask_about_common_ancestor
 327                         (NodeUFMData j p)
 328                         (NodeUFMData j' p'))
 329           where
 330                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 331                 --
 332                 --        j             j'                      (C j j')
 333                 --       / \    +      / \      ==>             /       \
 334                 --     t1   t2      t1'   t2'                  j         j'
 335                 --                                            / \       / \
 336                 --                                           t1  t2   t1'  t2'
 337                 -- Fast, Ehh !
 338                 --
 339           mix_branches (NewRoot nd False)
 340                 = mkLLNodeUFM nd left_t right_t
 341           mix_branches (NewRoot nd True)
 342                 = mkLLNodeUFM nd right_t left_t
 343
 344                 -- Now, if j == j':
 345                 --
 346                 --        j             j'                       j
 347                 --       / \    +      / \      ==>             / \
 348                 --     t1   t2      t1'   t2'           t1 + t1'   t2 + t2'
 349                 --
 350           mix_branches (SameRoot)
 351                 = mkSSNodeUFM (NodeUFMData j p)
 352                         (mix_trees t1 t1')
 353                         (mix_trees t2 t2')
 354                 -- Now the 4 different other ways; all like this:
 355                 --
 356                 -- Given j >^ j' (and, say,  j > j')
 357                 --
 358                 --        j             j'                       j
 359                 --       / \    +      / \      ==>             / \
 360                 --     t1   t2      t1'   t2'                 t1   t2 + j'
 361                 --                                                     / \
 362                 --                                                   t1'  t2'
 363           mix_branches (LeftRoot Leftt) --  | trace "LL" True
 364             = mkSLNodeUFM
 365                 (NodeUFMData j p)
 366                 (mix_trees t1 right_t)
 367                 t2
 368
 369           mix_branches (LeftRoot Rightt) --  | trace "LR" True
 370             = mkLSNodeUFM
 371                 (NodeUFMData j p)
 372                 t1
 373                 (mix_trees t2 right_t)
 374
 375           mix_branches (RightRoot Leftt) --  | trace "RL" True
 376             = mkSLNodeUFM
 377                 (NodeUFMData j' p')
 378                 (mix_trees left_t t1')
 379                 t2'
 380
 381           mix_branches (RightRoot Rightt) --  | trace "RR" True
 382             = mkLSNodeUFM
 383                 (NodeUFMData j' p')
 384                 t1'
 385                 (mix_trees left_t t2')
 386
 387         mix_trees _ _ = panic "EmptyUFM found when inserting into plusInt"
 388 \end{code}
 389
 390 And ways of subtracting them. First the base cases,
 391 then the full D&C approach.
 392
 393 \begin{code}
 394 minusUFM EmptyUFM _  = EmptyUFM
 395 minusUFM t1 EmptyUFM = t1
 396 minusUFM fm1 fm2     = minus_trees fm1 fm2
 397     where
 398         --
 399         -- Notice the asymetry of subtraction
 400         --
 401         minus_trees lf@(LeafUFM i a) t2 =
 402                 case lookUp t2 i of
 403                   Nothing -> lf
 404                   Just b -> EmptyUFM
 405
 406         minus_trees t1 (LeafUFM i _) = delete t1 i
 407
 408         minus_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 409           = minus_branches
 410                 (ask_about_common_ancestor
 411                         (NodeUFMData j p)
 412                         (NodeUFMData j' p'))
 413           where
 414                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 415                 --
 416                 --        j             j'                 j
 417                 --       / \    +      / \      ==>       / \
 418                 --     t1   t2      t1'   t2'            t1  t2
 419                 --
 420                 --
 421                 -- Fast, Ehh !
 422                 --
 423           minus_branches (NewRoot nd _) = left_t
 424
 425                 -- Now, if j == j':
 426                 --
 427                 --        j             j'                       j
 428                 --       / \    +      / \      ==>             / \
 429                 --     t1   t2      t1'   t2'           t1 + t1'   t2 + t2'
 430                 --
 431           minus_branches (SameRoot)
 432                 = mkSSNodeUFM (NodeUFMData j p)
 433                         (minus_trees t1 t1')
 434                         (minus_trees t2 t2')
 435                 -- Now the 4 different other ways; all like this:
 436                 -- again, with asymatry
 437
 438                 --
 439                 -- The left is above the right
 440                 --
 441           minus_branches (LeftRoot Leftt)
 442             = mkSLNodeUFM
 443                 (NodeUFMData j p)
 444                 (minus_trees t1 right_t)
 445                 t2
 446           minus_branches (LeftRoot Rightt)
 447             = mkLSNodeUFM
 448                 (NodeUFMData j p)
 449                 t1
 450                 (minus_trees t2 right_t)
 451
 452                 --
 453                 -- The right is above the left
 454                 --
 455           minus_branches (RightRoot Leftt)
 456             = minus_trees left_t t1'
 457           minus_branches (RightRoot Rightt)
 458             = minus_trees left_t t2'
 459
 460         minus_trees _ _ = panic "EmptyUFM found when insering into plusInt"
 461 \end{code}
 462
 463 And taking the intersection of two UniqFM's.
 464
 465 \begin{code}
 466 intersectUFM  t1 t2 = intersectUFM_C use_snd t1 t2
 467 intersectsUFM t1 t2 = isNullUFM (intersectUFM_C (\ _ _ -> error "urk") t1 t2)
 468
 469 intersectUFM_C f EmptyUFM _ = EmptyUFM
 470 intersectUFM_C f _ EmptyUFM = EmptyUFM
 471 intersectUFM_C f fm1 fm2    = intersect_trees fm1 fm2
 472     where
 473         intersect_trees (LeafUFM i a) t2 =
 474                 case lookUp t2 i of
 475                   Nothing -> EmptyUFM
 476                   Just b -> mkLeafUFM i (f a b)
 477
 478         intersect_trees t1 (LeafUFM i a) =
 479                 case lookUp t1 i of
 480                   Nothing -> EmptyUFM
 481                   Just b -> mkLeafUFM i (f b a)
 482
 483         intersect_trees left_t@(NodeUFM j p t1 t2) right_t@(NodeUFM j' p' t1' t2')
 484           = intersect_branches
 485                 (ask_about_common_ancestor
 486                         (NodeUFMData j p)
 487                         (NodeUFMData j' p'))
 488           where
 489                 -- Given a disjoint j,j' (p >^ p' && p' >^ p):
 490                 --
 491                 --        j             j'
 492                 --       / \    +      / \      ==>             EmptyUFM
 493                 --     t1   t2      t1'   t2'
 494                 --
 495                 -- Fast, Ehh !
 496                 --
 497           intersect_branches (NewRoot nd _) = EmptyUFM
 498
 499                 -- Now, if j == j':
 500                 --
 501                 --        j             j'                       j
 502                 --       / \    +      / \      ==>             / \
 503                 --     t1   t2      t1'   t2'           t1 x t1'   t2 x t2'
 504                 --
 505           intersect_branches (SameRoot)
 506                 = mkSSNodeUFM (NodeUFMData j p)
 507                         (intersect_trees t1 t1')
 508                         (intersect_trees t2 t2')
 509                 -- Now the 4 different other ways; all like this:
 510                 --
 511                 -- Given j >^ j' (and, say,  j > j')
 512                 --
 513                 --        j             j'                     t2 + j'
 514                 --       / \    +      / \      ==>                / \
 515                 --     t1   t2      t1'   t2'                    t1'  t2'
 516                 --
 517                 -- This does cut down the search space quite a bit.
 518
 519           intersect_branches (LeftRoot Leftt)
 520             = intersect_trees t1 right_t
 521           intersect_branches (LeftRoot Rightt)
 522             = intersect_trees t2 right_t
 523           intersect_branches (RightRoot Leftt)
 524             = intersect_trees left_t t1'
 525           intersect_branches (RightRoot Rightt)
 526             = intersect_trees left_t t2'
 527
 528         intersect_trees x y = panic ("EmptyUFM found when intersecting trees")
 529 \end{code}
 530
 531 Now the usual set of `collection' operators, like map, fold, etc.
 532
 533 \begin{code}
 534 foldUFM f a (NodeUFM _ _ t1 t2) = foldUFM f (foldUFM f a t2) t1
 535 foldUFM f a (LeafUFM _ obj)     = f obj a
 536 foldUFM f a EmptyUFM            = a
 537 \end{code}
 538
 539 \begin{code}
 540 mapUFM fn EmptyUFM    = EmptyUFM
 541 mapUFM fn fm          = map_tree fn fm
 542
 543 filterUFM fn EmptyUFM = EmptyUFM
 544 filterUFM fn fm       = filter_tree pred fm
 545         where
 546           pred (i::FastInt) e = fn e
 547
 548 filterUFM_Directly fn EmptyUFM = EmptyUFM
 549 filterUFM_Directly fn fm       = filter_tree pred fm
 550         where
 551           pred i e = fn (mkUniqueGrimily (iBox i)) e
 552 \end{code}
 553
 554 Note, this takes a long time, O(n), but
 555 because we dont want to do this very often, we put up with this.
 556 O'rable, but how often do we look at the size of
 557 a finite map?
 558
 559 \begin{code}
 560 sizeUFM EmptyUFM            = 0
 561 sizeUFM (NodeUFM _ _ t1 t2) = sizeUFM t1 + sizeUFM t2
 562 sizeUFM (LeafUFM _ _)       = 1
 563
 564 isNullUFM EmptyUFM = True
 565 isNullUFM _        = False
 566
 567 -- hashing is used in VarSet.uniqAway, and should be fast
 568 -- We use a cheap and cheerful method for now
 569 hashUFM EmptyUFM          = 0
 570 hashUFM (NodeUFM n _ _ _) = iBox n
 571 hashUFM (LeafUFM n _)     = iBox n
 572 \end{code}
 573
 574 looking up in a hurry is the {\em whole point} of this binary tree lark.
 575 Lookup up a binary tree is easy (and fast).
 576
 577 \begin{code}
 578 elemUFM          key fm = maybeToBool (lookupUFM fm key)
 579 elemUFM_Directly key fm = maybeToBool (lookupUFM_Directly fm key)
 580
 581 lookupUFM          fm key = lookUp fm (getKey# (getUnique key))
 582 lookupUFM_Directly fm key = lookUp fm (getKey# key)
 583
 584 lookupWithDefaultUFM fm deflt key
 585   = case lookUp fm (getKey# (getUnique key)) of
 586       Nothing  -> deflt
 587       Just elt -> elt
 588
 589 lookupWithDefaultUFM_Directly fm deflt key
 590   = case lookUp fm (getKey# key) of
 591       Nothing  -> deflt
 592       Just elt -> elt
 593
 594 lookUp EmptyUFM _   = Nothing
 595 lookUp fm i         = lookup_tree fm
 596   where
 597         lookup_tree :: UniqFM a -> Maybe a
 598
 599         lookup_tree (LeafUFM j b)
 600           | j ==# i     = Just b
 601           | otherwise   = Nothing
 602         lookup_tree (NodeUFM j p t1 t2)
 603           | j ># i      = lookup_tree t1
 604           | otherwise   = lookup_tree t2
 605
 606         lookup_tree EmptyUFM = panic "lookup Failed"
 607 \end{code}
 608
 609 folds are *wonderful* things.
 610
 611 \begin{code}
 612 eltsUFM fm = foldUFM (:) [] fm
 613
 614 ufmToList fm = fold_tree (\ iu elt rest -> (mkUniqueGrimily (iBox iu), elt) : rest) [] fm
 615
 616 keysUFM fm = fold_tree (\ iu elt rest -> mkUniqueGrimily (iBox iu) : rest) [] fm
 617
 618 fold_tree f a (NodeUFM _ _ t1 t2) = fold_tree f (fold_tree f a t2) t1
 619 fold_tree f a (LeafUFM iu obj)    = f iu obj a
 620 fold_tree f a EmptyUFM            = a
 621 \end{code}
 622
 623 %************************************************************************
 624 %*                                                                      *
 625 \subsubsection{The @UniqFM@ type, and its functions}
 626 %*                                                                      *
 627 %************************************************************************
 628
 629 You should always use these to build the tree.
 630 There are 4 versions of mkNodeUFM, depending on
 631 the strictness of the two sub-tree arguments.
 632 The strictness is used *both* to prune out
 633 empty trees, *and* to improve performance,
 634 stoping needless thunks lying around.
 635 The rule of thumb (from experence with these trees)
 636 is make thunks strict, but data structures lazy.
 637 If in doubt, use mkSSNodeUFM, which has the `strongest'
 638 functionality, but may do a few needless evaluations.
 639
 640 \begin{code}
 641 mkLeafUFM :: FastInt -> a -> UniqFM a
 642 mkLeafUFM i a     = LeafUFM i a
 643
 644 -- The *ONLY* ways of building a NodeUFM.
 645
 646 mkSSNodeUFM (NodeUFMData j p) EmptyUFM t2 = t2
 647 mkSSNodeUFM (NodeUFMData j p) t1 EmptyUFM = t1
 648 mkSSNodeUFM (NodeUFMData j p) t1 t2
 649   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 650     NodeUFM j p t1 t2
 651
 652 mkSLNodeUFM (NodeUFMData j p) EmptyUFM t2 = t2
 653 mkSLNodeUFM (NodeUFMData j p) t1 t2
 654   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 655     NodeUFM j p t1 t2
 656
 657 mkLSNodeUFM (NodeUFMData j p) t1 EmptyUFM = t1
 658 mkLSNodeUFM (NodeUFMData j p) t1 t2
 659   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 660     NodeUFM j p t1 t2
 661
 662 mkLLNodeUFM (NodeUFMData j p) t1 t2
 663   = ASSERT(correctNodeUFM (iBox j) (iBox p) t1 t2)
 664     NodeUFM j p t1 t2
 665
 666 correctNodeUFM
 667         :: Int
 668         -> Int
 669         -> UniqFM a
 670         -> UniqFM a
 671         -> Bool
 672
 673 correctNodeUFM j p t1 t2
 674   = correct (j-p) (j-1) p t1 && correct j ((j-1)+p) p t2
 675   where
 676     correct low high _ (LeafUFM i _)
 677       = low <= iBox i && iBox i <= high
 678     correct low high above_p (NodeUFM j p _ _)
 679       = low <= iBox j && iBox j <= high && above_p > iBox p
 680     correct _ _ _ EmptyUFM = panic "EmptyUFM stored inside a tree"
 681 \end{code}
 682
 683 Note: doing SAT on this by hand seems to make it worse. Todo: Investigate,
 684 and if necessary do $\lambda$ lifting on our functions that are bound.
 685
 686 \begin{code}
 687 insert_ele
 688         :: (a -> a -> a)        -- old -> new -> result
 689         -> UniqFM a
 690         -> FastInt
 691         -> a
 692         -> UniqFM a
 693
 694 insert_ele f EmptyUFM i new = mkLeafUFM i new
 695
 696 insert_ele f (LeafUFM j old) i new
 697   | j ># i =
 698           mkLLNodeUFM (getCommonNodeUFMData
 699                           (indexToRoot i)
 700                           (indexToRoot j))
 701                  (mkLeafUFM i new)
 702                  (mkLeafUFM j old)
 703   | j ==# i  = mkLeafUFM j (f old new)
 704   | otherwise =
 705           mkLLNodeUFM (getCommonNodeUFMData
 706                           (indexToRoot i)
 707                           (indexToRoot j))
 708                  (mkLeafUFM j old)
 709                  (mkLeafUFM i new)
 710
 711 insert_ele f n@(NodeUFM j p t1 t2) i a
 712   | i <# j
 713     = if (i >=# (j -# p))
 714       then mkSLNodeUFM (NodeUFMData j p) (insert_ele f t1 i a) t2
 715       else mkLLNodeUFM (getCommonNodeUFMData
 716                           (indexToRoot i)
 717                           ((NodeUFMData j p)))
 718                   (mkLeafUFM i a)
 719                   n
 720   | otherwise
 721     = if (i <=# ((j -# _ILIT(1)) +# p))
 722       then mkLSNodeUFM (NodeUFMData j p) t1 (insert_ele f t2 i a)
 723       else mkLLNodeUFM (getCommonNodeUFMData
 724                           (indexToRoot i)
 725                           ((NodeUFMData j p)))
 726                   n
 727                   (mkLeafUFM i a)
 728 \end{code}
 729
 730
 731
 732 \begin{code}
 733 map_tree f (NodeUFM j p t1 t2)
 734   = mkLLNodeUFM (NodeUFMData j p) (map_tree f t1) (map_tree f t2)
 735         -- NB. lazy! we know the tree is well-formed.
 736 map_tree f (LeafUFM i obj)
 737   = mkLeafUFM i (f obj)
 738 map_tree f _ = panic "map_tree failed"
 739 \end{code}
 740
 741 \begin{code}
 742 filter_tree :: (FastInt -> a -> Bool) -> UniqFM a -> UniqFM a
 743 filter_tree f nd@(NodeUFM j p t1 t2)
 744   = mkSSNodeUFM (NodeUFMData j p) (filter_tree f t1) (filter_tree f t2)
 745
 746 filter_tree f lf@(LeafUFM i obj)
 747   | f i obj = lf
 748   | otherwise = EmptyUFM
 749 filter_tree f _ = panic "filter_tree failed"
 750 \end{code}
 751
 752 %************************************************************************
 753 %*                                                                      *
 754 \subsubsection{The @UniqFM@ type, and signatures for the functions}
 755 %*                                                                      *
 756 %************************************************************************
 757
 758 Now some Utilities;
 759
 760 This is the information that is held inside a NodeUFM, packaged up for
 761 consumer use.
 762
 763 \begin{code}
 764 data NodeUFMData
 765   = NodeUFMData FastInt
 766                 FastInt
 767 \end{code}
 768
 769 This is the information used when computing new NodeUFMs.
 770
 771 \begin{code}
 772 data Side = Leftt | Rightt -- NB: avoid 1.3 names "Left" and "Right"
 773 data CommonRoot
 774   = LeftRoot  Side      -- which side is the right down ?
 775   | RightRoot Side      -- which side is the left down ?
 776   | SameRoot            -- they are the same !
 777   | NewRoot NodeUFMData -- here's the new, common, root
 778             Bool        -- do you need to swap left and right ?
 779 \end{code}
 780
 781 This specifies the relationship between NodeUFMData and CalcNodeUFMData.
 782
 783 \begin{code}
 784 indexToRoot :: FastInt -> NodeUFMData
 785
 786 indexToRoot i
 787   = let
 788         l = (_ILIT(1) :: FastInt)
 789     in
 790     NodeUFMData (((i `shiftR_` l) `shiftL_` l) +# _ILIT(1)) l
 791
 792 getCommonNodeUFMData :: NodeUFMData -> NodeUFMData -> NodeUFMData
 793
 794 getCommonNodeUFMData (NodeUFMData i p) (NodeUFMData i2 p2)
 795   | p ==# p2    = getCommonNodeUFMData_ p j j2
 796   | p <# p2     = getCommonNodeUFMData_ p2 (j `quotFastInt` (p2 `quotFastInt` p)) j2
 797   | otherwise   = getCommonNodeUFMData_ p j (j2 `quotFastInt` (p `quotFastInt` p2))
 798   where
 799     l  = (_ILIT(1) :: FastInt)
 800     j  = i  `quotFastInt` (p  `shiftL_` l)
 801     j2 = i2 `quotFastInt` (p2 `shiftL_` l)
 802
 803     getCommonNodeUFMData_ :: FastInt -> FastInt -> FastInt -> NodeUFMData
 804
 805     getCommonNodeUFMData_ p j j_
 806       | j ==# j_
 807       = NodeUFMData (((j `shiftL_` l) +# l) *# p) p
 808       | otherwise
 809       = getCommonNodeUFMData_ (p `shiftL_`  l) (j `shiftR_` l) (j_ `shiftR_` l)
 810
 811 ask_about_common_ancestor :: NodeUFMData -> NodeUFMData -> CommonRoot
 812
 813 ask_about_common_ancestor x@(NodeUFMData j p) y@(NodeUFMData j2 p2)
 814   | j ==# j2 = SameRoot
 815   | otherwise
 816   = case getCommonNodeUFMData x y of
 817       nd@(NodeUFMData j3 p3)
 818         | j3 ==# j  -> LeftRoot (decideSide (j ># j2))
 819         | j3 ==# j2 -> RightRoot (decideSide (j <# j2))
 820         | otherwise   -> NewRoot nd (j ># j2)
 821     where
 822         decideSide :: Bool -> Side
 823         decideSide True  = Leftt
 824         decideSide False = Rightt
 825 \end{code}
 826
 827 This might be better in Util.lhs ?
 828
 829
 830 Now the bit twiddling functions.
 831 \begin{code}
 832 shiftL_ :: FastInt -> FastInt -> FastInt
 833 shiftR_ :: FastInt -> FastInt -> FastInt
 834
 835 #if __GLASGOW_HASKELL__
 836 {-# INLINE shiftL_ #-}
 837 {-# INLINE shiftR_ #-}
 838 shiftL_ n p = word2Int#((int2Word# n) `uncheckedShiftL#` p)
 839 shiftR_ n p = word2Int#((int2Word# n) `uncheckedShiftRL#` p)
 840
 841 #else /* not GHC */
 842 shiftL_ n p = n * (2 ^ p)
 843 shiftR_ n p = n `quot` (2 ^ p)
 844
 845 #endif /* not GHC */
 846 \end{code}
 847
 848 \begin{code}
 849 use_snd :: a -> b -> b
 850 use_snd a b = b
 851 \end{code}