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