OccurAnal.lhs 27.2 KB
Newer Older
1
%
2
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3
4
5
6
7
8
9
%
%************************************************************************
%*									*
\section[OccurAnal]{Occurrence analysis pass}
%*									*
%************************************************************************

10
11
The occurrence analyser re-typechecks a core expression, returning a new
core expression with (hopefully) improved usage information.
12
13
14

\begin{code}
module OccurAnal (
15
	occurAnalysePgm, occurAnalyseExpr
16
17
    ) where

18
#include "HsVersions.h"
19
20

import CoreSyn
21
import CoreFVs		( idRuleVars )
22
import CoreUtils	( exprIsTrivial, isDefaultAlt )
23
import Id		( isDataConWorkId, isOneShotBndr, setOneShotLambda, 
24
			  idOccInfo, setIdOccInfo, isLocalId,
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
25
			  isExportedId, idArity, idHasRules,
26
			  idType, idUnique, Id
27
			)
28
import BasicTypes	( OccInfo(..), isOneOcc, InterestingCxt )
29
30
31
32

import VarSet
import VarEnv

33
34
import Type		( isFunTy, dropForAlls )
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
42
43
44
45
46
47
48
49
50
51
52
\end{code}


%************************************************************************
%*									*
\subsection[OccurAnal-main]{Counting occurrences: main function}
%*									*
%************************************************************************

Here's the externally-callable interface:

\begin{code}
53
54
occurAnalysePgm :: [CoreBind] -> [CoreBind]
occurAnalysePgm binds
55
  = snd (go initOccEnv binds)
56
57
58
59
60
61
62
  where
    go :: OccEnv -> [CoreBind] -> (UsageDetails, [CoreBind])
    go env [] 
	= (emptyDetails, [])
    go env (bind:binds) 
	= (final_usage, bind' ++ binds')
	where
