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