ghc/compiler/simplCore/OccurAnal.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 %************************************************************************
   5 %*                                                                      *
   6 \section[OccurAnal]{Occurrence analysis pass}
   7 %*                                                                      *
   8 %************************************************************************
   9
  10 The occurrence analyser re-typechecks a core expression, returning a new
  11 core expression with (hopefully) improved usage information.
  12
  13 \begin{code}
  14 module OccurAnal (
  15         occurAnalysePgm, occurAnalyseGlobalExpr, occurAnalyseRule,
  16     ) where
  17
  18 #include "HsVersions.h"
  19
  20 import CoreSyn
  21 import CoreFVs          ( idRuleVars )
  22 import CoreUtils        ( exprIsTrivial )
  23 import Id               ( isDataConWorkId, isOneShotBndr, setOneShotLambda,
  24                           idOccInfo, setIdOccInfo,
  25                           isExportedId, idArity, idSpecialisation,
  26                           idType, idUnique, Id
  27                         )
  28 import BasicTypes       ( OccInfo(..), isOneOcc )
  29
  30 import VarSet
  31 import VarEnv
  32
  33 import Type             ( isFunTy, dropForAlls )
  34 import Maybes           ( orElse )
  35 import Digraph          ( stronglyConnCompR, SCC(..) )
  36 import PrelNames        ( buildIdKey, foldrIdKey, runSTRepIdKey, augmentIdKey )
  37 import Unique           ( Unique )
  38 import UniqFM           ( keysUFM )
  39 import Util             ( zipWithEqual, mapAndUnzip )
  40 import Outputable
  41 \end{code}
  42
  43
  44 %************************************************************************
  45 %*                                                                      *
  46 \subsection[OccurAnal-main]{Counting occurrences: main function}
  47 %*                                                                      *
  48 %************************************************************************
  49
  50 Here's the externally-callable interface:
  51
  52 \begin{code}
  53 occurAnalysePgm :: [CoreBind] -> [CoreBind]
  54 occurAnalysePgm binds
  55   = snd (go (initOccEnv emptyVarSet) binds)
  56   where
  57     go :: OccEnv -> [CoreBind] -> (UsageDetails, [CoreBind])
  58     go env []
  59         = (emptyDetails, [])
  60     go env (bind:binds)
  61         = (final_usage, bind' ++ binds')
  62         where
  63            new_env              = env `addNewCands` (bindersOf bind)
  64            (bs_usage, binds')   = go new_env binds
  65            (final_usage, bind') = occAnalBind env bind bs_usage
  66
  67 occurAnalyseGlobalExpr :: CoreExpr -> CoreExpr
  68 occurAnalyseGlobalExpr expr
  69   =     -- Top level expr, so no interesting free vars, and
  70         -- discard occurence info returned
  71     snd (occAnal (initOccEnv emptyVarSet) expr)
  72
  73 occurAnalyseRule :: CoreRule -> CoreRule
  74 occurAnalyseRule rule@(BuiltinRule _ _) = rule
  75 occurAnalyseRule (Rule str act tpl_vars tpl_args rhs)
  76                 -- Add occ info to tpl_vars, rhs
  77   = Rule str act tpl_vars' tpl_args rhs'
  78   where
  79     (rhs_uds, rhs') = occAnal (initOccEnv (mkVarSet tpl_vars)) rhs
  80     (_, tpl_vars')  = tagBinders rhs_uds tpl_vars
  81 \end{code}
  82
  83
  84 %************************************************************************
  85 %*                                                                      *
  86 \subsection[OccurAnal-main]{Counting occurrences: main function}
  87 %*                                                                      *
  88 %************************************************************************
  89
  90 Bindings
  91 ~~~~~~~~
  92
  93 \begin{code}
  94 type IdWithOccInfo = Id                 -- An Id with fresh PragmaInfo attached
  95
  96 type Node details = (details, Unique, [Unique]) -- The Ints are gotten from the Unique,
  97                                                 -- which is gotten from the Id.
  98 type Details1     = (Id, UsageDetails, CoreExpr)
  99 type Details2     = (IdWithOccInfo, CoreExpr)
 100
 101
 102 occAnalBind :: OccEnv
 103             -> CoreBind
 104             -> UsageDetails             -- Usage details of scope
 105             -> (UsageDetails,           -- Of the whole let(rec)
 106                 [CoreBind])
 107
 108 occAnalBind env (NonRec binder rhs) body_usage
 109   | not (binder `usedIn` body_usage)            -- It's not mentioned
 110   = (body_usage, [])
 111
 112   | otherwise                   -- It's mentioned in the body
 113   = (final_body_usage `combineUsageDetails` rhs_usage,
 114      [NonRec tagged_binder rhs'])
 115
 116   where
 117     (final_body_usage, tagged_binder) = tagBinder body_usage binder
 118     (rhs_usage, rhs')                 = occAnalRhs env tagged_binder rhs
 119 \end{code}
 120
 121 Dropping dead code for recursive bindings is done in a very simple way:
 122
 123         the entire set of bindings is dropped if none of its binders are
 124         mentioned in its body; otherwise none are.
 125
 126 This seems to miss an obvious improvement.
 127 @
 128         letrec  f = ...g...
 129                 g = ...f...
 130         in
 131         ...g...
 132
 133 ===>
 134
 135         letrec f = ...g...
 136                g = ...(...g...)...
 137         in
 138         ...g...
 139 @
 140
 141 Now @f@ is unused. But dependency analysis will sort this out into a
 142 @letrec@ for @g@ and a @let@ for @f@, and then @f@ will get dropped.
 143 It isn't easy to do a perfect job in one blow.  Consider
 144
 145 @
 146         letrec f = ...g...
 147                g = ...h...
 148                h = ...k...
 149                k = ...m...
 150                m = ...m...
 151         in
 152         ...m...
 153 @
 154
 155
 156 \begin{code}
 157 occAnalBind env (Rec pairs) body_usage
 158   = foldr (_scc_ "occAnalBind.dofinal" do_final_bind) (body_usage, []) sccs
 159   where
 160     binders = map fst pairs
 161     rhs_env = env `addNewCands` binders
 162
 163     analysed_pairs :: [Details1]
 164     analysed_pairs  = [ (bndr, rhs_usage, rhs')
 165                       | (bndr, rhs) <- pairs,
 166                         let (rhs_usage, rhs') = occAnalRhs rhs_env bndr rhs
 167                       ]
 168
 169     sccs :: [SCC (Node Details1)]
 170     sccs = _scc_ "occAnalBind.scc" stronglyConnCompR edges
 171
 172
 173     ---- stuff for dependency analysis of binds -------------------------------
 174     edges :: [Node Details1]
 175     edges = _scc_ "occAnalBind.assoc"
 176             [ (details, idUnique id, edges_from rhs_usage)
 177             | details@(id, rhs_usage, rhs) <- analysed_pairs
 178             ]
 179
 180         -- (a -> b) means a mentions b
 181         -- Given the usage details (a UFM that gives occ info for each free var of
 182         -- the RHS) we can get the list of free vars -- or rather their Int keys --
 183         -- by just extracting the keys from the finite map.  Grimy, but fast.
 184         -- Previously we had this:
 185         --      [ bndr | bndr <- bndrs,
 186         --               maybeToBool (lookupVarEnv rhs_usage bndr)]
 187         -- which has n**2 cost, and this meant that edges_from alone
 188         -- consumed 10% of total runtime!
 189     edges_from :: UsageDetails -> [Unique]
 190     edges_from rhs_usage = _scc_ "occAnalBind.edges_from"
 191                            keysUFM rhs_usage
 192
 193     ---- stuff to "re-constitute" bindings from dependency-analysis info ------
 194
 195         -- Non-recursive SCC
 196     do_final_bind (AcyclicSCC ((bndr, rhs_usage, rhs'), _, _)) (body_usage, binds_so_far)
 197       | not (bndr `usedIn` body_usage)
 198       = (body_usage, binds_so_far)                      -- Dead code
 199       | otherwise
 200       = (combined_usage, new_bind : binds_so_far)
 201       where
 202         total_usage                   = combineUsageDetails body_usage rhs_usage
 203         (combined_usage, tagged_bndr) = tagBinder total_usage bndr
 204         new_bind                      = NonRec tagged_bndr rhs'
 205
 206         -- Recursive SCC
 207     do_final_bind (CyclicSCC cycle) (body_usage, binds_so_far)
 208       | not (any (`usedIn` body_usage) bndrs)           -- NB: look at body_usage, not total_usage
 209       = (body_usage, binds_so_far)                      -- Dead code
 210       | otherwise
 211       = (combined_usage, final_bind:binds_so_far)
 212       where
 213         details                        = [details   | (details, _, _) <- cycle]
 214         bndrs                          = [bndr      | (bndr, _, _)      <- details]
 215         rhs_usages                     = [rhs_usage | (_, rhs_usage, _) <- details]
 216         total_usage                    = foldr combineUsageDetails body_usage rhs_usages
 217         (combined_usage, tagged_bndrs) = tagBinders total_usage bndrs
 218         final_bind                     = Rec (reOrderRec env new_cycle)
 219
 220         new_cycle = CyclicSCC (zipWithEqual "occAnalBind" mk_new_bind tagged_bndrs cycle)
 221         mk_new_bind tagged_bndr ((_, _, rhs'), key, keys) = ((tagged_bndr, rhs'), key, keys)
 222 \end{code}
 223
 224 @reOrderRec@ is applied to the list of (binder,rhs) pairs for a cyclic
 225 strongly connected component (there's guaranteed to be a cycle).  It returns the
 226 same pairs, but
 227         a) in a better order,
 228         b) with some of the Ids having a IMustNotBeINLINEd pragma
 229
 230 The "no-inline" Ids are sufficient to break all cycles in the SCC.  This means
 231 that the simplifier can guarantee not to loop provided it never records an inlining
 232 for these no-inline guys.
 233
 234 Furthermore, the order of the binds is such that if we neglect dependencies
 235 on the no-inline Ids then the binds are topologically sorted.  This means
 236 that the simplifier will generally do a good job if it works from top bottom,
 237 recording inlinings for any Ids which aren't marked as "no-inline" as it goes.
 238
 239 ==============
 240 [June 98: I don't understand the following paragraphs, and I've
 241           changed the a=b case again so that it isn't a special case any more.]
 242
 243 Here's a case that bit me:
 244
 245         letrec
 246                 a = b
 247                 b = \x. BIG
 248         in
 249         ...a...a...a....
 250
 251 Re-ordering doesn't change the order of bindings, but there was no loop-breaker.
 252
 253 My solution was to make a=b bindings record b as Many, rather like INLINE bindings.
 254 Perhaps something cleverer would suffice.
 255 ===============
 256
 257 You might think that you can prevent non-termination simply by making
 258 sure that we simplify a recursive binding's RHS in an environment that
 259 simply clones the recursive Id.  But no.  Consider
 260
 261                 letrec f = \x -> let z = f x' in ...
 262
 263                 in
 264                 let n = f y
 265                 in
 266                 case n of { ... }
 267
 268 We bind n to its *simplified* RHS, we then *re-simplify* it when
 269 we inline n.  Then we may well inline f; and then the same thing
 270 happens with z!
 271
 272 I don't think it's possible to prevent non-termination by environment
 273 manipulation in this way.  Apart from anything else, successive
 274 iterations of the simplifier may unroll recursive loops in cases like
 275 that above.  The idea of beaking every recursive loop with an
 276 IMustNotBeINLINEd pragma is much much better.
 277
 278
 279 \begin{code}
 280 reOrderRec
 281         :: OccEnv
 282         -> SCC (Node Details2)
 283         -> [Details2]
 284                         -- Sorted into a plausible order.  Enough of the Ids have
 285                         --      dontINLINE pragmas that there are no loops left.
 286
 287         -- Non-recursive case
 288 reOrderRec env (AcyclicSCC (bind, _, _)) = [bind]
 289
 290         -- Common case of simple self-recursion
 291 reOrderRec env (CyclicSCC [bind])
 292   = [(setIdOccInfo tagged_bndr IAmALoopBreaker, rhs)]
 293   where
 294     ((tagged_bndr, rhs), _, _) = bind
 295
 296 reOrderRec env (CyclicSCC (bind : binds))
 297   =     -- Choose a loop breaker, mark it no-inline,
 298         -- do SCC analysis on the rest, and recursively sort them out
 299     concat (map (reOrderRec env) (stronglyConnCompR unchosen))
 300     ++
 301     [(setIdOccInfo tagged_bndr IAmALoopBreaker, rhs)]
 302
 303   where
 304     (chosen_pair, unchosen) = choose_loop_breaker bind (score bind) [] binds
 305     (tagged_bndr, rhs)      = chosen_pair
 306
 307         -- This loop looks for the bind with the lowest score
 308         -- to pick as the loop  breaker.  The rest accumulate in
 309     choose_loop_breaker (details,_,_) loop_sc acc []
 310         = (details, acc)        -- Done
 311
 312     choose_loop_breaker loop_bind loop_sc acc (bind : binds)
 313         | sc < loop_sc  -- Lower score so pick this new one
 314         = choose_loop_breaker bind sc (loop_bind : acc) binds
 315
 316         | otherwise     -- No lower so don't pick it
 317         = choose_loop_breaker loop_bind loop_sc (bind : acc) binds
 318         where
 319           sc = score bind
 320
 321     score :: Node Details2 -> Int       -- Higher score => less likely to be picked as loop breaker
 322     score ((bndr, rhs), _, _)
 323         | exprIsTrivial rhs        = 4  -- Practically certain to be inlined
 324                 -- Used to have also: && not (isExportedId bndr)
 325                 -- But I found this sometimes cost an extra iteration when we have
 326                 --      rec { d = (a,b); a = ...df...; b = ...df...; df = d }
 327                 -- where df is the exported dictionary. Then df makes a really
 328                 -- bad choice for loop breaker
 329
 330         | not_fun_ty (idType bndr) = 3  -- Data types help with cases
 331                 -- This used to have a lower score than inlineCandidate, but
 332                 -- it's *really* helpful if dictionaries get inlined fast,
 333                 -- so I'm experimenting with giving higher priority to data-typed things
 334
 335         | inlineCandidate bndr rhs = 2  -- Likely to be inlined
 336
 337         | not (isEmptyCoreRules (idSpecialisation bndr)) = 1
 338                 -- Avoid things with specialisations; we'd like
 339                 -- to take advantage of them in the subsequent bindings
 340
 341         | otherwise = 0
 342
 343     inlineCandidate :: Id -> CoreExpr -> Bool
 344     inlineCandidate id (Note InlineMe _) = True
 345     inlineCandidate id rhs               = isOneOcc (idOccInfo id)
 346
 347         -- Real example (the Enum Ordering instance from PrelBase):
 348         --      rec     f = \ x -> case d of (p,q,r) -> p x
 349         --              g = \ x -> case d of (p,q,r) -> q x
 350         --              d = (v, f, g)
 351         --
 352         -- Here, f and g occur just once; but we can't inline them into d.
 353         -- On the other hand we *could* simplify those case expressions if
 354         -- we didn't stupidly choose d as the loop breaker.
 355         -- But we won't because constructor args are marked "Many".
 356
 357     not_fun_ty ty = not (isFunTy (dropForAlls ty))
 358 \end{code}
 359
 360 @occAnalRhs@ deals with the question of bindings where the Id is marked
 361 by an INLINE pragma.  For these we record that anything which occurs
 362 in its RHS occurs many times.  This pessimistically assumes that ths
 363 inlined binder also occurs many times in its scope, but if it doesn't
 364 we'll catch it next time round.  At worst this costs an extra simplifier pass.
 365 ToDo: try using the occurrence info for the inline'd binder.
 366
 367 [March 97] We do the same for atomic RHSs.  Reason: see notes with reOrderRec.
 368 [June 98, SLPJ]  I've undone this change; I don't understand it.  See notes with reOrderRec.
 369
 370
 371 \begin{code}
 372 occAnalRhs :: OccEnv
 373            -> Id -> CoreExpr    -- Binder and rhs
 374                                 -- For non-recs the binder is alrady tagged
 375                                 -- with occurrence info
 376            -> (UsageDetails, CoreExpr)
 377
 378 occAnalRhs env id rhs
 379   = (final_usage, rhs')
 380   where
 381     (rhs_usage, rhs') = occAnal ctxt rhs
 382     ctxt | certainly_inline id = env
 383          | otherwise           = rhsCtxt env
 384         -- Note that we generally use an rhsCtxt.  This tells the occ anal n
 385         -- that it's looking at an RHS, which has an effect in occAnalApp
 386         --
 387         -- But there's a problem.  Consider
 388         --      x1 = a0 : []
 389         --      x2 = a1 : x1
 390         --      x3 = a2 : x2
 391         --      g  = f x3
 392         -- First time round, it looks as if x1 and x2 occur as an arg of a
 393         -- let-bound constructor ==> give them a many-occurrence.
 394         -- But then x3 is inlined (unconditionally as it happens) and
 395         -- next time round, x2 will be, and the next time round x1 will be
 396         -- Result: multiple simplifier iterations.  Sigh.
 397         -- Crude solution: use rhsCtxt for things that occur just once...
 398
 399     certainly_inline id = case idOccInfo id of
 400                             OneOcc in_lam one_br -> not in_lam && one_br
 401                             other                -> False
 402
 403         -- [March 98] A new wrinkle is that if the binder has specialisations inside
 404         -- it then we count the specialised Ids as "extra rhs's".  That way
 405         -- the "parent" keeps the specialised "children" alive.  If the parent
 406         -- dies (because it isn't referenced any more), then the children will
 407         -- die too unless they are already referenced directly.
 408
 409     final_usage = addRuleUsage rhs_usage id
 410
 411 addRuleUsage :: UsageDetails -> Id -> UsageDetails
 412 -- Add the usage from RULES in Id to the usage
 413 addRuleUsage usage id
 414   = foldVarSet add usage (idRuleVars id)
 415   where
 416     add v u = addOneOcc u v NoOccInfo           -- Give a non-committal binder info
 417                                                 -- (i.e manyOcc) because many copies
 418                                                 -- of the specialised thing can appear
 419 \end{code}
 420
 421 Expressions
 422 ~~~~~~~~~~~
 423 \begin{code}
 424 occAnal :: OccEnv
 425         -> CoreExpr
 426         -> (UsageDetails,       -- Gives info only about the "interesting" Ids
 427             CoreExpr)
 428
 429 occAnal env (Type t)  = (emptyDetails, Type t)
 430
 431 occAnal env (Var v)
 432   = (var_uds, Var v)
 433   where
 434     var_uds | isCandidate env v = unitVarEnv v oneOcc
 435             | otherwise         = emptyDetails
 436
 437     -- At one stage, I gathered the idRuleVars for v here too,
 438     -- which in a way is the right thing to do.
 439     -- But that went wrong right after specialisation, when
 440     -- the *occurrences* of the overloaded function didn't have any
 441     -- rules in them, so the *specialised* versions looked as if they
 442     -- weren't used at all.
 443
 444 \end{code}
 445
 446 We regard variables that occur as constructor arguments as "dangerousToDup":
 447
 448 \begin{verbatim}
 449 module A where
 450 f x = let y = expensive x in
 451       let z = (True,y) in
 452       (case z of {(p,q)->q}, case z of {(p,q)->q})
 453 \end{verbatim}
 454
 455 We feel free to duplicate the WHNF (True,y), but that means
 456 that y may be duplicated thereby.
 457
 458 If we aren't careful we duplicate the (expensive x) call!
 459 Constructors are rather like lambdas in this way.
 460
 461 \begin{code}
 462 occAnal env expr@(Lit lit) = (emptyDetails, expr)
 463 \end{code}
 464
 465 \begin{code}
 466 occAnal env (Note InlineMe body)
 467   = case occAnal env body of { (usage, body') ->
 468     (mapVarEnv markMany usage, Note InlineMe body')
 469     }
 470
 471 occAnal env (Note note@(SCC cc) body)
 472   = case occAnal env body of { (usage, body') ->
 473     (mapVarEnv markInsideSCC usage, Note note body')
 474     }
 475
 476 occAnal env (Note note body)
 477   = case occAnal env body of { (usage, body') ->
 478     (usage, Note note body')
 479     }
 480 \end{code}
 481
 482 \begin{code}
 483 occAnal env app@(App fun arg)
 484   = occAnalApp env (collectArgs app) False
 485
 486 -- Ignore type variables altogether
 487 --   (a) occurrences inside type lambdas only not marked as InsideLam
 488 --   (b) type variables not in environment
 489
 490 occAnal env expr@(Lam x body) | isTyVar x
 491   = case occAnal env body of { (body_usage, body') ->
 492     (body_usage, Lam x body')
 493     }
 494
 495 -- For value lambdas we do a special hack.  Consider
 496 --      (\x. \y. ...x...)
 497 -- If we did nothing, x is used inside the \y, so would be marked
 498 -- as dangerous to dup.  But in the common case where the abstraction
 499 -- is applied to two arguments this is over-pessimistic.
 500 -- So instead, we just mark each binder with its occurrence
 501 -- info in the *body* of the multiple lambda.
 502 -- Then, the simplifier is careful when partially applying lambdas.
 503
 504 occAnal env expr@(Lam _ _)
 505   = case occAnal env_body body of { (body_usage, body') ->
 506     let
 507         (final_usage, tagged_binders) = tagBinders body_usage binders
 508         --      URGH!  Sept 99: we don't seem to be able to use binders' here, because
 509         --      we get linear-typed things in the resulting program that we can't handle yet.
 510         --      (e.g. PrelShow)  TODO
 511
 512         really_final_usage = if linear then
 513                                 final_usage
 514                              else
 515                                 mapVarEnv markInsideLam final_usage
 516     in
 517     (really_final_usage,
 518      mkLams tagged_binders body') }
 519   where
 520     (binders, body)   = collectBinders expr
 521     (linear, env1, _) = oneShotGroup env binders
 522     env2              = env1 `addNewCands` binders      -- Add in-scope binders
 523     env_body          = vanillaCtxt env2                -- Body is (no longer) an RhsContext
 524
 525 occAnal env (Case scrut bndr ty alts)
 526   = case mapAndUnzip (occAnalAlt alt_env bndr) alts of { (alts_usage_s, alts')   ->
 527     case occAnal (vanillaCtxt env) scrut            of { (scrut_usage, scrut') ->
 528         -- No need for rhsCtxt
 529     let
 530         alts_usage  = foldr1 combineAltsUsageDetails alts_usage_s
 531         alts_usage' = addCaseBndrUsage alts_usage
 532         (alts_usage1, tagged_bndr) = tagBinder alts_usage' bndr
 533         total_usage = scrut_usage `combineUsageDetails` alts_usage1
 534     in
 535     total_usage `seq` (total_usage, Case scrut' tagged_bndr ty alts') }}
 536   where
 537     alt_env = env `addNewCand` bndr
 538
 539         -- The case binder gets a usage of either "many" or "dead", never "one".
 540         -- Reason: we like to inline single occurrences, to eliminate a binding,
 541         -- but inlining a case binder *doesn't* eliminate a binding.
 542         -- We *don't* want to transform
 543         --      case x of w { (p,q) -> f w }
 544         -- into
 545         --      case x of w { (p,q) -> f (p,q) }
 546     addCaseBndrUsage usage = case lookupVarEnv usage bndr of
 547                                 Nothing  -> usage
 548                                 Just occ -> extendVarEnv usage bndr (markMany occ)
 549
 550 occAnal env (Let bind body)
 551   = case occAnal new_env body            of { (body_usage, body') ->
 552     case occAnalBind env bind body_usage of { (final_usage, new_binds) ->
 553        (final_usage, mkLets new_binds body') }}
 554   where
 555     new_env = env `addNewCands` (bindersOf bind)
 556
 557 occAnalArgs env args
 558   = case mapAndUnzip (occAnal arg_env) args of  { (arg_uds_s, args') ->
 559     (foldr combineUsageDetails emptyDetails arg_uds_s, args')}
 560   where
 561     arg_env = vanillaCtxt env
 562 \end{code}
 563
 564 Applications are dealt with specially because we want
 565 the "build hack" to work.
 566
 567 \begin{code}
 568 -- Hack for build, fold, runST
 569 occAnalApp env (Var fun, args) is_rhs
 570   = case args_stuff of { (args_uds, args') ->
 571     let
 572         -- We mark the free vars of the argument of a constructor or PAP
 573         -- as "many", if it is the RHS of a let(rec).
 574         -- This means that nothing gets inlined into a constructor argument
 575         -- position, which is what we want.  Typically those constructor
 576         -- arguments are just variables, or trivial expressions.
 577         --
 578         -- This is the *whole point* of the isRhsEnv predicate
 579         final_args_uds
 580                 | isRhsEnv env,
 581                   isDataConWorkId fun || valArgCount args < idArity fun
 582                 = mapVarEnv markMany args_uds
 583                 | otherwise = args_uds
 584     in
 585     (fun_uds `combineUsageDetails` final_args_uds, mkApps (Var fun) args') }
 586   where
 587     fun_uniq = idUnique fun
 588
 589     fun_uds | isCandidate env fun = unitVarEnv fun oneOcc
 590             | otherwise           = emptyDetails
 591
 592     args_stuff  | fun_uniq == buildIdKey    = appSpecial env 2 [True,True]  args
 593                 | fun_uniq == augmentIdKey  = appSpecial env 2 [True,True]  args
 594                 | fun_uniq == foldrIdKey    = appSpecial env 3 [False,True] args
 595                 | fun_uniq == runSTRepIdKey = appSpecial env 2 [True]       args
 596                         -- (foldr k z xs) may call k many times, but it never
 597                         -- shares a partial application of k; hence [False,True]
 598                         -- This means we can optimise
 599                         --      foldr (\x -> let v = ...x... in \y -> ...v...) z xs
 600                         -- by floating in the v
 601
 602                 | otherwise = occAnalArgs env args
 603
 604
 605 occAnalApp env (fun, args) is_rhs
 606   = case occAnal (addAppCtxt env args) fun of   { (fun_uds, fun') ->
 607         -- The addAppCtxt is a bit cunning.  One iteration of the simplifier
 608         -- often leaves behind beta redexs like
 609         --      (\x y -> e) a1 a2
 610         -- Here we would like to mark x,y as one-shot, and treat the whole
 611         -- thing much like a let.  We do this by pushing some True items
 612         -- onto the context stack.
 613
 614     case occAnalArgs env args of        { (args_uds, args') ->
 615     let
 616         final_uds = fun_uds `combineUsageDetails` args_uds
 617     in
 618     (final_uds, mkApps fun' args') }}
 619
 620 appSpecial :: OccEnv
 621            -> Int -> CtxtTy     -- Argument number, and context to use for it
 622            -> [CoreExpr]
 623            -> (UsageDetails, [CoreExpr])
 624 appSpecial env n ctxt args
 625   = go n args
 626   where
 627     arg_env = vanillaCtxt env
 628
 629     go n [] = (emptyDetails, [])        -- Too few args
 630
 631     go 1 (arg:args)                     -- The magic arg
 632       = case occAnal (setCtxt arg_env ctxt) arg of      { (arg_uds, arg') ->
 633         case occAnalArgs env args of                    { (args_uds, args') ->
 634         (combineUsageDetails arg_uds args_uds, arg':args') }}
 635
 636     go n (arg:args)
 637       = case occAnal arg_env arg of     { (arg_uds, arg') ->
 638         case go (n-1) args of           { (args_uds, args') ->
 639         (combineUsageDetails arg_uds args_uds, arg':args') }}
 640 \end{code}
 641
 642
 643 Case alternatives
 644 ~~~~~~~~~~~~~~~~~
 645 If the case binder occurs at all, the other binders effectively do too.
 646 For example
 647         case e of x { (a,b) -> rhs }
 648 is rather like
 649         let x = (a,b) in rhs
 650 If e turns out to be (e1,e2) we indeed get something like
 651         let a = e1; b = e2; x = (a,b) in rhs
 652
 653 \begin{code}
 654 occAnalAlt env case_bndr (con, bndrs, rhs)
 655   = case occAnal (env `addNewCands` bndrs) rhs of { (rhs_usage, rhs') ->
 656     let
 657         (final_usage, tagged_bndrs) = tagBinders rhs_usage bndrs
 658         final_bndrs | case_bndr `elemVarEnv` final_usage = bndrs
 659                     | otherwise                         = tagged_bndrs
 660                 -- Leave the binders untagged if the case
 661                 -- binder occurs at all; see note above
 662     in
 663     (final_usage, (con, final_bndrs, rhs')) }
 664 \end{code}
 665
 666
 667 %************************************************************************
 668 %*                                                                      *
 669 \subsection[OccurAnal-types]{OccEnv}
 670 %*                                                                      *
 671 %************************************************************************
 672
 673 \begin{code}
 674 data OccEnv
 675   = OccEnv IdSet        -- In-scope Ids; we gather info about these only
 676            OccEncl      -- Enclosing context information
 677            CtxtTy       -- Tells about linearity
 678
 679 -- OccEncl is used to control whether to inline into constructor arguments
 680 -- For example:
 681 --      x = (p,q)               -- Don't inline p or q
 682 --      y = /\a -> (p a, q a)   -- Still don't inline p or q
 683 --      z = f (p,q)             -- Do inline p,q; it may make a rule fire
 684 -- So OccEncl tells enought about the context to know what to do when
 685 -- we encounter a contructor application or PAP.
 686
 687 data OccEncl
 688   = OccRhs              -- RHS of let(rec), albeit perhaps inside a type lambda
 689                         -- Don't inline into constructor args here
 690   | OccVanilla          -- Argument of function, body of lambda, scruintee of case etc.
 691                         -- Do inline into constructor args here
 692
 693 type CtxtTy = [Bool]
 694         -- []           No info
 695         --
 696         -- True:ctxt    Analysing a function-valued expression that will be
 697         --                      applied just once
 698         --
 699         -- False:ctxt   Analysing a function-valued expression that may
 700         --                      be applied many times; but when it is,
 701         --                      the CtxtTy inside applies
 702
 703 initOccEnv :: VarSet -> OccEnv
 704 initOccEnv vars = OccEnv vars OccRhs []
 705
 706 isRhsEnv (OccEnv _ OccRhs     _) = True
 707 isRhsEnv (OccEnv _ OccVanilla _) = False
 708
 709 isCandidate :: OccEnv -> Id -> Bool
 710 isCandidate (OccEnv cands encl _) id = id `elemVarSet` cands
 711
 712 addNewCands :: OccEnv -> [Id] -> OccEnv
 713 addNewCands (OccEnv cands encl ctxt) ids
 714   = OccEnv (extendVarSetList cands ids) encl ctxt
 715
 716 addNewCand :: OccEnv -> Id -> OccEnv
 717 addNewCand (OccEnv cands encl ctxt) id
 718   = OccEnv (extendVarSet cands id) encl ctxt
 719
 720 setCtxt :: OccEnv -> CtxtTy -> OccEnv
 721 setCtxt (OccEnv cands encl _) ctxt = OccEnv cands encl ctxt
 722
 723 oneShotGroup :: OccEnv -> [CoreBndr] -> (Bool, OccEnv, [CoreBndr])
 724         -- True <=> this is a one-shot linear lambda group
 725         -- The [CoreBndr] are the binders.
 726
 727         -- The result binders have one-shot-ness set that they might not have had originally.
 728         -- This happens in (build (\cn -> e)).  Here the occurrence analyser
 729         -- linearity context knows that c,n are one-shot, and it records that fact in
 730         -- the binder. This is useful to guide subsequent float-in/float-out tranformations
 731
 732 oneShotGroup (OccEnv cands encl ctxt) bndrs
 733   = case go ctxt bndrs [] of
 734         (new_ctxt, new_bndrs) -> (all is_one_shot new_bndrs, OccEnv cands encl new_ctxt, new_bndrs)
 735   where
 736     is_one_shot b = isId b && isOneShotBndr b
 737
 738     go ctxt [] rev_bndrs = (ctxt, reverse rev_bndrs)
 739
 740     go (lin_ctxt:ctxt) (bndr:bndrs) rev_bndrs
 741         | isId bndr = go ctxt bndrs (bndr':rev_bndrs)
 742         where
 743           bndr' | lin_ctxt  = setOneShotLambda bndr
 744                 | otherwise = bndr
 745
 746     go ctxt (bndr:bndrs) rev_bndrs = go ctxt bndrs (bndr:rev_bndrs)
 747
 748
 749 vanillaCtxt (OccEnv cands _ _) = OccEnv cands OccVanilla []
 750 rhsCtxt     (OccEnv cands _ _) = OccEnv cands OccRhs     []
 751
 752 addAppCtxt (OccEnv cands encl ctxt) args
 753   = OccEnv cands encl (replicate (valArgCount args) True ++ ctxt)
 754 \end{code}
 755
 756 %************************************************************************
 757 %*                                                                      *
 758 \subsection[OccurAnal-types]{OccEnv}
 759 %*                                                                      *
 760 %************************************************************************
 761
 762 \begin{code}
 763 type UsageDetails = IdEnv OccInfo       -- A finite map from ids to their usage
 764
 765 combineUsageDetails, combineAltsUsageDetails
 766         :: UsageDetails -> UsageDetails -> UsageDetails
 767
 768 combineUsageDetails usage1 usage2
 769   = plusVarEnv_C addOccInfo usage1 usage2
 770
 771 combineAltsUsageDetails usage1 usage2
 772   = plusVarEnv_C orOccInfo usage1 usage2
 773
 774 addOneOcc :: UsageDetails -> Id -> OccInfo -> UsageDetails
 775 addOneOcc usage id info
 776   = plusVarEnv_C addOccInfo usage (unitVarEnv id info)
 777         -- ToDo: make this more efficient
 778
 779 emptyDetails = (emptyVarEnv :: UsageDetails)
 780
 781 usedIn :: Id -> UsageDetails -> Bool
 782 v `usedIn` details =  isExportedId v || v `elemVarEnv` details
 783
 784 tagBinders :: UsageDetails          -- Of scope
 785            -> [Id]                  -- Binders
 786            -> (UsageDetails,        -- Details with binders removed
 787               [IdWithOccInfo])    -- Tagged binders
 788
 789 tagBinders usage binders
 790  = let
 791      usage' = usage `delVarEnvList` binders
 792      uss    = map (setBinderOcc usage) binders
 793    in
 794    usage' `seq` (usage', uss)
 795
 796 tagBinder :: UsageDetails           -- Of scope
 797           -> Id                     -- Binders
 798           -> (UsageDetails,         -- Details with binders removed
 799               IdWithOccInfo)        -- Tagged binders
 800
 801 tagBinder usage binder
 802  = let
 803      usage'  = usage `delVarEnv` binder
 804      binder' = setBinderOcc usage binder
 805    in
 806    usage' `seq` (usage', binder')
 807
 808 setBinderOcc :: UsageDetails -> CoreBndr -> CoreBndr
 809 setBinderOcc usage bndr
 810   | isTyVar bndr      = bndr
 811   | isExportedId bndr = case idOccInfo bndr of
 812                           NoOccInfo -> bndr
 813                           other     -> setIdOccInfo bndr NoOccInfo
 814             -- Don't use local usage info for visible-elsewhere things
 815             -- BUT *do* erase any IAmALoopBreaker annotation, because we're
 816             -- about to re-generate it and it shouldn't be "sticky"
 817
 818   | otherwise = setIdOccInfo bndr occ_info
 819   where
 820     occ_info = lookupVarEnv usage bndr `orElse` IAmDead
 821 \end{code}
 822
 823
 824 %************************************************************************
 825 %*                                                                      *
 826 \subsection{Operations over OccInfo}
 827 %*                                                                      *
 828 %************************************************************************
 829
 830 \begin{code}
 831 oneOcc :: OccInfo
 832 oneOcc = OneOcc False True
 833
 834 markMany, markInsideLam, markInsideSCC :: OccInfo -> OccInfo
 835
 836 markMany IAmDead = IAmDead
 837 markMany other   = NoOccInfo
 838
 839 markInsideSCC occ = markMany occ
 840
 841 markInsideLam (OneOcc _ one_br) = OneOcc True one_br
 842 markInsideLam occ               = occ
 843
 844 addOccInfo, orOccInfo :: OccInfo -> OccInfo -> OccInfo
 845
 846 addOccInfo IAmDead info2 = info2
 847 addOccInfo info1 IAmDead = info1
 848 addOccInfo info1 info2   = NoOccInfo
 849
 850 -- (orOccInfo orig new) is used
 851 -- when combining occurrence info from branches of a case
 852
 853 orOccInfo IAmDead info2 = info2
 854 orOccInfo info1 IAmDead = info1
 855 orOccInfo (OneOcc in_lam1 one_branch1)
 856           (OneOcc in_lam2 one_branch2)
 857   = OneOcc (in_lam1 || in_lam2)
 858            False        -- False, because it occurs in both branches
 859
 860 orOccInfo info1 info2 = NoOccInfo
 861 \end{code}