SpecConstr.lhs 18.8 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
%
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section[SpecConstr]{Specialise over constructors}

\begin{code}
module SpecConstr(
	specConstrProgram	
    ) where

#include "HsVersions.h"

import CoreSyn
import CoreLint		( showPass, endPass )
15
import CoreUtils	( exprType, tcEqExpr, mkPiTypes )
16
import CoreFVs 		( exprsFreeVars )
17
18
import CoreTidy		( tidyRules )
import PprCore		( pprRules )
19
import WwLib		( mkWorkerArgs )
20
21
import DataCon		( dataConRepArity )
import Type		( tyConAppArgs )
22
import Id		( Id, idName, idType, isDataConWorkId_maybe, 
23
			  mkUserLocal, mkSysLocal )
24
25
26
27
import Var		( Var )
import VarEnv
import VarSet
import Name		( nameOccName, nameSrcLoc )
28
import Rules		( addIdSpecialisations, mkLocalRule, rulesOfBinds )
29
30
import OccName		( mkSpecOcc )
import ErrUtils		( dumpIfSet_dyn )
31
import DynFlags	( DynFlags, DynFlag(..) )
32
import BasicTypes	( Activation(..) )
33
34
35
import Outputable

import Maybes		( orElse )
sof's avatar
sof committed
36
import Util		( mapAccumL, lengthAtLeast, notNull )
37
38
39
import List		( nubBy, partition )
import UniqSupply
import Outputable
40
import FastString
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
\end{code}

-----------------------------------------------------
			Game plan
-----------------------------------------------------

Consider
	drop n []     = []
	drop 0 xs     = []
	drop n (x:xs) = drop (n-1) xs

