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