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