ghc/compiler/utils/UniqFM.lhs

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