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