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