After the first time round, we could pass n unboxed.  This happens in
numerical code too.  Here's what it looks like in Core:

	drop n xs = case xs of
		      []     -> []
		      (y:ys) -> case n of 
				  I# n# -> case n# of
					     0 -> []
					     _ -> drop (I# (n# -# 1#)) xs

Notice that the recursive call has an explicit constructor as argument.
Noticing this, we can make a specialised version of drop
	
	RULE: drop (I# n#) xs ==> drop' n# xs

	drop' n# xs = let n = I# n# in ...orig RHS...

Now the simplifier will apply the specialisation in the rhs of drop', giving

	drop' n# xs = case xs of
		      []     -> []
		      (y:ys) -> case n# of
				  0 -> []
				  _ -> drop (n# -# 1#) xs

Much better!  

We'd also like to catch cases where a parameter is carried along unchanged,
but evaluated each time round the loop:

	f i n = if i>0 || i>n then i else f (i*2) n

Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
In Core, by the time we've w/wd (f is strict in i) we get

	f i# n = case i# ># 0 of
		   False -> I# i#
		   True  -> case n of n' { I# n# ->
			    case i# ># n# of
				False -> I# i#
				True  -> f (i# *# 2#) n'

At the call to f, we see that the argument, n is know to be (I# n#),
and n is evaluated elsewhere in the body of f, so we can play the same
trick as above.  However we don't want to do that if the boxed version
of n is needed (else we'd avoid the eval but pay more for re-boxing n).
So in this case we want that the *only* uses of n are in case statements.


So we look for

* A self-recursive function.  Ignore mutual recursion for now, 
  because it's less common, and the code is simpler for self-recursion.

* EITHER

   a) At a recursive call, one or more parameters is an explicit 
      constructor application
	AND
      That same parameter is scrutinised by a case somewhere in 
      the RHS of the function

  OR

    b) At a recursive call, one or more parameters has an unfolding
       that is an explicit constructor application
	AND
      That same parameter is scrutinised by a case somewhere in 
      the RHS of the function
	AND
      Those are the only uses of the parameter


There's a bit of a complication with type arguments.  If the call
site looks like

	f p = ...f ((:) [a] x xs)...

then our specialised function look like

	f_spec x xs = let p = (:) [a] x xs in ....as before....

This only makes sense if either
  a) the type variable 'a' is in scope at the top of f, or
  b) the type variable 'a' is an argument to f (and hence fs)

Actually, (a) may hold for value arguments too, in which case
we may not want to pass them.  Supose 'x' is in scope at f's
defn, but xs is not.  Then we'd like

	f_spec xs = let p = (:) [a] x xs in ....as before....

Similarly (b) may hold too.  If x is already an argument at the
call, no need to pass it again.

Finally, if 'a' is not in scope at the call site, we could abstract
it as we do the term variables:

	f_spec a x xs = let p = (:) [a] x xs in ...as before...

So the grand plan is:

	* abstract the call site to a constructor-only pattern
	  e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)

	* Find the free variables of the abstracted pattern

	* Pass these variables, less any that are in scope at
	  the fn defn.


NOTICE that we only abstract over variables that are not in scope,
so we're in no danger of shadowing variables used in "higher up"
in f_spec's RHS.


%************************************************************************
%*									*
\subsection{Top level wrapper stuff}
%*									*
%************************************************************************

\begin{code}
specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
specConstrProgram dflags us binds
  = do
	showPass dflags "SpecConstr"

	let (binds', _) = initUs us (go emptyScEnv binds)

	endPass dflags "SpecConstr" Opt_D_dump_spec binds'

	dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
185
		  (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
186
187
188
189
190
191
192
193
194
195
196
197

	return binds'
  where
    go env []	        = returnUs []
    go env (bind:binds) = scBind env bind 	`thenUs` \ (env', _, bind') ->
			  go env' binds 	`thenUs` \ binds' ->
			  returnUs (bind' : binds')
\end{code}


%************************************************************************
%*									*
198
\subsection{Environment: goes downwards}
199
200
201
202
%*									*
%************************************************************************

\begin{code}
203
204
data ScEnv = SCE { scope :: VarEnv HowBound,
			-- Binds all non-top-level variables in scope
205

206
207
208
209
210
211
212
213
		   cons  :: ConstrEnv
	     }

type ConstrEnv = IdEnv (AltCon, [CoreArg])
	-- Variables known to be bound to a constructor
	-- in a particular case alternative

emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
214
215
216

data HowBound = RecFun		-- These are the recursive functions for which 
				-- we seek interesting call patterns
217

218
219
	      | RecArg		-- These are those functions' arguments; we are
				-- interested to see if those arguments are scrutinised
220

221
	      | Other		-- We track all others so we know what's in scope
222
223
224
				-- This is used in spec_one to check what needs to be
				-- passed as a parameter and what is in scope at the 
				-- function definition site
225

226
227
228
229
230
instance Outputable HowBound where
  ppr RecFun = text "RecFun"
  ppr RecArg = text "RecArg"
  ppr Other = text "Other"

231
232
233
234
235
236
237
238
239
lookupScopeEnv env v = lookupVarEnv (scope env) v

extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }

    -- When we encounter
    --	case scrut of b
    --	    C x y -> ...
    -- we want to bind b, and perhaps scrut too, to (C x y)
240
241
242
243
244
extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
  = extendBndrs env (case_bndr : alt_bndrs)

extendCaseBndrs env case_bndr scrut con alt_bndrs
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
  = case scrut of
	Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
		   -- Also forget if the scrutinee is a RecArg, because we're
		   -- now in the branch of a case, and we don't want to
		   -- record a non-scrutinee use of v if we have
		   --	case v of { (a,b) -> ...(f v)... }
		 SCE { scope = extendVarEnv (scope env1) v Other,
		       cons  = extendVarEnv (cons env1)  v (con,args) }
	other -> env1

  where
    env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
		 cons  = extendVarEnv     (cons  env) case_bndr (con,args) }

    args = map Type (tyConAppArgs (idType case_bndr)) ++
	   map varToCoreExpr alt_bndrs

    -- When we encounter a recursive function binding
    --	f = \x y -> ...
    -- we want to extend the scope env with bindings 
    -- that record that f is a RecFn and x,y are RecArgs
extendRecBndr env fn bndrs
  =  env { scope = scope env `extendVarEnvList` 
		   ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
\end{code}


%************************************************************************
%*									*
\subsection{Usage information: flows upwards}
%*									*
%************************************************************************
277

278
\begin{code}
279
280
data ScUsage
   = SCU {
281
282
283
	calls :: !(IdEnv ([Call])),	-- Calls
					-- The functions are a subset of the 
					-- 	RecFuns in the ScEnv
284
285
286
287
288

	occs :: !(IdEnv ArgOcc)		-- Information on argument occurrences
     }					-- The variables are a subset of the 
					--	RecArg in the ScEnv

289
290
291
292
type Call = (ConstrEnv, [CoreArg])
	-- The arguments of the call, together with the
	-- env giving the constructor bindings at the call site

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }

combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
			   occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }

