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