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