ghc/compiler/utils/UniqFM.lhs

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