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