ghc/compiler/utils/UniqFM.lhs

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