6a2cd92224d7287fdf22f5088d83a0f0733d0a2d
[ghc-hetmet.git] / ghc / compiler / specialise / SpecConstr.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 \section[SpecConstr]{Specialise over constructors}
5
6 \begin{code}
7 module SpecConstr(
8         specConstrProgram       
9     ) where
10
11 #include "HsVersions.h"
12
13 import CoreSyn
14 import CoreLint         ( showPass, endPass )
15 import CoreUtils        ( exprType, tcEqExpr, mkPiTypes )
16 import CoreFVs          ( exprsFreeVars )
17 import CoreTidy         ( tidyRules )
18 import PprCore          ( pprRules )
19 import WwLib            ( mkWorkerArgs )
20 import DataCon          ( dataConRepArity )
21 import Type             ( tyConAppArgs )
22 import Id               ( Id, idName, idType, isDataConWorkId_maybe, 
23                           mkUserLocal, mkSysLocal )
24 import Var              ( Var )
25 import VarEnv
26 import VarSet
27 import Name             ( nameOccName, nameSrcLoc )
28 import Rules            ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
29 import OccName          ( mkSpecOcc )
30 import ErrUtils         ( dumpIfSet_dyn )
31 import DynFlags         ( DynFlags, DynFlag(..) )
32 import BasicTypes       ( Activation(..) )
33 import Maybes           ( orElse )
34 import Util             ( mapAccumL, lengthAtLeast, notNull )
35 import List             ( nubBy, partition )
36 import UniqSupply
37 import Outputable
38 import FastString
39 \end{code}
40
41 -----------------------------------------------------
42                         Game plan
43 -----------------------------------------------------
44
45 Consider
46         drop n []     = []
47         drop 0 xs     = []
48         drop n (x:xs) = drop (n-1) xs
49
50 After the first time round, we could pass n unboxed.  This happens in
51 numerical code too.  Here's what it looks like in Core:
52
53         drop n xs = case xs of
54                       []     -> []
55                       (y:ys) -> case n of 
56                                   I# n# -> case n# of
57                                              0 -> []
58                                              _ -> drop (I# (n# -# 1#)) xs
59
60 Notice that the recursive call has an explicit constructor as argument.
61 Noticing this, we can make a specialised version of drop
62         
63         RULE: drop (I# n#) xs ==> drop' n# xs
64
65         drop' n# xs = let n = I# n# in ...orig RHS...
66
67 Now the simplifier will apply the specialisation in the rhs of drop', giving
68
69         drop' n# xs = case xs of
70                       []     -> []
71                       (y:ys) -> case n# of
72                                   0 -> []
73                                   _ -> drop (n# -# 1#) xs
74
75 Much better!  
76
77 We'd also like to catch cases where a parameter is carried along unchanged,
78 but evaluated each time round the loop:
79
80         f i n = if i>0 || i>n then i else f (i*2) n
81
82 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
83 In Core, by the time we've w/wd (f is strict in i) we get
84
85         f i# n = case i# ># 0 of
86                    False -> I# i#
87                    True  -> case n of n' { I# n# ->
88                             case i# ># n# of
89                                 False -> I# i#
90                                 True  -> f (i# *# 2#) n'
91
92 At the call to f, we see that the argument, n is know to be (I# n#),
93 and n is evaluated elsewhere in the body of f, so we can play the same
94 trick as above.  However we don't want to do that if the boxed version
95 of n is needed (else we'd avoid the eval but pay more for re-boxing n).
96 So in this case we want that the *only* uses of n are in case statements.
97
98
99 So we look for
100
101 * A self-recursive function.  Ignore mutual recursion for now, 
102   because it's less common, and the code is simpler for self-recursion.
103
104 * EITHER
105
106    a) At a recursive call, one or more parameters is an explicit 
107       constructor application
108         AND
109       That same parameter is scrutinised by a case somewhere in 
110       the RHS of the function
111
112   OR
113
114     b) At a recursive call, one or more parameters has an unfolding
115        that is an explicit constructor application
116         AND
117       That same parameter is scrutinised by a case somewhere in 
118       the RHS of the function
119         AND
120       Those are the only uses of the parameter
121
122
123 There's a bit of a complication with type arguments.  If the call
124 site looks like
125
126         f p = ...f ((:) [a] x xs)...
127
128 then our specialised function look like
129
130         f_spec x xs = let p = (:) [a] x xs in ....as before....
131
132 This only makes sense if either
133   a) the type variable 'a' is in scope at the top of f, or
134   b) the type variable 'a' is an argument to f (and hence fs)
135
136 Actually, (a) may hold for value arguments too, in which case
137 we may not want to pass them.  Supose 'x' is in scope at f's
138 defn, but xs is not.  Then we'd like
139
140         f_spec xs = let p = (:) [a] x xs in ....as before....
141
142 Similarly (b) may hold too.  If x is already an argument at the
143 call, no need to pass it again.
144
145 Finally, if 'a' is not in scope at the call site, we could abstract
146 it as we do the term variables:
147
148         f_spec a x xs = let p = (:) [a] x xs in ...as before...
149
150 So the grand plan is:
151
152         * abstract the call site to a constructor-only pattern
153           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
154
155         * Find the free variables of the abstracted pattern
156
157         * Pass these variables, less any that are in scope at
158           the fn defn.
159
160
161 NOTICE that we only abstract over variables that are not in scope,
162 so we're in no danger of shadowing variables used in "higher up"
163 in f_spec's RHS.
164
165
166 %************************************************************************
167 %*                                                                      *
168 \subsection{Top level wrapper stuff}
169 %*                                                                      *
170 %************************************************************************
171
172 \begin{code}
173 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
174 specConstrProgram dflags us binds
175   = do
176         showPass dflags "SpecConstr"
177
178         let (binds', _) = initUs us (go emptyScEnv binds)
179
180         endPass dflags "SpecConstr" Opt_D_dump_spec binds'
181
182         dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
183                   (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
184
185         return binds'
186   where
187     go env []           = returnUs []
188     go env (bind:binds) = scBind env bind       `thenUs` \ (env', _, bind') ->
189                           go env' binds         `thenUs` \ binds' ->
190                           returnUs (bind' : binds')
191 \end{code}
192
193
194 %************************************************************************
195 %*                                                                      *
196 \subsection{Environment: goes downwards}
197 %*                                                                      *
198 %************************************************************************
199
200 \begin{code}
201 data ScEnv = SCE { scope :: VarEnv HowBound,
202                         -- Binds all non-top-level variables in scope
203
204                    cons  :: ConstrEnv
205              }
206
207 type ConstrEnv = IdEnv (AltCon, [CoreArg])
208         -- Variables known to be bound to a constructor
209         -- in a particular case alternative
210
211 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
212
213 data HowBound = RecFun          -- These are the recursive functions for which 
214                                 -- we seek interesting call patterns
215
216               | RecArg          -- These are those functions' arguments; we are
217                                 -- interested to see if those arguments are scrutinised
218
219               | Other           -- We track all others so we know what's in scope
220                                 -- This is used in spec_one to check what needs to be
221                                 -- passed as a parameter and what is in scope at the 
222                                 -- function definition site
223
224 instance Outputable HowBound where
225   ppr RecFun = text "RecFun"
226   ppr RecArg = text "RecArg"
227   ppr Other = text "Other"
228
229 lookupScopeEnv env v = lookupVarEnv (scope env) v
230
231 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
232 extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }
233
234     -- When we encounter
235     --  case scrut of b
236     --      C x y -> ...
237     -- we want to bind b, and perhaps scrut too, to (C x y)
238 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
239 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
240   = extendBndrs env (case_bndr : alt_bndrs)
241
242 extendCaseBndrs env case_bndr scrut con alt_bndrs
243   = case scrut of
244         Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
245                    -- Also forget if the scrutinee is a RecArg, because we're
246                    -- now in the branch of a case, and we don't want to
247                    -- record a non-scrutinee use of v if we have
248                    --   case v of { (a,b) -> ...(f v)... }
249                  SCE { scope = extendVarEnv (scope env1) v Other,
250                        cons  = extendVarEnv (cons env1)  v (con,args) }
251         other -> env1
252
253   where
254     env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
255                  cons  = extendVarEnv     (cons  env) case_bndr (con,args) }
256
257     args = map Type (tyConAppArgs (idType case_bndr)) ++
258            map varToCoreExpr alt_bndrs
259
260     -- When we encounter a recursive function binding
261     --  f = \x y -> ...
262     -- we want to extend the scope env with bindings 
263     -- that record that f is a RecFn and x,y are RecArgs
264 extendRecBndr env fn bndrs
265   =  env { scope = scope env `extendVarEnvList` 
266                    ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
267 \end{code}
268
269
270 %************************************************************************
271 %*                                                                      *
272 \subsection{Usage information: flows upwards}
273 %*                                                                      *
274 %************************************************************************
275
276 \begin{code}
277 data ScUsage
278    = SCU {
279         calls :: !(IdEnv ([Call])),     -- Calls
280                                         -- The functions are a subset of the 
281                                         --      RecFuns in the ScEnv
282
283         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
284      }                                  -- The variables are a subset of the 
285                                         --      RecArg in the ScEnv
286
287 type Call = (ConstrEnv, [CoreArg])
288         -- The arguments of the call, together with the
289         -- env giving the constructor bindings at the call site
290
291 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
292
293 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
294                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
295
296 combineUsages [] = nullUsage
297 combineUsages us = foldr1 combineUsage us
298
299 data ArgOcc = CaseScrut 
300             | OtherOcc
301             | Both
302
303 instance Outputable ArgOcc where
304   ppr CaseScrut = ptext SLIT("case-scrut")
305   ppr OtherOcc  = ptext SLIT("other-occ")
306   ppr Both      = ptext SLIT("case-scrut and other")
307
308 combineOcc CaseScrut CaseScrut = CaseScrut
309 combineOcc OtherOcc  OtherOcc  = OtherOcc
310 combineOcc _         _         = Both
311 \end{code}
312
313
314 %************************************************************************
315 %*                                                                      *
316 \subsection{The main recursive function}
317 %*                                                                      *
318 %************************************************************************
319
320 The main recursive function gathers up usage information, and
321 creates specialised versions of functions.
322
323 \begin{code}
324 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
325         -- The unique supply is needed when we invent
326         -- a new name for the specialised function and its args
327
328 scExpr env e@(Type t) = returnUs (nullUsage, e)
329 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
330 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
331 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
332                         returnUs (usg, Note n e')
333 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
334                         returnUs (usg, Lam b e')
335
336 scExpr env (Case scrut b ty alts) 
337   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
338     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
339     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
340               Case scrut' b ty alts')
341   where
342     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
343     sc_scrut e         = scExpr env e
344
345     sc_alt (con,bs,rhs) = scExpr env1 rhs       `thenUs` \ (usg,rhs') ->
346                           returnUs (usg, (con,bs,rhs'))
347                         where
348                           env1 = extendCaseBndrs env b scrut con bs
349
350 scExpr env (Let bind body)
351   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
352     scExpr env' body    `thenUs` \ (body_usg, body') ->
353     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
354
355 scExpr env e@(App _ _) 
356   = let 
357         (fn, args) = collectArgs e
358     in
359     mapAndUnzipUs (scExpr env) args     `thenUs` \ (usgs, args') ->
360     let
361         arg_usg = combineUsages usgs
362         fn_usg  | Var f <- fn,
363                   Just RecFun <- lookupScopeEnv env f
364                 = SCU { calls = unitVarEnv f [(cons env, args)], 
365                         occs  = emptyVarEnv }
366                 | otherwise
367                 = nullUsage
368     in
369     returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
370         -- Don't bother to look inside fn;
371         -- it's almost always a variable
372
373 ----------------------
374 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
375 scBind env (Rec [(fn,rhs)])
376   | notNull val_bndrs
377   = scExpr env_fn_body body             `thenUs` \ (usg, body') ->
378     let
379         SCU { calls = calls, occs = occs } = usg
380     in
381     specialise env fn bndrs body usg    `thenUs` \ (rules, spec_prs) ->
382     returnUs (extendBndr env fn,        -- For the body of the letrec, just
383                                         -- extend the env with Other to record 
384                                         -- that it's in scope; no funny RecFun business
385               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
386               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
387   where
388     (bndrs,body) = collectBinders rhs
389     val_bndrs    = filter isId bndrs
390     env_fn_body  = extendRecBndr env fn bndrs
391
392 scBind env (Rec prs)
393   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
394     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
395   where
396     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
397                         returnUs (usg, (bndr,rhs'))
398
399 scBind env (NonRec bndr rhs)
400   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
401     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
402
403 ----------------------
404 varUsage env v use 
405   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv, 
406                                                 occs = unitVarEnv v use }
407   | otherwise                           = nullUsage
408 \end{code}
409
410
411 %************************************************************************
412 %*                                                                      *
413 \subsection{The specialiser}
414 %*                                                                      *
415 %************************************************************************
416
417 \begin{code}
418 specialise :: ScEnv
419            -> Id                        -- Functionn
420            -> [CoreBndr] -> CoreExpr    -- Its RHS
421            -> ScUsage                   -- Info on usage
422            -> UniqSM ([CoreRule],       -- Rules
423                       [(Id,CoreExpr)])  -- Bindings
424
425 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
426   = getUs               `thenUs` \ us ->
427     let
428         all_calls = lookupVarEnv calls fn `orElse` []
429
430         good_calls :: [[CoreArg]]
431         good_calls = [ pats
432                      | (con_env, call_args) <- all_calls,
433                        call_args `lengthAtLeast` n_bndrs,           -- App is saturated
434                        let call = (bndrs `zip` call_args),
435                        any (good_arg con_env occs) call,    -- At least one arg is a constr app
436                        let (_, pats) = argsToPats con_env us call_args
437                      ]
438     in
439     mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
440                   (nubBy same_call good_calls `zip` [1..])
441   where
442     n_bndrs  = length bndrs
443     same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
444
445 ---------------------
446 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
447 good_arg con_env arg_occs (bndr, arg)
448   = case is_con_app_maybe con_env arg of        
449         Just _ ->  bndr_usg_ok arg_occs bndr arg
450         other   -> False
451
452 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
453 bndr_usg_ok arg_occs bndr arg
454   = case lookupVarEnv arg_occs bndr of
455         Just CaseScrut -> True                  -- Used only by case scrutiny
456         Just Both      -> case arg of           -- Used by case and elsewhere
457                             App _ _ -> True     -- so the arg should be an explicit con app
458                             other   -> False
459         other -> False                          -- Not used, or used wonkily
460     
461
462 ---------------------
463 spec_one :: ScEnv
464          -> Id                                  -- Function
465          -> CoreExpr                            -- Rhs of the original function
466          -> ([CoreArg], Int)
467          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
468
469 -- spec_one creates a specialised copy of the function, together
470 -- with a rule for using it.  I'm very proud of how short this
471 -- function is, considering what it does :-).
472
473 {- 
474   Example
475   
476      In-scope: a, x::a   
477      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
478           [c::*, v::(b,c) are presumably bound by the (...) part]
479   ==>
480      f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
481                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
482   
483      RULE:  forall b::* c::*,           -- Note, *not* forall a, x
484                    v::(b,c),
485                    hw::[(a,(b,c))] .
486   
487             f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
488 -}
489
490 spec_one env fn rhs (pats, rule_number)
491   = getUniqueUs                 `thenUs` \ spec_uniq ->
492     let 
493         fn_name      = idName fn
494         fn_loc       = nameSrcLoc fn_name
495         spec_occ     = mkSpecOcc (nameOccName fn_name)
496         pat_fvs      = varSetElems (exprsFreeVars pats)
497         vars_to_bind = filter not_avail pat_fvs
498         not_avail v  = not (v `elemVarEnv` scope env)
499                 -- Put the type variables first; the type of a term
500                 -- variable may mention a type variable
501         (tvs, ids)   = partition isTyVar vars_to_bind
502         bndrs        = tvs ++ ids
503         spec_body    = mkApps rhs pats
504         body_ty      = exprType spec_body
505         
506         (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
507                 -- Usual w/w hack to avoid generating 
508                 -- a spec_rhs of unlifted type and no args
509         
510         rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
511         spec_rhs  = mkLams spec_lam_args spec_body
512         spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
513         rule_rhs  = mkVarApps (Var spec_id) spec_call_args
514         rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
515     in
516     returnUs (rule, (spec_id, spec_rhs))
517
518 -- In which phase should the specialise-constructor rules be active?
519 -- Originally I made them always-active, but Manuel found that
520 -- this defeated some clever user-written rules.  So Plan B
521 -- is to make them active only in Phase 0; after all, currently,
522 -- the specConstr transformation is only run after the simplifier
523 -- has reached Phase 0.  In general one would want it to be 
524 -- flag-controllable, but for now I'm leaving it baked in
525 --                                      [SLPJ Oct 01]
526 specConstrActivation :: Activation
527 specConstrActivation = ActiveAfter 0    -- Baked in; see comments above
528 \end{code}
529
530 %************************************************************************
531 %*                                                                      *
532 \subsection{Argument analysis}
533 %*                                                                      *
534 %************************************************************************
535
536 This code deals with analysing call-site arguments to see whether
537 they are constructor applications.
538
539 \begin{code}
540     -- argToPat takes an actual argument, and returns an abstracted
541     -- version, consisting of just the "constructor skeleton" of the
542     -- argument, with non-constructor sub-expression replaced by new
543     -- placeholder variables.  For example:
544     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
545
546 argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
547 argToPat env us (Type ty) 
548   = (us, Type ty)
549
550 argToPat env us arg
551   | Just (dc,args) <- is_con_app_maybe env arg
552   = let
553         (us',args') = argsToPats env us args
554     in
555     (us', mk_con_app dc args')
556
557 argToPat env us (Var v) -- Don't uniqify existing vars,
558   = (us, Var v)         -- so that we can spot when we pass them twice
559
560 argToPat env us arg
561   = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
562   where
563     (us1,us2) = splitUniqSupply us
564
565 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
566 argsToPats env us args = mapAccumL (argToPat env) us args
567 \end{code}
568
569
570 \begin{code}
571 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
572 is_con_app_maybe env (Var v)
573   = lookupVarEnv env v
574         -- You might think we could look in the idUnfolding here
575         -- but that doesn't take account of which branch of a 
576         -- case we are in, which is the whole point
577
578 is_con_app_maybe env (Lit lit)
579   = Just (LitAlt lit, [])
580
581 is_con_app_maybe env expr
582   = case collectArgs expr of
583         (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
584                           args `lengthAtLeast` dataConRepArity con
585                 -- Might be > because the arity excludes type args
586                         -> Just (DataAlt con,args)
587
588         other -> Nothing
589
590 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
591 mk_con_app (LitAlt lit)  []   = Lit lit
592 mk_con_app (DataAlt con) args = mkConApp con args
593 \end{code}