ghc/compiler/utils/UniqFM.lhs

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