compiler/utils/UniqFM.lhs

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