63
	   (bs_usage, binds')   = go env binds
64
65
	   (final_usage, bind') = occAnalBind env bind bs_usage

66
67
68
occurAnalyseExpr :: CoreExpr -> CoreExpr
	-- Do occurrence analysis, and discard occurence info returned
occurAnalyseExpr expr = snd (occAnal initOccEnv expr)
69
70
\end{code}

71

72
73
74
75
76
77
78
79
80
81
%************************************************************************
%*									*
\subsection[OccurAnal-main]{Counting occurrences: main function}
%*									*
%************************************************************************

Bindings
~~~~~~~~

\begin{code}
82
83
type IdWithOccInfo = Id			-- An Id with fresh PragmaInfo attached

84
type Node details = (details, Unique, [Unique])	-- The Ints are gotten from the Unique,
sof's avatar
sof committed
85
						-- which is gotten from the Id.
86
87
type Details1	  = (Id, UsageDetails, CoreExpr)
type Details2	  = (IdWithOccInfo, CoreExpr)
sof's avatar
sof committed
88
89


90
occAnalBind :: OccEnv
91
	    -> CoreBind
92
93
	    -> UsageDetails		-- Usage details of scope
	    -> (UsageDetails,		-- Of the whole let(rec)
94
		[CoreBind])
95

96
occAnalBind env (NonRec binder rhs) body_usage
97
  | not (binder `usedIn` body_usage)		-- It's not mentioned
98
99
100
  = (body_usage, [])

  | otherwise			-- It's mentioned in the body
101
  = (final_body_usage `combineUsageDetails` rhs_usage,
102
     [NonRec tagged_binder rhs'])
103
104

  where
105
    (final_body_usage, tagged_binder) = tagBinder body_usage binder
106
    (rhs_usage, rhs')		      = occAnalRhs env tagged_binder rhs
107
108
109
110
111
112
113
114
115
\end{code}

Dropping dead code for recursive bindings is done in a very simple way:

	the entire set of bindings is dropped if none of its binders are
	mentioned in its body; otherwise none are.

This seems to miss an obvious improvement.
@
116
117
118
	letrec  f = ...g...
		g = ...f...
	in
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
	...g...

===>

	letrec f = ...g...
	       g = ...(...g...)...
	in
	...g...
@

Now @f@ is unused. But dependency analysis will sort this out into a
@letrec@ for @g@ and a @let@ for @f@, and then @f@ will get dropped.
It isn't easy to do a perfect job in one blow.  Consider

@
	letrec f = ...g...
	       g = ...h...
	       h = ...k...
	       k = ...m...
	       m = ...m...
	in
	...m...
@


\begin{code}
145
occAnalBind env (Rec pairs) body_usage
sof's avatar
sof committed
146
  = foldr (_scc_ "occAnalBind.dofinal" do_final_bind) (body_usage, []) sccs
147
  where
sof's avatar
sof committed
148
    analysed_pairs :: [Details1]
149
    analysed_pairs  = [ (bndr, rhs_usage, rhs')
150
		      | (bndr, rhs) <- pairs,
151
			let (rhs_usage, rhs') = occAnalRhs env bndr rhs
152
		      ]
153

sof's avatar
sof committed
154
155
    sccs :: [SCC (Node Details1)]
    sccs = _scc_ "occAnalBind.scc" stronglyConnCompR edges
156
157
158


    ---- stuff for dependency analysis of binds -------------------------------
sof's avatar
sof committed
159
160
    edges :: [Node Details1]
    edges = _scc_ "occAnalBind.assoc"
161
	    [ (details, idUnique id, edges_from rhs_usage)
162
	    | details@(id, rhs_usage, rhs) <- analysed_pairs
sof's avatar
sof committed
163
164
165
166
167
168
169
170
	    ]

	-- (a -> b) means a mentions b
	-- Given the usage details (a UFM that gives occ info for each free var of
	-- the RHS) we can get the list of free vars -- or rather their Int keys --
	-- by just extracting the keys from the finite map.  Grimy, but fast.
	-- Previously we had this:
	-- 	[ bndr | bndr <- bndrs,
171
	--		 maybeToBool (lookupVarEnv rhs_usage bndr)]
sof's avatar
sof committed
172
173
	-- which has n**2 cost, and this meant that edges_from alone 
	-- consumed 10% of total runtime!
174
    edges_from :: UsageDetails -> [Unique]
sof's avatar
sof committed
175
176
    edges_from rhs_usage = _scc_ "occAnalBind.edges_from"
			   keysUFM rhs_usage
177
178
179

    ---- stuff to "re-constitute" bindings from dependency-analysis info ------

sof's avatar
sof committed
180
	-- Non-recursive SCC
181
    do_final_bind (AcyclicSCC ((bndr, rhs_usage, rhs'), _, _)) (body_usage, binds_so_far)
182
      | not (bndr `usedIn` body_usage)
sof's avatar
sof committed
183
      = (body_usage, binds_so_far)			-- Dead code
184
185
      | otherwise
      = (combined_usage, new_bind : binds_so_far)	
186
      where
sof's avatar
sof committed
187
188
189
190
191
192
	total_usage       	      = combineUsageDetails body_usage rhs_usage
	(combined_usage, tagged_bndr) = tagBinder total_usage bndr
	new_bind		      = NonRec tagged_bndr rhs'

	-- Recursive SCC
    do_final_bind (CyclicSCC cycle) (body_usage, binds_so_far)
193
      | not (any (`usedIn` body_usage) bndrs)		-- NB: look at body_usage, not total_usage
sof's avatar
sof committed
194
      = (body_usage, binds_so_far)			-- Dead code
195
196
      | otherwise
      = (combined_usage, final_bind:binds_so_far)
197
      where
198
199
200
201
202
203
204
205
206
	details			       = [details   | (details, _, _) <- cycle]
	bndrs			       = [bndr      | (bndr, _, _)      <- details]
	rhs_usages		       = [rhs_usage | (_, rhs_usage, _) <- details]
	total_usage		       = foldr combineUsageDetails body_usage rhs_usages
	(combined_usage, tagged_bndrs) = tagBinders total_usage bndrs
	final_bind		       = Rec (reOrderRec env new_cycle)

	new_cycle = CyclicSCC (zipWithEqual "occAnalBind" mk_new_bind tagged_bndrs cycle)
	mk_new_bind tagged_bndr ((_, _, rhs'), key, keys) = ((tagged_bndr, rhs'), key, keys)
sof's avatar
sof committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
\end{code}

@reOrderRec@ is applied to the list of (binder,rhs) pairs for a cyclic
strongly connected component (there's guaranteed to be a cycle).  It returns the
same pairs, but 
	a) in a better order,
	b) with some of the Ids having a IMustNotBeINLINEd pragma

The "no-inline" Ids are sufficient to break all cycles in the SCC.  This means
that the simplifier can guarantee not to loop provided it never records an inlining
for these no-inline guys.

Furthermore, the order of the binds is such that if we neglect dependencies
on the no-inline Ids then the binds are topologically sorted.  This means
that the simplifier will generally do a good job if it works from top bottom,
recording inlinings for any Ids which aren't marked as "no-inline" as it goes.

224
225
226
227
==============
[June 98: I don't understand the following paragraphs, and I've 
	  changed the a=b case again so that it isn't a special case any more.]

sof's avatar
sof committed
228
Here's a case that bit me:
229

sof's avatar
sof committed
230
231
232
233
234
235
236
	letrec
		a = b
		b = \x. BIG
	in
	...a...a...a....

Re-ordering doesn't change the order of bindings, but there was no loop-breaker.
sof's avatar
sof committed
237

sof's avatar
sof committed
238
239
My solution was to make a=b bindings record b as Many, rather like INLINE bindings.
Perhaps something cleverer would suffice.
240
===============
sof's avatar
sof committed
241

sof's avatar
sof committed
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
You might think that you can prevent non-termination simply by making
sure that we simplify a recursive binding's RHS in an environment that
simply clones the recursive Id.  But no.  Consider

		letrec f = \x -> let z = f x' in ...

		in
		let n = f y
		in
		case n of { ... }

We bind n to its *simplified* RHS, we then *re-simplify* it when
we inline n.  Then we may well inline f; and then the same thing
happens with z!

I don't think it's possible to prevent non-termination by environment
manipulation in this way.  Apart from anything else, successive
iterations of the simplifier may unroll recursive loops in cases like
that above.  The idea of beaking every recursive loop with an
IMustNotBeINLINEd pragma is much much better.


sof's avatar
sof committed
264
265
266
267
268
269
270
271
272
\begin{code}
reOrderRec
	:: OccEnv
	-> SCC (Node Details2)
	-> [Details2]
			-- Sorted into a plausible order.  Enough of the Ids have
			--	dontINLINE pragmas that there are no loops left.

	-- Non-recursive case
273
reOrderRec env (AcyclicSCC (bind, _, _)) = [bind]
sof's avatar
sof committed
274
275
276

	-- Common case of simple self-recursion
reOrderRec env (CyclicSCC [bind])
277
  = [(setIdOccInfo tagged_bndr IAmALoopBreaker, rhs)]
sof's avatar
sof committed
278
  where
279
    ((tagged_bndr, rhs), _, _) = bind
sof's avatar
sof committed
280

281
reOrderRec env (CyclicSCC (bind : binds))
sof's avatar
sof committed
282
283
284
285
  = 	-- Choose a loop breaker, mark it no-inline,
	-- do SCC analysis on the rest, and recursively sort them out
    concat (map (reOrderRec env) (stronglyConnCompR unchosen))
    ++ 
286
    [(setIdOccInfo tagged_bndr IAmALoopBreaker, rhs)]
sof's avatar
sof committed
287
288

  where
289
290
291
292
293
294
295
296
297
298
299
300
301
302
    (chosen_pair, unchosen) = choose_loop_breaker bind (score bind) [] binds
    (tagged_bndr, rhs)      = chosen_pair

	-- This loop looks for the bind with the lowest score
	-- to pick as the loop  breaker.  The rest accumulate in 
    choose_loop_breaker (details,_,_) loop_sc acc []
	= (details, acc)	-- Done

    choose_loop_breaker loop_bind loop_sc acc (bind : binds)
	| sc < loop_sc	-- Lower score so pick this new one
	= choose_loop_breaker bind sc (loop_bind : acc) binds

	| otherwise	-- No lower so don't pick it
	= choose_loop_breaker loop_bind loop_sc (bind : acc) binds
sof's avatar
sof committed
303
	where
304
305
306
307
	  sc = score bind
	  
    score :: Node Details2 -> Int	-- Higher score => less likely to be picked as loop breaker
    score ((bndr, rhs), _, _)
308
309
310
311
312
313
314
315
316
317
318
319
320
321
	| exprIsTrivial rhs 	   = 4	-- Practically certain to be inlined
		-- Used to have also: && not (isExportedId bndr)
		-- But I found this sometimes cost an extra iteration when we have
		--	rec { d = (a,b); a = ...df...; b = ...df...; df = d }
		-- where df is the exported dictionary. Then df makes a really
		-- bad choice for loop breaker
	  
	| not_fun_ty (idType bndr) = 3	-- Data types help with cases
		-- This used to have a lower score than inlineCandidate, but
		-- it's *really* helpful if dictionaries get inlined fast,
		-- so I'm experimenting with giving higher priority to data-typed things

	| inlineCandidate bndr rhs = 2	-- Likely to be inlined

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
322
	| idHasRules bndr = 1
323
324
		-- Avoid things with specialisations; we'd like
		-- to take advantage of them in the subsequent bindings
325

326
327
	| otherwise = 0

328
329
    inlineCandidate :: Id -> CoreExpr -> Bool
    inlineCandidate id (Note InlineMe _) = True
330
    inlineCandidate id rhs	         = isOneOcc (idOccInfo id)
331
332

	-- Real example (the Enum Ordering instance from PrelBase):
sof's avatar
sof committed
333
334
335
336
337
338
339
	--	rec	f = \ x -> case d of (p,q,r) -> p x
	--		g = \ x -> case d of (p,q,r) -> q x
	--		d = (v, f, g)
	--
	-- Here, f and g occur just once; but we can't inline them into d.
	-- On the other hand we *could* simplify those case expressions if
	-- we didn't stupidly choose d as the loop breaker.
340
	-- But we won't because constructor args are marked "Many".
sof's avatar
sof committed
341

342
    not_fun_ty ty = not (isFunTy (dropForAlls ty))
343
344
345
346
347
348
349
350
351
\end{code}

@occAnalRhs@ deals with the question of bindings where the Id is marked
by an INLINE pragma.  For these we record that anything which occurs
in its RHS occurs many times.  This pessimistically assumes that ths
inlined binder also occurs many times in its scope, but if it doesn't
we'll catch it next time round.  At worst this costs an extra simplifier pass.
ToDo: try using the occurrence info for the inline'd binder.

sof's avatar
sof committed
352
[March 97] We do the same for atomic RHSs.  Reason: see notes with reOrderRec.
353
[June 98, SLPJ]  I've undone this change; I don't understand it.  See notes with reOrderRec.
sof's avatar
sof committed
354

355

356
357
\begin{code}
occAnalRhs :: OccEnv
sof's avatar
sof committed
358
	   -> Id -> CoreExpr	-- Binder and rhs
359
360
				-- For non-recs the binder is alrady tagged
				-- with occurrence info
361
	   -> (UsageDetails, CoreExpr)
362
363

occAnalRhs env id rhs
364
  = (final_usage, rhs')
365
  where
366
367
    (rhs_usage, rhs') = occAnal ctxt rhs
    ctxt | certainly_inline id = env
368
	 | otherwise	       = rhsCtxt
369
370
	-- Note that we generally use an rhsCtxt.  This tells the occ anal n
	-- that it's looking at an RHS, which has an effect in occAnalApp
371
372
373
374
375
	--
	-- But there's a problem.  Consider
	--	x1 = a0 : []
	--	x2 = a1 : x1
	--	x3 = a2 : x2
376
	--	g  = f x3
377
378
379
380
381
	-- First time round, it looks as if x1 and x2 occur as an arg of a 
	-- let-bound constructor ==> give them a many-occurrence.
	-- But then x3 is inlined (unconditionally as it happens) and
	-- next time round, x2 will be, and the next time round x1 will be
	-- Result: multiple simplifier iterations.  Sigh.  
382
383
384
	-- Crude solution: use rhsCtxt for things that occur just once...

    certainly_inline id = case idOccInfo id of
385
386
			    OneOcc in_lam one_br _ -> not in_lam && one_br
			    other		   -> False
387

388
389
390
391
392
393
	-- [March 98] A new wrinkle is that if the binder has specialisations inside
	-- it then we count the specialised Ids as "extra rhs's".  That way
	-- the "parent" keeps the specialised "children" alive.  If the parent
	-- dies (because it isn't referenced any more), then the children will
	-- die too unless they are already referenced directly.

394
395
396
397
398
399
400
    final_usage = addRuleUsage rhs_usage id

addRuleUsage :: UsageDetails -> Id -> UsageDetails
-- Add the usage from RULES in Id to the usage
addRuleUsage usage id
  = foldVarSet add usage (idRuleVars id)
  where
401
    add v u = addOneOcc u v NoOccInfo		-- Give a non-committal binder info
402
403
						-- (i.e manyOcc) because many copies
						-- of the specialised thing can appear
404
405
406
407
408
409
\end{code}

Expressions
~~~~~~~~~~~
\begin{code}
occAnal :: OccEnv
410
	-> CoreExpr
411
	-> (UsageDetails,	-- Gives info only about the "interesting" Ids
412
	    CoreExpr)
413

414
occAnal env (Type t)  = (emptyDetails, Type t)
415
occAnal env (Var v)   = (mkOneOcc env v False, Var v)
416
417
    -- At one stage, I gathered the idRuleVars for v here too,
    -- which in a way is the right thing to do.
418
    -- Btu that went wrong right after specialisation, when
419
420
421
    -- the *occurrences* of the overloaded function didn't have any
    -- rules in them, so the *specialised* versions looked as if they
    -- weren't used at all.
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
\end{code}

We regard variables that occur as constructor arguments as "dangerousToDup":

\begin{verbatim}
module A where
f x = let y = expensive x in 
      let z = (True,y) in 
      (case z of {(p,q)->q}, case z of {(p,q)->q})
\end{verbatim}

We feel free to duplicate the WHNF (True,y), but that means
that y may be duplicated thereby.

If we aren't careful we duplicate the (expensive x) call!
Constructors are rather like lambdas in this way.

\begin{code}
440
occAnal env expr@(Lit lit) = (emptyDetails, expr)
441
\end{code}
442

443
\begin{code}
444
445
446
447
448
occAnal env (Note InlineMe body)
  = case occAnal env body of { (usage, body') -> 
    (mapVarEnv markMany usage, Note InlineMe body')
    }

449
occAnal env (Note note@(SCC cc) body)
450
451
452
  = case occAnal env body of { (usage, body') ->
    (mapVarEnv markInsideSCC usage, Note note body')
    }
453

454
occAnal env (Note note body)
455
456
457
458
  = case occAnal env body of { (usage, body') ->
    (usage, Note note body')
    }
\end{code}
459

460
\begin{code}
461
occAnal env app@(App fun arg)
462
  = occAnalApp env (collectArgs app) False
463

464
465
466
467
468
469
470
471
-- Ignore type variables altogether
--   (a) occurrences inside type lambdas only not marked as InsideLam
--   (b) type variables not in environment

occAnal env expr@(Lam x body) | isTyVar x
  = case occAnal env body of { (body_usage, body') ->
    (body_usage, Lam x body')
    }
472

473
474
475
476
477
-- For value lambdas we do a special hack.  Consider
-- 	(\x. \y. ...x...)
-- If we did nothing, x is used inside the \y, so would be marked
-- as dangerous to dup.  But in the common case where the abstraction
-- is applied to two arguments this is over-pessimistic.
478
479
480
481
482
-- So instead, we just mark each binder with its occurrence
-- info in the *body* of the multiple lambda.
-- Then, the simplifier is careful when partially applying lambdas.

occAnal env expr@(Lam _ _)
483
  = case occAnal env_body body of { (body_usage, body') ->
484
485
    let
        (final_usage, tagged_binders) = tagBinders body_usage binders
486
487
488
489
	--	URGH!  Sept 99: we don't seem to be able to use binders' here, because
	--	we get linear-typed things in the resulting program that we can't handle yet.
	--	(e.g. PrelShow)  TODO 

490
491
492
493
	really_final_usage = if linear then
				final_usage
			     else
				mapVarEnv markInsideLam final_usage
494
    in
495
    (really_final_usage,
496
     mkLams tagged_binders body') }
497
  where
498
499
500
501
502
    env_body	    = vanillaCtxt			-- Body is (no longer) an RhsContext
    (binders, body) = collectBinders expr
    binders' 	    = oneShotGroup env binders
    linear	    = all is_one_shot binders'
    is_one_shot b   = isId b && isOneShotBndr b
503

504
occAnal env (Case scrut bndr ty alts)
505
506
  = case occ_anal_scrut scrut alts		    of { (scrut_usage, scrut') ->
    case mapAndUnzip (occAnalAlt alt_env bndr) alts of { (alts_usage_s, alts')   -> 
507
508
    let
	alts_usage  = foldr1 combineAltsUsageDetails alts_usage_s
509
510
	alts_usage' = addCaseBndrUsage alts_usage
	(alts_usage1, tagged_bndr) = tagBinder alts_usage' bndr
511
512
        total_usage = scrut_usage `combineUsageDetails` alts_usage1
    in
513
    total_usage `seq` (total_usage, Case scrut' tagged_bndr ty alts') }}
514
  where
515
516
517
518
519
520
521
522
523
524
525
	-- The case binder gets a usage of either "many" or "dead", never "one".
	-- Reason: we like to inline single occurrences, to eliminate a binding,
	-- but inlining a case binder *doesn't* eliminate a binding.
	-- We *don't* want to transform
	--	case x of w { (p,q) -> f w }
	-- into
	--	case x of w { (p,q) -> f (p,q) }
    addCaseBndrUsage usage = case lookupVarEnv usage bndr of
				Nothing  -> usage
				Just occ -> extendVarEnv usage bndr (markMany occ)

526
527
528
529
    alt_env = setVanillaCtxt env
	-- Consider 	x = case v of { True -> (p,q); ... }
	-- Then it's fine to inline p and q

530
531
532
533
534
535
    occ_anal_scrut (Var v) (alt1 : other_alts)
				| not (null other_alts) || not (isDefaultAlt alt1)
			        = (mkOneOcc env v True, Var v)
    occ_anal_scrut scrut alts   = occAnal vanillaCtxt scrut
					-- No need for rhsCtxt

536
occAnal env (Let bind body)
537
  = case occAnal env body     	         of { (body_usage, body') ->
538
    case occAnalBind env bind body_usage of { (final_usage, new_binds) ->
539
       (final_usage, mkLets new_binds body') }}
540
541

occAnalArgs env args
542
  = case mapAndUnzip (occAnal arg_env) args of	{ (arg_uds_s, args') ->
543
    (foldr combineUsageDetails emptyDetails arg_uds_s, args')}
544
  where
545
    arg_env = vanillaCtxt
546
547
\end{code}

548
549
550
551
Applications are dealt with specially because we want
the "build hack" to work.

\begin{code}
552
occAnalApp env (Var fun, args) is_rhs
553
554
  = case args_stuff of { (args_uds, args') ->
    let
555
556
557
558
559
560
561
562
563
	-- We mark the free vars of the argument of a constructor or PAP 
	-- as "many", if it is the RHS of a let(rec).
	-- This means that nothing gets inlined into a constructor argument
	-- position, which is what we want.  Typically those constructor
	-- arguments are just variables, or trivial expressions.
	--
	-- This is the *whole point* of the isRhsEnv predicate
        final_args_uds
		| isRhsEnv env,
564
		  isDataConWorkId fun || valArgCount args < idArity fun
565
566
		= mapVarEnv markMany args_uds
		| otherwise = args_uds
567
    in
568
    (fun_uds `combineUsageDetails` final_args_uds, mkApps (Var fun) args') }
569
570
  where
    fun_uniq = idUnique fun
571
    fun_uds  = mkOneOcc env fun (valArgCount args > 0)
572
573

		-- Hack for build, fold, runST
574
575
576
    args_stuff	| fun_uniq == buildIdKey    = appSpecial env 2 [True,True]  args
		| fun_uniq == augmentIdKey  = appSpecial env 2 [True,True]  args
		| fun_uniq == foldrIdKey    = appSpecial env 3 [False,True] args
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
		| fun_uniq == runSTRepIdKey = appSpecial env 2 [True]	    args
			-- (foldr k z xs) may call k many times, but it never
			-- shares a partial application of k; hence [False,True]
			-- This means we can optimise
			--	foldr (\x -> let v = ...x... in \y -> ...v...) z xs
			-- by floating in the v

		| otherwise = occAnalArgs env args


occAnalApp env (fun, args) is_rhs
  = case occAnal (addAppCtxt env args) fun of	{ (fun_uds, fun') ->
	-- The addAppCtxt is a bit cunning.  One iteration of the simplifier
	-- often leaves behind beta redexs like
	--	(\x y -> e) a1 a2
	-- Here we would like to mark x,y as one-shot, and treat the whole
	-- thing much like a let.  We do this by pushing some True items
	-- onto the context stack.

    case occAnalArgs env args of	{ (args_uds, args') ->
597
598
599
600
601
    let
	final_uds = fun_uds `combineUsageDetails` args_uds
    in
    (final_uds, mkApps fun' args') }}
    
602
603
604
605
appSpecial :: OccEnv 
	   -> Int -> CtxtTy	-- Argument number, and context to use for it
	   -> [CoreExpr]
	   -> (UsageDetails, [CoreExpr])
606
607
608
appSpecial env n ctxt args
  = go n args
  where
609
    arg_env = vanillaCtxt
610

611
612
613
    go n [] = (emptyDetails, [])	-- Too few args

    go 1 (arg:args)			-- The magic arg
614
615
      = case occAnal (setCtxt arg_env ctxt) arg of	{ (arg_uds, arg') ->
	case occAnalArgs env args of			{ (args_uds, args') ->
616
617
618
	(combineUsageDetails arg_uds args_uds, arg':args') }}
    
    go n (arg:args)
619
      = case occAnal arg_env arg of	{ (arg_uds, arg') ->
620
621
622
623
624
	case go (n-1) args of		{ (args_uds, args') ->
	(combineUsageDetails arg_uds args_uds, arg':args') }}
\end{code}

    
625
626
Case alternatives
~~~~~~~~~~~~~~~~~
627
628
629
630
631
632
633
634
If the case binder occurs at all, the other binders effectively do too.  
For example
	case e of x { (a,b) -> rhs }
is rather like
	let x = (a,b) in rhs
If e turns out to be (e1,e2) we indeed get something like
	let a = e1; b = e2; x = (a,b) in rhs

635
\begin{code}
636
occAnalAlt env case_bndr (con, bndrs, rhs)
637
  = case occAnal env rhs of { (rhs_usage, rhs') ->
638
639
    let
        (final_usage, tagged_bndrs) = tagBinders rhs_usage bndrs
640
641
642
643
	final_bndrs | case_bndr `elemVarEnv` final_usage = bndrs
		    | otherwise				= tagged_bndrs
		-- Leave the binders untagged if the case 
		-- binder occurs at all; see note above
644
    in
645
    (final_usage, (con, final_bndrs, rhs')) }
646
\end{code}
647
648
649
650


%************************************************************************
%*									*
651
\subsection[OccurAnal-types]{OccEnv}
652
653
654
655
%*									*
%************************************************************************

\begin{code}
656
data OccEnv
657
  = OccEnv OccEncl	-- Enclosing context information
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
	   CtxtTy	-- Tells about linearity

-- OccEncl is used to control whether to inline into constructor arguments
-- For example:
--	x = (p,q)		-- Don't inline p or q
--	y = /\a -> (p a, q a)	-- Still don't inline p or q
--	z = f (p,q)		-- Do inline p,q; it may make a rule fire
-- So OccEncl tells enought about the context to know what to do when
-- we encounter a contructor application or PAP.

data OccEncl
  = OccRhs 		-- RHS of let(rec), albeit perhaps inside a type lambda
			-- Don't inline into constructor args here
  | OccVanilla		-- Argument of function, body of lambda, scruintee of case etc.
			-- Do inline into constructor args here
673

674
675
676
677
678
679
680
681
682
type CtxtTy = [Bool]
	-- []	 	No info
	--
	-- True:ctxt  	Analysing a function-valued expression that will be
	--			applied just once
	--
	-- False:ctxt	Analysing a function-valued expression that may
	--			be applied many times; but when it is, 
	--			the CtxtTy inside applies
683

684
685
initOccEnv :: OccEnv
initOccEnv = OccEnv OccRhs []
686

687
688
vanillaCtxt = OccEnv OccVanilla []
rhsCtxt     = OccEnv OccRhs     []
689

690
691
isRhsEnv (OccEnv OccRhs     _) = True
isRhsEnv (OccEnv OccVanilla _) = False
692

693
694
695
696
setVanillaCtxt :: OccEnv -> OccEnv
setVanillaCtxt (OccEnv OccRhs ctxt_ty) = OccEnv OccVanilla ctxt_ty
setVanillaCtxt other_env	       = other_env

697
setCtxt :: OccEnv -> CtxtTy -> OccEnv
698
setCtxt (OccEnv encl _) ctxt = OccEnv encl ctxt
699

700
oneShotGroup :: OccEnv -> [CoreBndr] -> [CoreBndr]
701
702
703
704
705
	-- The result binders have one-shot-ness set that they might not have had originally.
	-- This happens in (build (\cn -> e)).  Here the occurrence analyser
	-- linearity context knows that c,n are one-shot, and it records that fact in
	-- the binder. This is useful to guide subsequent float-in/float-out tranformations

706
707
oneShotGroup (OccEnv encl ctxt) bndrs 
  = go ctxt bndrs []
708
  where
709
    go ctxt [] rev_bndrs = reverse rev_bndrs
710
711
712
713
714
715
716
717
718

    go (lin_ctxt:ctxt) (bndr:bndrs) rev_bndrs
	| isId bndr = go ctxt bndrs (bndr':rev_bndrs)
	where
	  bndr' | lin_ctxt  = setOneShotLambda bndr
		| otherwise = bndr

    go ctxt (bndr:bndrs) rev_bndrs = go ctxt bndrs (bndr:rev_bndrs)

719
720
addAppCtxt (OccEnv encl ctxt) args 
  = OccEnv encl (replicate (valArgCount args) True ++ ctxt)
721
722
723
724
725
726
727
\end{code}

%************************************************************************
%*									*
\subsection[OccurAnal-types]{OccEnv}
%*									*
%************************************************************************
728

729
\begin{code}
730
type UsageDetails = IdEnv OccInfo	-- A finite map from ids to their usage
731
732
733
734
735

combineUsageDetails, combineAltsUsageDetails
	:: UsageDetails -> UsageDetails -> UsageDetails

combineUsageDetails usage1 usage2
736
  = plusVarEnv_C addOccInfo usage1 usage2
737
738

combineAltsUsageDetails usage1 usage2
739
  = plusVarEnv_C orOccInfo usage1 usage2
740

741
addOneOcc :: UsageDetails -> Id -> OccInfo -> UsageDetails
742
addOneOcc usage id info
743
  = plusVarEnv_C addOccInfo usage (unitVarEnv id info)
744
745
	-- ToDo: make this more efficient

746
emptyDetails = (emptyVarEnv :: UsageDetails)
747

748
usedIn :: Id -> UsageDetails -> Bool
749
v `usedIn` details =  isExportedId v || v `elemVarEnv` details
750

751
752
753
tagBinders :: UsageDetails	    -- Of scope
	   -> [Id]		    -- Binders
	   -> (UsageDetails, 	    -- Details with binders removed
754
755
756
757
758
	      [IdWithOccInfo])    -- Tagged binders

tagBinders usage binders
 = let
     usage' = usage `delVarEnvList` binders
759
     uss    = map (setBinderOcc usage) binders
760
761
762
   in
   usage' `seq` (usage', uss)

763
764
765
tagBinder :: UsageDetails	    -- Of scope
	  -> Id			    -- Binders
	  -> (UsageDetails, 	    -- Details with binders removed
766
767
768
769
770
	      IdWithOccInfo)	    -- Tagged binders

tagBinder usage binder
 = let
     usage'  = usage `delVarEnv` binder
771
     binder' = setBinderOcc usage binder
772
773
774
   in
   usage' `seq` (usage', binder')

775
776
777
setBinderOcc :: UsageDetails -> CoreBndr -> CoreBndr
setBinderOcc usage bndr
  | isTyVar bndr      = bndr
778
779
780
781
782
783
  | isExportedId bndr = case idOccInfo bndr of
			  NoOccInfo -> bndr
			  other     -> setIdOccInfo bndr NoOccInfo
  	    -- Don't use local usage info for visible-elsewhere things
	    -- BUT *do* erase any IAmALoopBreaker annotation, because we're
	    -- about to re-generate it and it shouldn't be "sticky"
784
			  
785
  | otherwise = setIdOccInfo bndr occ_info
786
  where
787
788
789
790
791
792
793
794
795
796
797
    occ_info = lookupVarEnv usage bndr `orElse` IAmDead
\end{code}


%************************************************************************
%*									*
\subsection{Operations over OccInfo}
%*									*
%************************************************************************

\begin{code}
798
799
800
801
mkOneOcc :: OccEnv -> Id -> InterestingCxt -> UsageDetails
mkOneOcc env id int_cxt
  | isLocalId id = unitVarEnv id (OneOcc False True int_cxt)
  | otherwise    = emptyDetails
802
803
804
805
806
807
808
809

markMany, markInsideLam, markInsideSCC :: OccInfo -> OccInfo

markMany IAmDead = IAmDead
markMany other   = NoOccInfo

markInsideSCC occ = markMany occ

810
811
markInsideLam (OneOcc _ one_br int_cxt) = OneOcc True one_br int_cxt
markInsideLam occ		 	= occ
812
813
814
815
816
817
818
819
820
821
822
823

addOccInfo, orOccInfo :: OccInfo -> OccInfo -> OccInfo

addOccInfo IAmDead info2 = info2
addOccInfo info1 IAmDead = info1
addOccInfo info1 info2   = NoOccInfo

-- (orOccInfo orig new) is used
-- when combining occurrence info from branches of a case

orOccInfo IAmDead info2 = info2
orOccInfo info1 IAmDead = info1
824
825
orOccInfo (OneOcc in_lam1 one_branch1 int_cxt1)
	  (OneOcc in_lam2 one_branch2 int_cxt2)
826
827
  = OneOcc (in_lam1 || in_lam2)
	   False	-- False, because it occurs in both branches
828
	   (int_cxt1 && int_cxt2)
829

830
orOccInfo info1 info2 = NoOccInfo
831
\end{code}