ghc/compiler/utils/UniqFM.lhs

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