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     ) where
  45
  46 #include "HsVersions.h"
  47
  48 import {-# SOURCE #-} Name      ( Name )
  49
  50 import Unique           ( Uniquable(..), Unique, u2i, mkUniqueGrimily )
  51 import Panic
  52 import GlaExts          -- Lots of Int# operations
  53
  54 #if ! OMIT_NATIVE_CODEGEN
  55 #define IF_NCG(a) a
  56 #else
  57 #define IF_NCG(a) {--}
  58 #endif
  59 \end{code}
  60
  61 %************************************************************************
  62 %*                                                                      *
  63 \subsection{The @UniqFM@ type, and signatures for the functions}
  64 %*                                                                      *
  65 %************************************************************************
  66
  67 We use @FiniteMaps@, with a (@getUnique@-able) @Unique@ as ``key''.
  68
  69 \begin{code}
  70 emptyUFM        :: UniqFM elt
  71 isNullUFM       :: UniqFM elt -> Bool
  72 unitUFM         :: Uniquable key => key -> elt -> UniqFM elt
  73 unitDirectlyUFM -- got the Unique already
  74                 :: Unique -> elt -> UniqFM elt
  75 listToUFM       :: Uniquable key => [(key,elt)] -> UniqFM elt
  76 listToUFM_Directly
  77                 :: [(Unique, elt)] -> UniqFM elt
  78
  79 addToUFM        :: Uniquable key => UniqFM elt -> key -> elt  -> UniqFM elt
  80 addListToUFM    :: Uniquable key => UniqFM elt -> [(key,elt)] -> UniqFM elt
  81 addToUFM_Directly
  82                 :: UniqFM elt -> Unique -> elt -> UniqFM elt
  83
  84 addToUFM_C      :: Uniquable key => (elt -> elt -> elt) -- old -> new -> result
  85                            -> UniqFM elt                -- old
  86                            -> key -> elt                -- new
  87                            -> UniqFM elt                -- result
  88
  89 addListToUFM_C  :: Uniquable key => (elt -> elt -> elt)
  90                            -> UniqFM elt -> [(key,elt)]
  91                            -> UniqFM elt
  92
  93 delFromUFM      :: Uniquable key => UniqFM elt -> key    -> UniqFM elt
  94 delListFromUFM  :: Uniquable key => UniqFM elt -> [key] -> UniqFM elt
  95 delFromUFM_Directly :: UniqFM elt -> Unique -> UniqFM elt
  96
  97 plusUFM         :: UniqFM elt -> UniqFM elt -> UniqFM elt
  98
  99 plusUFM_C       :: (elt -> elt -> elt)
 100                 -> UniqFM elt -> UniqFM elt -> UniqFM elt
 101
 102 minusUFM        :: UniqFM elt -> UniqFM elt -> UniqFM elt
 103
 104 intersectUFM    :: UniqFM elt -> UniqFM elt -> UniqFM elt
 105 intersectUFM_C  :: (elt -> elt -> elt)
 106                 -> UniqFM elt -> UniqFM elt -> UniqFM elt
 107 foldUFM         :: (elt -> a -> a) -> a -> UniqFM elt -> a
 108 mapUFM          :: (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
 109 filterUFM       :: (elt -> Bool) -> UniqFM elt -> UniqFM elt
 110
 111 sizeUFM         :: UniqFM elt -> Int
 112 elemUFM         :: Uniquable key => key -> UniqFM elt -> Bool
 113
 114 lookupUFM       :: Uniquable key => UniqFM elt -> key -> Maybe elt
 115 lookupUFM_Directly  -- when you've got the Unique already
 116                 :: UniqFM elt -> Unique -> Maybe elt
 117 lookupWithDefaultUFM
 118                 :: Uniquable key => UniqFM elt -> elt -> key -> elt
 119 lookupWithDefaultUFM_Directly
 120                 :: UniqFM elt -> elt -> Unique -> elt
 121
 122 keysUFM         :: UniqFM elt -> [Int]          -- Get the keys
 123 eltsUFM         :: UniqFM elt -> [elt]
 124 ufmToList       :: UniqFM elt -> [(Unique, elt)]
 125 \end{code}
 126
 127 %************************************************************************
 128 %*                                                                      *
 129 \subsection{The @IdFinMap@ and @TyVarFinMap@ specialisations for Ids/TyVars}
 130 %*                                                                      *
 131 %************************************************************************
 132
 133 \begin{code}
 134 -- Turn off for now, these need to be updated (SDM 4/98)
 135
 136 #if 0
 137 #ifdef __GLASGOW_HASKELL__
 138 -- I don't think HBC was too happy about this (WDP 94/10)
 139
 140 {-# SPECIALIZE
 141     addListToUFM :: UniqFM elt -> [(Name,   elt)] -> UniqFM elt
 142   #-}
 143 {-# SPECIALIZE
 144     addListToUFM_C :: (elt -> elt -> elt) -> UniqFM elt -> [(Name,  elt)] -> UniqFM elt
 145   #-}
 146 {-# SPECIALIZE
 147     addToUFM    :: UniqFM elt -> Unique -> elt  -> UniqFM elt
 148   #-}
 149 {-# SPECIALIZE
 150     listToUFM   :: [(Unique, elt)]     -> UniqFM elt
 151   #-}
 152 {-# SPECIALIZE
 153     lookupUFM   :: UniqFM elt -> Name   -> Maybe elt
 154                  , UniqFM elt -> Unique -> Maybe elt
 155   #-}
 156
 157 #endif {- __GLASGOW_HASKELL__ -}
 158 #endif
 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 (getUnique 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 (getUnique 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 (getUnique 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 (getUnique 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 elemUFM key fm = case lookUp fm (u2i (getUnique key)) of
 539                         Nothing -> False
 540                         Just _  -> True
 541
 542 lookupUFM          fm key = lookUp fm (u2i (getUnique key))
 543 lookupUFM_Directly fm key = lookUp fm (u2i key)
 544
 545 lookupWithDefaultUFM fm deflt key
 546   = case lookUp fm (u2i (getUnique key)) of
 547       Nothing  -> deflt
 548       Just elt -> elt
 549
 550 lookupWithDefaultUFM_Directly fm deflt key
 551   = case lookUp fm (u2i key) of
 552       Nothing  -> deflt
 553       Just elt -> elt
 554
 555 lookUp EmptyUFM _   = Nothing
 556 lookUp fm i         = lookup_tree fm
 557   where
 558         lookup_tree :: UniqFM a -> Maybe a
 559
 560         lookup_tree (LeafUFM j b)
 561           | j _EQ_ i    = Just b
 562           | otherwise   = Nothing
 563         lookup_tree (NodeUFM j p t1 t2)
 564           | j _GT_ i    = lookup_tree t1
 565           | otherwise   = lookup_tree t2
 566
 567         lookup_tree EmptyUFM = panic "lookup Failed"
 568 \end{code}
 569
 570 folds are *wonderful* things.
 571
 572 \begin{code}
 573 eltsUFM fm = foldUFM (:) [] fm
 574
 575 ufmToList fm = fold_tree (\ iu elt rest -> (mkUniqueGrimily iu, elt) : rest) [] fm
 576
 577 keysUFM fm = fold_tree (\ iu elt rest -> IBOX(iu) : rest) [] fm
 578
 579 fold_tree f a (NodeUFM _ _ t1 t2) = fold_tree f (fold_tree f a t2) t1
 580 fold_tree f a (LeafUFM iu obj)    = f iu obj a
 581 fold_tree f a EmptyUFM            = a
 582 \end{code}
 583
 584 %************************************************************************
 585 %*                                                                      *
 586 \subsubsection{The @UniqFM@ type, and its functions}
 587 %*                                                                      *
 588 %************************************************************************
 589
 590 You should always use these to build the tree.
 591 There are 4 versions of mkNodeUFM, depending on
 592 the strictness of the two sub-tree arguments.
 593 The strictness is used *both* to prune out
 594 empty trees, *and* to improve performance,
 595 stoping needless thunks lying around.
 596 The rule of thumb (from experence with these trees)
 597 is make thunks strict, but data structures lazy.
 598 If in doubt, use mkSSNodeUFM, which has the `strongest'
 599 functionality, but may do a few needless evaluations.
 600
 601 \begin{code}
 602 mkLeafUFM :: FAST_INT -> a -> UniqFM a
 603 mkLeafUFM i a     = LeafUFM i a
 604
 605 -- The *ONLY* ways of building a NodeUFM.
 606
 607 mkSSNodeUFM (NodeUFMData j p) EmptyUFM t2 = t2
 608 mkSSNodeUFM (NodeUFMData j p) t1 EmptyUFM = t1
 609 mkSSNodeUFM (NodeUFMData j p) t1 t2
 610   = ASSERT(correctNodeUFM (IBOX(j)) (IBOX(p)) t1 t2)
 611     NodeUFM j p t1 t2
 612
 613 mkSLNodeUFM (NodeUFMData j p) EmptyUFM t2 = t2
 614 mkSLNodeUFM (NodeUFMData j p) t1 t2
 615   = ASSERT(correctNodeUFM (IBOX(j)) (IBOX(p)) t1 t2)
 616     NodeUFM j p t1 t2
 617
 618 mkLSNodeUFM (NodeUFMData j p) t1 EmptyUFM = t1
 619 mkLSNodeUFM (NodeUFMData j p) t1 t2
 620   = ASSERT(correctNodeUFM (IBOX(j)) (IBOX(p)) t1 t2)
 621     NodeUFM j p t1 t2
 622
 623 mkLLNodeUFM (NodeUFMData j p) t1 t2
 624   = ASSERT(correctNodeUFM (IBOX(j)) (IBOX(p)) t1 t2)
 625     NodeUFM j p t1 t2
 626
 627 correctNodeUFM
 628         :: Int
 629         -> Int
 630         -> UniqFM a
 631         -> UniqFM a
 632         -> Bool
 633
 634 correctNodeUFM j p t1 t2
 635   = correct (j-p) (j-1) p t1 && correct j ((j-1)+p) p t2
 636   where
 637     correct low high _ (LeafUFM i _)
 638       = low <= IBOX(i) && IBOX(i) <= high
 639     correct low high above_p (NodeUFM j p _ _)
 640       = low <= IBOX(j) && IBOX(j) <= high && above_p > IBOX(p)
 641     correct _ _ _ EmptyUFM = panic "EmptyUFM stored inside a tree"
 642 \end{code}
 643
 644 Note: doing SAT on this by hand seems to make it worse. Todo: Investigate,
 645 and if necessary do $\lambda$ lifting on our functions that are bound.
 646
 647 \begin{code}
 648 insert_ele
 649         :: (a -> a -> a)
 650         -> UniqFM a
 651         -> FAST_INT
 652         -> a
 653         -> UniqFM a
 654
 655 insert_ele f EmptyUFM i new = mkLeafUFM i new
 656
 657 insert_ele f (LeafUFM j old) i new
 658   | j _GT_ i =
 659           mkLLNodeUFM (getCommonNodeUFMData
 660                           (indexToRoot i)
 661                           (indexToRoot j))
 662                  (mkLeafUFM i new)
 663                  (mkLeafUFM j old)
 664   | j _EQ_ i  = mkLeafUFM j (f old new)
 665   | otherwise =
 666           mkLLNodeUFM (getCommonNodeUFMData
 667                           (indexToRoot i)
 668                           (indexToRoot j))
 669                  (mkLeafUFM j old)
 670                  (mkLeafUFM i new)
 671
 672 insert_ele f n@(NodeUFM j p t1 t2) i a
 673   | i _LT_ j
 674     = if (i _GE_ (j _SUB_ p))
 675       then mkSLNodeUFM (NodeUFMData j p) (insert_ele f t1 i a) t2
 676       else mkLLNodeUFM (getCommonNodeUFMData
 677                           (indexToRoot i)
 678                           ((NodeUFMData j p)))
 679                   (mkLeafUFM i a)
 680                   n
 681   | otherwise
 682     = if (i _LE_ ((j _SUB_ ILIT(1)) _ADD_ p))
 683       then mkLSNodeUFM (NodeUFMData j p) t1 (insert_ele f t2 i a)
 684       else mkLLNodeUFM (getCommonNodeUFMData
 685                           (indexToRoot i)
 686                           ((NodeUFMData j p)))
 687                   n
 688                   (mkLeafUFM i a)
 689 \end{code}
 690
 691
 692
 693 \begin{code}
 694 map_tree f (NodeUFM j p t1 t2)
 695   = mkSSNodeUFM (NodeUFMData j p) (map_tree f t1) (map_tree f t2)
 696 map_tree f (LeafUFM i obj)
 697   = mkLeafUFM i (f obj)
 698
 699 map_tree f _ = panic "map_tree failed"
 700 \end{code}
 701
 702 \begin{code}
 703 filter_tree f nd@(NodeUFM j p t1 t2)
 704   = mkSSNodeUFM (NodeUFMData j p) (filter_tree f t1) (filter_tree f t2)
 705
 706 filter_tree f lf@(LeafUFM i obj)
 707   | f obj = lf
 708   | otherwise = EmptyUFM
 709 filter_tree f _ = panic "filter_tree failed"
 710 \end{code}
 711
 712 %************************************************************************
 713 %*                                                                      *
 714 \subsubsection{The @UniqFM@ type, and signatures for the functions}
 715 %*                                                                      *
 716 %************************************************************************
 717
 718 Now some Utilities;
 719
 720 This is the information that is held inside a NodeUFM, packaged up for
 721 consumer use.
 722
 723 \begin{code}
 724 data NodeUFMData
 725   = NodeUFMData FAST_INT
 726                 FAST_INT
 727 \end{code}
 728
 729 This is the information used when computing new NodeUFMs.
 730
 731 \begin{code}
 732 data Side = Leftt | Rightt -- NB: avoid 1.3 names "Left" and "Right"
 733 data CommonRoot
 734   = LeftRoot  Side      -- which side is the right down ?
 735   | RightRoot Side      -- which side is the left down ?
 736   | SameRoot            -- they are the same !
 737   | NewRoot NodeUFMData -- here's the new, common, root
 738             Bool        -- do you need to swap left and right ?
 739 \end{code}
 740
 741 This specifies the relationship between NodeUFMData and CalcNodeUFMData.
 742
 743 \begin{code}
 744 indexToRoot :: FAST_INT -> NodeUFMData
 745
 746 indexToRoot i
 747   = let
 748         l = (ILIT(1) :: FAST_INT)
 749     in
 750     NodeUFMData (((i `shiftR_` l) `shiftL_` l) _ADD_ ILIT(1)) l
 751
 752 getCommonNodeUFMData :: NodeUFMData -> NodeUFMData -> NodeUFMData
 753
 754 getCommonNodeUFMData (NodeUFMData i p) (NodeUFMData i2 p2)
 755   | p _EQ_ p2   = getCommonNodeUFMData_ p j j2
 756   | p _LT_ p2   = getCommonNodeUFMData_ p2 (j _QUOT_ (p2 _QUOT_ p)) j2
 757   | otherwise   = getCommonNodeUFMData_ p j (j2 _QUOT_ (p _QUOT_ p2))
 758   where
 759     l  = (ILIT(1) :: FAST_INT)
 760     j  = i  _QUOT_ (p  `shiftL_` l)
 761     j2 = i2 _QUOT_ (p2 `shiftL_` l)
 762
 763     getCommonNodeUFMData_ :: FAST_INT -> FAST_INT -> FAST_INT -> NodeUFMData
 764
 765     getCommonNodeUFMData_ p j j_
 766       | j _EQ_ j_
 767       = NodeUFMData (((j `shiftL_` l) _ADD_ l) _MUL_ p) p
 768       | otherwise
 769       = getCommonNodeUFMData_ (p `shiftL_`  l) (j `shiftR_` l) (j_ `shiftR_` l)
 770
 771 ask_about_common_ancestor :: NodeUFMData -> NodeUFMData -> CommonRoot
 772
 773 ask_about_common_ancestor x@(NodeUFMData j p) y@(NodeUFMData j2 p2)
 774   | j _EQ_ j2 = SameRoot
 775   | otherwise
 776   = case getCommonNodeUFMData x y of
 777       nd@(NodeUFMData j3 p3)
 778         | j3 _EQ_ j  -> LeftRoot (decideSide (j _GT_ j2))
 779         | j3 _EQ_ j2 -> RightRoot (decideSide (j _LT_ j2))
 780         | otherwise   -> NewRoot nd (j _GT_ j2)
 781     where
 782         decideSide :: Bool -> Side
 783         decideSide True  = Leftt
 784         decideSide False = Rightt
 785 \end{code}
 786
 787 This might be better in Util.lhs ?
 788
 789
 790 Now the bit twiddling functions.
 791 \begin{code}
 792 shiftL_ :: FAST_INT -> FAST_INT -> FAST_INT
 793 shiftR_ :: FAST_INT -> FAST_INT -> FAST_INT
 794
 795 #if __GLASGOW_HASKELL__
 796 {-# INLINE shiftL_ #-}
 797 {-# INLINE shiftR_ #-}
 798 shiftL_ n p = word2Int#((int2Word# n) `shiftL#` p)
 799 shiftR_ n p = word2Int#((int2Word# n) `shiftr` p)
 800   where
 801     shiftr x y = shiftRL# x y
 802
 803 #else {- not GHC -}
 804 shiftL_ n p = n * (2 ^ p)
 805 shiftR_ n p = n `quot` (2 ^ p)
 806
 807 #endif {- not GHC -}
 808 \end{code}
 809
 810 \begin{code}
 811 use_snd :: a -> b -> b
 812 use_snd a b = b
 813 \end{code}