2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[SpecConstr]{Specialise over constructors}
11 #include "HsVersions.h"
14 import CoreLint ( showPass, endPass )
15 import CoreUtils ( exprType, eqExpr )
16 import CoreFVs ( exprsFreeVars )
17 import DataCon ( dataConRepArity )
18 import Type ( tyConAppArgs )
19 import PprCore ( pprCoreRules, pprCoreRule )
20 import Id ( Id, idName, idType, idSpecialisation,
22 mkUserLocal, mkSysLocal )
26 import Name ( nameOccName, nameSrcLoc )
27 import Rules ( addIdSpecialisations )
28 import OccName ( mkSpecOcc )
29 import ErrUtils ( dumpIfSet_dyn )
30 import CmdLineOpts ( DynFlags, DynFlag(..) )
33 import Maybes ( orElse )
34 import Util ( mapAccumL )
35 import List ( nubBy, partition )
40 -----------------------------------------------------
42 -----------------------------------------------------
47 drop n (x:xs) = drop (n-1) xs
49 After the first time round, we could pass n unboxed. This happens in
50 numerical code too. Here's what it looks like in Core:
52 drop n xs = case xs of
57 _ -> drop (I# (n# -# 1#)) xs
59 Notice that the recursive call has an explicit constructor as argument.
60 Noticing this, we can make a specialised version of drop
62 RULE: drop (I# n#) xs ==> drop' n# xs
64 drop' n# xs = let n = I# n# in ...orig RHS...
66 Now the simplifier will apply the specialisation in the rhs of drop', giving
68 drop' n# xs = case xs of
72 _ -> drop (n# -# 1#) xs
76 We'd also like to catch cases where a parameter is carried along unchanged,
77 but evaluated each time round the loop:
79 f i n = if i>0 || i>n then i else f (i*2) n
81 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
82 In Core, by the time we've w/wd (f is strict in i) we get
84 f i# n = case i# ># 0 of
86 True -> case n of n' { I# n# ->
89 True -> f (i# *# 2#) n'
91 At the call to f, we see that the argument, n is know to be (I# n#),
92 and n is evaluated elsewhere in the body of f, so we can play the same
93 trick as above. However we don't want to do that if the boxed version
94 of n is needed (else we'd avoid the eval but pay more for re-boxing n).
95 So in this case we want that the *only* uses of n are in case statements.
100 * A self-recursive function. Ignore mutual recursion for now,
101 because it's less common, and the code is simpler for self-recursion.
105 a) At a recursive call, one or more parameters is an explicit
106 constructor application
108 That same parameter is scrutinised by a case somewhere in
109 the RHS of the function
113 b) At a recursive call, one or more parameters has an unfolding
114 that is an explicit constructor application
116 That same parameter is scrutinised by a case somewhere in
117 the RHS of the function
119 Those are the only uses of the parameter
122 There's a bit of a complication with type arguments. If the call
125 f p = ...f ((:) [a] x xs)...
127 then our specialised function look like
129 f_spec x xs = let p = (:) [a] x xs in ....as before....
131 This only makes sense if either
132 a) the type variable 'a' is in scope at the top of f, or
133 b) the type variable 'a' is an argument to f (and hence fs)
135 Actually, (a) may hold for value arguments too, in which case
136 we may not want to pass them. Supose 'x' is in scope at f's
137 defn, but xs is not. Then we'd like
139 f_spec xs = let p = (:) [a] x xs in ....as before....
141 Similarly (b) may hold too. If x is already an argument at the
142 call, no need to pass it again.
144 Finally, if 'a' is not in scope at the call site, we could abstract
145 it as we do the term variables:
147 f_spec a x xs = let p = (:) [a] x xs in ...as before...
149 So the grand plan is:
151 * abstract the call site to a constructor-only pattern
152 e.g. C x (D (f p) (g q)) ==> C s1 (D s2 s3)
154 * Find the free variables of the abstracted pattern
156 * Pass these variables, less any that are in scope at
160 NOTICE that we only abstract over variables that are not in scope,
161 so we're in no danger of shadowing variables used in "higher up"
165 %************************************************************************
167 \subsection{Top level wrapper stuff}
169 %************************************************************************
172 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
173 specConstrProgram dflags us binds
175 showPass dflags "SpecConstr"
177 let (binds', _) = initUs us (go emptyScEnv binds)
179 endPass dflags "SpecConstr" Opt_D_dump_spec binds'
181 dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
182 (vcat (map dump_specs (concat (map bindersOf binds'))))
186 go env [] = returnUs []
187 go env (bind:binds) = scBind env bind `thenUs` \ (env', _, bind') ->
188 go env' binds `thenUs` \ binds' ->
189 returnUs (bind' : binds')
191 dump_specs var = pprCoreRules var (idSpecialisation var)
195 %************************************************************************
197 \subsection{Environment: goes downwards}
199 %************************************************************************
202 data ScEnv = SCE { scope :: VarEnv HowBound,
203 -- Binds all non-top-level variables in scope
208 type ConstrEnv = IdEnv (AltCon, [CoreArg])
209 -- Variables known to be bound to a constructor
210 -- in a particular case alternative
212 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
214 data HowBound = RecFun -- These are the recursive functions for which
215 -- we seek interesting call patterns
217 | RecArg -- These are those functions' arguments; we are
218 -- interested to see if those arguments are scrutinised
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
225 lookupScopeEnv env v = lookupVarEnv (scope env) v
227 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
228 extendBndr env bndr = env { scope = extendVarEnv (scope env) bndr Other }
233 -- we want to bind b, and perhaps scrut too, to (C x y)
234 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
235 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
236 = extendBndrs env (case_bndr : alt_bndrs)
238 extendCaseBndrs env case_bndr scrut con alt_bndrs
240 Var v -> -- Bind the scrutinee in the ConstrEnv if it's a variable
241 -- Also forget if the scrutinee is a RecArg, because we're
242 -- now in the branch of a case, and we don't want to
243 -- record a non-scrutinee use of v if we have
244 -- case v of { (a,b) -> ...(f v)... }
245 SCE { scope = extendVarEnv (scope env1) v Other,
246 cons = extendVarEnv (cons env1) v (con,args) }
250 env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
251 cons = extendVarEnv (cons env) case_bndr (con,args) }
253 args = map Type (tyConAppArgs (idType case_bndr)) ++
254 map varToCoreExpr alt_bndrs
256 -- When we encounter a recursive function binding
258 -- we want to extend the scope env with bindings
259 -- that record that f is a RecFn and x,y are RecArgs
260 extendRecBndr env fn bndrs
261 = env { scope = scope env `extendVarEnvList`
262 ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
266 %************************************************************************
268 \subsection{Usage information: flows upwards}
270 %************************************************************************
275 calls :: !(IdEnv ([Call])), -- Calls
276 -- The functions are a subset of the
277 -- RecFuns in the ScEnv
279 occs :: !(IdEnv ArgOcc) -- Information on argument occurrences
280 } -- The variables are a subset of the
281 -- RecArg in the ScEnv
283 type Call = (ConstrEnv, [CoreArg])
284 -- The arguments of the call, together with the
285 -- env giving the constructor bindings at the call site
287 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
289 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
290 occs = plusVarEnv_C combineOcc (occs u1) (occs u2) }
292 combineUsages [] = nullUsage
293 combineUsages us = foldr1 combineUsage us
295 data ArgOcc = CaseScrut
299 instance Outputable ArgOcc where
300 ppr CaseScrut = ptext SLIT("case-scrut")
301 ppr OtherOcc = ptext SLIT("other-occ")
302 ppr Both = ptext SLIT("case-scrut and other")
304 combineOcc CaseScrut CaseScrut = CaseScrut
305 combineOcc OtherOcc OtherOcc = OtherOcc
306 combineOcc _ _ = Both
310 %************************************************************************
312 \subsection{The main recursive function}
314 %************************************************************************
316 The main recursive function gathers up usage information, and
317 creates specialised versions of functions.
320 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
321 -- The unique supply is needed when we invent
322 -- a new name for the specialised function and its args
324 scExpr env e@(Type t) = returnUs (nullUsage, e)
325 scExpr env e@(Lit l) = returnUs (nullUsage, e)
326 scExpr env e@(Var v) = returnUs (varUsage env v OtherOcc, e)
327 scExpr env (Note n e) = scExpr env e `thenUs` \ (usg,e') ->
328 returnUs (usg, Note n e')
329 scExpr env (Lam b e) = scExpr (extendBndr env b) e `thenUs` \ (usg,e') ->
330 returnUs (usg, Lam b e')
332 scExpr env (Case scrut b alts)
333 = sc_scrut scrut `thenUs` \ (scrut_usg, scrut') ->
334 mapAndUnzipUs sc_alt alts `thenUs` \ (alts_usgs, alts') ->
335 returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
338 sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
339 sc_scrut e = scExpr env e
341 sc_alt (con,bs,rhs) = scExpr env1 rhs `thenUs` \ (usg,rhs') ->
342 returnUs (usg, (con,bs,rhs'))
344 env1 = extendCaseBndrs env b scrut con bs
346 scExpr env (Let bind body)
347 = scBind env bind `thenUs` \ (env', bind_usg, bind') ->
348 scExpr env' body `thenUs` \ (body_usg, body') ->
349 returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
351 scExpr env e@(App _ _)
353 (fn, args) = collectArgs e
355 mapAndUnzipUs (scExpr env) args `thenUs` \ (usgs, args') ->
357 arg_usg = combineUsages usgs
358 fn_usg | Var f <- fn,
359 Just RecFun <- lookupScopeEnv env f
360 = SCU { calls = unitVarEnv f [(cons env, args)],
365 returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
366 -- Don't bother to look inside fn;
367 -- it's almost always a variable
369 ----------------------
370 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
371 scBind env (Rec [(fn,rhs)])
372 | not (null val_bndrs)
373 = scExpr env' body `thenUs` \ (usg, body') ->
375 SCU { calls = calls, occs = occs } = usg
377 specialise env fn bndrs body usg `thenUs` \ (rules, spec_prs) ->
378 returnUs (extendBndrs env bndrs,
379 SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
380 Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
382 (bndrs,body) = collectBinders rhs
383 val_bndrs = filter isId bndrs
384 env' = extendRecBndr env fn bndrs
387 = mapAndUnzipUs do_one prs `thenUs` \ (usgs, prs') ->
388 returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
390 do_one (bndr,rhs) = scExpr env rhs `thenUs` \ (usg, rhs') ->
391 returnUs (usg, (bndr,rhs'))
393 scBind env (NonRec bndr rhs)
394 = scExpr env rhs `thenUs` \ (usg, rhs') ->
395 returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
397 ----------------------
399 | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
400 occs = unitVarEnv v use }
401 | otherwise = nullUsage
405 %************************************************************************
407 \subsection{The specialiser}
409 %************************************************************************
414 -> [CoreBndr] -> CoreExpr -- Its RHS
415 -> ScUsage -- Info on usage
416 -> UniqSM ([CoreRule], -- Rules
417 [(Id,CoreExpr)]) -- Bindings
419 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
420 = getUs `thenUs` \ us ->
422 all_calls = lookupVarEnv calls fn `orElse` []
424 good_calls :: [[CoreArg]]
426 | (con_env, call_args) <- all_calls,
427 length call_args >= n_bndrs, -- App is saturated
428 let call = (bndrs `zip` call_args),
429 any (good_arg con_env occs) call, -- At least one arg is a constr app
430 let (_, pats) = argsToPats con_env us call_args
433 mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
434 (nubBy same_call good_calls `zip` [1..])
436 n_bndrs = length bndrs
437 same_call as1 as2 = and (zipWith eqExpr as1 as2)
439 ---------------------
440 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
441 good_arg con_env arg_occs (bndr, arg)
442 = case is_con_app_maybe con_env arg of
443 Just _ -> bndr_usg_ok arg_occs bndr arg
446 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
447 bndr_usg_ok arg_occs bndr arg
448 = case lookupVarEnv arg_occs bndr of
449 Just CaseScrut -> True -- Used only by case scrutiny
450 Just Both -> case arg of -- Used by case and elsewhere
451 App _ _ -> True -- so the arg should be an explicit con app
453 other -> False -- Not used, or used wonkily
456 ---------------------
459 -> CoreExpr -- Rhs of the original function
461 -> UniqSM (CoreRule, (Id,CoreExpr)) -- Rule and binding
463 -- spec_one creates a specialised copy of the function, together
464 -- with a rule for using it. I'm very proud of how short this
465 -- function is, considering what it does :-).
471 f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
472 [c is presumably bound by the (...) part]
474 f_spec = /\ b c \ v::(a,(b,c)) ->
475 (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
480 h::(a,(b,c))->[(a,(b,c))] .
482 f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
485 spec_one env fn rhs (pats, n)
486 = getUniqueUs `thenUs` \ spec_uniq ->
489 fn_loc = nameSrcLoc fn_name
490 spec_occ = mkSpecOcc (nameOccName fn_name)
491 pat_fvs = varSetElems (exprsFreeVars pats)
492 vars_to_bind = filter not_avail pat_fvs
493 not_avail v = not (v `elemVarEnv` scope env)
494 -- Put the type variables first just for tidiness
495 (tvs, ids) = partition isTyVar vars_to_bind
498 rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
499 spec_rhs = mkLams bndrs (mkApps rhs pats)
500 spec_id = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
501 rule = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
503 pprTrace "SpecConstr" (pprCoreRule (ppr fn) rule) $
504 returnUs (rule, (spec_id, spec_rhs))
507 %************************************************************************
509 \subsection{Argument analysis}
511 %************************************************************************
513 This code deals with analysing call-site arguments to see whether
514 they are constructor applications.
517 -- argToPat takes an actual argument, and returns an abstracted
518 -- version, consisting of just the "constructor skeleton" of the
519 -- argument, with non-constructor sub-expression replaced by new
520 -- placeholder variables. For example:
521 -- C a (D (f x) (g y)) ==> C p1 (D p2 p3)
523 argToPat :: ConstrEnv -> UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
524 argToPat env us (Type ty)
528 | Just (dc,args) <- is_con_app_maybe env arg
530 (us',args') = argsToPats env us args
532 (us', mk_con_app dc args')
534 argToPat env us (Var v) -- Don't uniqify existing vars,
535 = (us, Var v) -- so that we can spot when we pass them twice
538 = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
540 (us1,us2) = splitUniqSupply us
542 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
543 argsToPats env us args = mapAccumL (argToPat env) us args
548 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
549 is_con_app_maybe env (Var v)
551 -- You might think we could look in the idUnfolding here
552 -- but that doesn't take account of which branch of a
553 -- case we are in, which is the whole point
555 is_con_app_maybe env (Lit lit)
556 = Just (LitAlt lit, [])
558 is_con_app_maybe env expr
559 = case collectArgs expr of
560 (Var fun, args) | Just con <- isDataConId_maybe fun,
561 length args >= dataConRepArity con
562 -- Might be > because the arity excludes type args
563 -> Just (DataAlt con,args)
567 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
568 mk_con_app (LitAlt lit) [] = Lit lit
569 mk_con_app (DataAlt con) args = mkConApp con args