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