compiler/utils/UniqFM.lhs

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