combineUsages [] = nullUsage
combineUsages us = foldr1 combineUsage us

data ArgOcc = CaseScrut 
	    | OtherOcc
	    | Both

instance Outputable ArgOcc where
  ppr CaseScrut = ptext SLIT("case-scrut")
  ppr OtherOcc  = ptext SLIT("other-occ")
  ppr Both      = ptext SLIT("case-scrut and other")

combineOcc CaseScrut CaseScrut = CaseScrut
combineOcc OtherOcc  OtherOcc  = OtherOcc
combineOcc _	     _	       = Both
\end{code}


%************************************************************************
%*									*
\subsection{The main recursive function}
%*									*
%************************************************************************

322
323
324
The main recursive function gathers up usage information, and
creates specialised versions of functions.

325
326
327
328
329
330
331
332
333
334
335
336
337
\begin{code}
scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
	-- The unique supply is needed when we invent
	-- a new name for the specialised function and its args

scExpr env e@(Type t) = returnUs (nullUsage, e)
scExpr env e@(Lit l)  = returnUs (nullUsage, e)
scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
scExpr env (Note n e) = scExpr env e	`thenUs` \ (usg,e') ->
			returnUs (usg, Note n e')
scExpr env (Lam b e)  = scExpr (extendBndr env b) e	`thenUs` \ (usg,e') ->
			returnUs (usg, Lam b e')

338
scExpr env (Case scrut b ty alts) 
339
340
341
  = sc_scrut scrut		`thenUs` \ (scrut_usg, scrut') ->
    mapAndUnzipUs sc_alt alts	`thenUs` \ (alts_usgs, alts') ->
    returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
342
	      Case scrut' b ty alts')
343
344
345
346
  where
    sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
    sc_scrut e	       = scExpr env e

347
    sc_alt (con,bs,rhs) = scExpr env1 rhs	`thenUs` \ (usg,rhs') ->
348
349
			  returnUs (usg, (con,bs,rhs'))
			where
350
			  env1 = extendCaseBndrs env b scrut con bs
351
352
353
354
355
356
357
358
359
360
361
362
363
364

scExpr env (Let bind body)
  = scBind env bind	`thenUs` \ (env', bind_usg, bind') ->
    scExpr env' body	`thenUs` \ (body_usg, body') ->
    returnUs (bind_usg `combineUsage` body_usg, Let bind' body')

scExpr env e@(App _ _) 
  = let 
	(fn, args) = collectArgs e
    in
    mapAndUnzipUs (scExpr env) args	`thenUs` \ (usgs, args') ->
    let
	arg_usg = combineUsages usgs
	fn_usg  | Var f <- fn,
365
366
367
		  Just RecFun <- lookupScopeEnv env f
		= SCU { calls = unitVarEnv f [(cons env, args)], 
			occs  = emptyVarEnv }
368
369
370
371
372
373
374
375
376
377
		| otherwise
		= nullUsage
    in
    returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
	-- Don't bother to look inside fn;
	-- it's almost always a variable

----------------------
scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
scBind env (Rec [(fn,rhs)])
sof's avatar
sof committed
378
  | notNull val_bndrs
379
  = scExpr env_fn_body body		`thenUs` \ (usg, body') ->
380
381
382
    let
	SCU { calls = calls, occs = occs } = usg
    in
383
    specialise env fn bndrs body usg	`thenUs` \ (rules, spec_prs) ->
384
385
386
    returnUs (extendBndr env fn,	-- For the body of the letrec, just
					-- extend the env with Other to record 
					-- that it's in scope; no funny RecFun business
387
388
389
390
391
	      SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
	      Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
  where
    (bndrs,body) = collectBinders rhs
    val_bndrs    = filter isId bndrs
392
    env_fn_body	 = extendRecBndr env fn bndrs
393
394
395
396
397
398
399
400
401
402
403
404
405
406

scBind env (Rec prs)
  = mapAndUnzipUs do_one prs	`thenUs` \ (usgs, prs') ->
    returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
  where
    do_one (bndr,rhs) = scExpr env rhs	`thenUs` \ (usg, rhs') ->
		        returnUs (usg, (bndr,rhs'))

scBind env (NonRec bndr rhs)
  = scExpr env rhs	`thenUs` \ (usg, rhs') ->
    returnUs (extendBndr env bndr, usg, NonRec bndr rhs')

----------------------
varUsage env v use 
407
408
409
  | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv, 
						occs = unitVarEnv v use }
  | otherwise		   	        = nullUsage
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
\end{code}


%************************************************************************
%*									*
\subsection{The specialiser}
%*									*
%************************************************************************

\begin{code}
specialise :: ScEnv
	   -> Id 			-- Functionn
	   -> [CoreBndr] -> CoreExpr	-- Its RHS
	   -> ScUsage			-- Info on usage
	   -> UniqSM ([CoreRule], 	-- Rules
		      [(Id,CoreExpr)])	-- Bindings

specialise env fn bndrs body (SCU {calls=calls, occs=occs})
  = getUs		`thenUs` \ us ->
    let
	all_calls = lookupVarEnv calls fn `orElse` []

	good_calls :: [[CoreArg]]
	good_calls = [ pats
434
		     | (con_env, call_args) <- all_calls,
sof's avatar
sof committed
435
		       call_args `lengthAtLeast` n_bndrs,	    -- App is saturated
436
		       let call = (bndrs `zip` call_args),
437
438
		       any (good_arg con_env occs) call,    -- At least one arg is a constr app
		       let (_, pats) = argsToPats con_env us call_args
439
440
441
442
443
444
		     ]
    in
    mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
		  (nubBy same_call good_calls `zip` [1..])
  where
    n_bndrs  = length bndrs
445
    same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
446
447

---------------------
448
449
450
451
good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
good_arg con_env arg_occs (bndr, arg)
  = case is_con_app_maybe con_env arg of	
	Just _ ->  bndr_usg_ok arg_occs bndr arg
452
453
454
455
	other   -> False

bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
bndr_usg_ok arg_occs bndr arg
456
  = case lookupVarEnv arg_occs bndr of
457
458
459
460
461
462
463
464
465
466
467
468
469
470
	Just CaseScrut -> True			-- Used only by case scrutiny
	Just Both      -> case arg of		-- Used by case and elsewhere
			    App _ _ -> True	-- so the arg should be an explicit con app
			    other   -> False
	other -> False				-- Not used, or used wonkily
    

---------------------
spec_one :: ScEnv
	 -> Id					-- Function
	 -> CoreExpr				-- Rhs of the original function
	 -> ([CoreArg], Int)
	 -> UniqSM (CoreRule, (Id,CoreExpr))	-- Rule and binding

471
472
473
474
-- spec_one creates a specialised copy of the function, together
-- with a rule for using it.  I'm very proud of how short this
-- function is, considering what it does :-).

475
476
477
478
{- 
  Example
  
     In-scope: a, x::a   
479
480
     f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
	  [c::*, v::(b,c) are presumably bound by the (...) part]
481
  ==>
482
483
     f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
		  (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
484
  
485
486
487
     RULE:  forall b::* c::*,		-- Note, *not* forall a, x
		   v::(b,c),
		   hw::[(a,(b,c))] .
488
  
489
	    f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
490
491
-}

492
493
spec_one env fn rhs (pats, rule_number)
  = getUniqueUs 		`thenUs` \ spec_uniq ->
494
495
496
497
498
499
    let 
	fn_name      = idName fn
	fn_loc       = nameSrcLoc fn_name
	spec_occ     = mkSpecOcc (nameOccName fn_name)
	pat_fvs	     = varSetElems (exprsFreeVars pats)
	vars_to_bind = filter not_avail pat_fvs
500
	not_avail v  = not (v `elemVarEnv` scope env)
501
502
		-- Put the type variables first; the type of a term
		-- variable may mention a type variable
503
504
	(tvs, ids)   = partition isTyVar vars_to_bind
	bndrs  	     = tvs ++ ids
505
506
	spec_body    = mkApps rhs pats
	body_ty	     = exprType spec_body
507
	
508
509
510
511
	(spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
		-- Usual w/w hack to avoid generating 
		-- a spec_rhs of unlifted type and no args
	
512
	rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
513
514
	spec_rhs  = mkLams spec_lam_args spec_body
	spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
515
516
	rule_rhs  = mkVarApps (Var spec_id) spec_call_args
	rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
517
518
    in
    returnUs (rule, (spec_id, spec_rhs))
519
520
521
522
523
524
525
526
527
528
529

-- In which phase should the specialise-constructor rules be active?
-- Originally I made them always-active, but Manuel found that
-- this defeated some clever user-written rules.  So Plan B
-- is to make them active only in Phase 0; after all, currently,
-- the specConstr transformation is only run after the simplifier
-- has reached Phase 0.  In general one would want it to be 
-- flag-controllable, but for now I'm leaving it baked in
--					[SLPJ Oct 01]
specConstrActivation :: Activation
specConstrActivation = ActiveAfter 0	-- Baked in; see comments above
530
\end{code}
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562

%************************************************************************
%*									*
\subsection{Argument analysis}
%*									*
%************************************************************************

This code deals with analysing call-site arguments to see whether
they are constructor applications.

\begin{code}
    -- argToPat takes an actual argument, and returns an abstracted
    -- version, consisting of just the "constructor skeleton" of the
    -- argument, with non-constructor sub-expression replaced by new
    -- placeholder variables.  For example:
    --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)

argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
argToPat env us (Type ty) 
  = (us, Type ty)

argToPat env us arg
  | Just (dc,args) <- is_con_app_maybe env arg
  = let
	(us',args') = argsToPats env us args
    in
    (us', mk_con_app dc args')

argToPat env us (Var v)	-- Don't uniqify existing vars,
  = (us, Var v)		-- so that we can spot when we pass them twice

argToPat env us arg
563
  = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
  where
    (us1,us2) = splitUniqSupply us

argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
argsToPats env us args = mapAccumL (argToPat env) us args
\end{code}


\begin{code}
is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
is_con_app_maybe env (Var v)
  = lookupVarEnv env v
	-- You might think we could look in the idUnfolding here
	-- but that doesn't take account of which branch of a 
	-- case we are in, which is the whole point

is_con_app_maybe env (Lit lit)
  = Just (LitAlt lit, [])

is_con_app_maybe env expr
  = case collectArgs expr of
585
	(Var fun, args) | Just con <- isDataConWorkId_maybe fun,
sof's avatar
sof committed
586
			  args `lengthAtLeast` dataConRepArity con
587
588
589
590
591
592
593
594
595
		-- Might be > because the arity excludes type args
		        -> Just (DataAlt con,args)

	other -> Nothing

mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
mk_con_app (LitAlt lit)  []   = Lit lit
mk_con_app (DataAlt con) args = mkConApp con args
\end{code}