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