compiler/utils/UniqFM.lhs

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