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, tcEqExpr, mkPiTypes )
16 import CoreFVs ( exprsFreeVars )
17 import CoreSubst ( Subst, mkSubst, substExpr )
18 import CoreTidy ( tidyRules )
19 import PprCore ( pprRules )
20 import WwLib ( mkWorkerArgs )
21 import DataCon ( dataConRepArity, isVanillaDataCon )
22 import Type ( tyConAppArgs, tyVarsOfTypes )
23 import Unify ( coreRefineTys )
24 import Id ( Id, idName, idType, isDataConWorkId_maybe,
25 mkUserLocal, mkSysLocal )
29 import Name ( nameOccName, nameSrcLoc )
30 import Rules ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
31 import OccName ( mkSpecOcc )
32 import ErrUtils ( dumpIfSet_dyn )
33 import DynFlags ( DynFlags, DynFlag(..) )
34 import BasicTypes ( Activation(..) )
35 import Maybes ( orElse )
36 import Util ( mapAccumL, lengthAtLeast, notNull )
37 import List ( nubBy, partition )
43 -----------------------------------------------------
45 -----------------------------------------------------
50 drop n (x:xs) = drop (n-1) xs
52 After the first time round, we could pass n unboxed. This happens in
53 numerical code too. Here's what it looks like in Core:
55 drop n xs = case xs of
60 _ -> drop (I# (n# -# 1#)) xs
62 Notice that the recursive call has an explicit constructor as argument.
63 Noticing this, we can make a specialised version of drop
65 RULE: drop (I# n#) xs ==> drop' n# xs
67 drop' n# xs = let n = I# n# in ...orig RHS...
69 Now the simplifier will apply the specialisation in the rhs of drop', giving
71 drop' n# xs = case xs of
75 _ -> drop (n# -# 1#) xs
79 We'd also like to catch cases where a parameter is carried along unchanged,
80 but evaluated each time round the loop:
82 f i n = if i>0 || i>n then i else f (i*2) n
84 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
85 In Core, by the time we've w/wd (f is strict in i) we get
87 f i# n = case i# ># 0 of
89 True -> case n of n' { I# n# ->
92 True -> f (i# *# 2#) n'
94 At the call to f, we see that the argument, n is know to be (I# n#),
95 and n is evaluated elsewhere in the body of f, so we can play the same
96 trick as above. However we don't want to do that if the boxed version
97 of n is needed (else we'd avoid the eval but pay more for re-boxing n).
98 So in this case we want that the *only* uses of n are in case statements.
103 * A self-recursive function. Ignore mutual recursion for now,
104 because it's less common, and the code is simpler for self-recursion.
108 a) At a recursive call, one or more parameters is an explicit
109 constructor application
111 That same parameter is scrutinised by a case somewhere in
112 the RHS of the function
116 b) At a recursive call, one or more parameters has an unfolding
117 that is an explicit constructor application
119 That same parameter is scrutinised by a case somewhere in
120 the RHS of the function
122 Those are the only uses of the parameter
125 What to abstract over
126 ~~~~~~~~~~~~~~~~~~~~~
127 There's a bit of a complication with type arguments. If the call
130 f p = ...f ((:) [a] x xs)...
132 then our specialised function look like
134 f_spec x xs = let p = (:) [a] x xs in ....as before....
136 This only makes sense if either
137 a) the type variable 'a' is in scope at the top of f, or
138 b) the type variable 'a' is an argument to f (and hence fs)
140 Actually, (a) may hold for value arguments too, in which case
141 we may not want to pass them. Supose 'x' is in scope at f's
142 defn, but xs is not. Then we'd like
144 f_spec xs = let p = (:) [a] x xs in ....as before....
146 Similarly (b) may hold too. If x is already an argument at the
147 call, no need to pass it again.
149 Finally, if 'a' is not in scope at the call site, we could abstract
150 it as we do the term variables:
152 f_spec a x xs = let p = (:) [a] x xs in ...as before...
154 So the grand plan is:
156 * abstract the call site to a constructor-only pattern
157 e.g. C x (D (f p) (g q)) ==> C s1 (D s2 s3)
159 * Find the free variables of the abstracted pattern
161 * Pass these variables, less any that are in scope at
162 the fn defn. But see Note [Shadowing] below.
165 NOTICE that we only abstract over variables that are not in scope,
166 so we're in no danger of shadowing variables used in "higher up"
172 In this pass we gather up usage information that may mention variables
173 that are bound between the usage site and the definition site; or (more
174 seriously) may be bound to something different at the definition site.
177 f x = letrec g y v = let x = ...
180 Since 'x' is in scope at the call site, we may make a rewrite rule that
182 RULE forall a,b. g (a,b) x = ...
183 But this rule will never match, because it's really a different 'x' at
184 the call site -- and that difference will be manifest by the time the
185 simplifier gets to it. [A worry: the simplifier doesn't *guarantee*
186 no-shadowing, so perhaps it may not be distinct?]
188 Anyway, the rule isn't actually wrong, it's just not useful. One possibility
189 is to run deShadowBinds before running SpecConstr, but instead we run the
190 simplifier. That gives the simplest possible program for SpecConstr to
191 chew on; and it virtually guarantees no shadowing.
194 %************************************************************************
196 \subsection{Top level wrapper stuff}
198 %************************************************************************
201 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
202 specConstrProgram dflags us binds
204 showPass dflags "SpecConstr"
206 let (binds', _) = initUs us (go emptyScEnv binds)
208 endPass dflags "SpecConstr" Opt_D_dump_spec binds'
210 dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
211 (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
215 go env [] = returnUs []
216 go env (bind:binds) = scBind env bind `thenUs` \ (env', _, bind') ->
217 go env' binds `thenUs` \ binds' ->
218 returnUs (bind' : binds')
222 %************************************************************************
224 \subsection{Environment: goes downwards}
226 %************************************************************************
229 data ScEnv = SCE { scope :: VarEnv HowBound,
230 -- Binds all non-top-level variables in scope
235 type ConstrEnv = IdEnv ConValue
236 data ConValue = CV AltCon [CoreArg]
237 -- Variables known to be bound to a constructor
238 -- in a particular case alternative
241 instance Outputable ConValue where
242 ppr (CV con args) = ppr con <+> interpp'SP args
244 refineConstrEnv :: Subst -> ConstrEnv -> ConstrEnv
245 -- The substitution is a type substitution only
246 refineConstrEnv subst env = mapVarEnv refine_con_value env
248 refine_con_value (CV con args) = CV con (map (substExpr subst) args)
250 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
252 data HowBound = RecFun -- These are the recursive functions for which
253 -- we seek interesting call patterns
255 | RecArg -- These are those functions' arguments; we are
256 -- interested to see if those arguments are scrutinised
258 | Other -- We track all others so we know what's in scope
259 -- This is used in spec_one to check what needs to be
260 -- passed as a parameter and what is in scope at the
261 -- function definition site
263 instance Outputable HowBound where
264 ppr RecFun = text "RecFun"
265 ppr RecArg = text "RecArg"
266 ppr Other = text "Other"
268 lookupScopeEnv env v = lookupVarEnv (scope env) v
270 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
271 extendBndr env bndr = env { scope = extendVarEnv (scope env) bndr Other }
276 -- we want to bind b, and perhaps scrut too, to (C x y)
277 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
278 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
279 = extendBndrs env (case_bndr : alt_bndrs)
281 extendCaseBndrs env case_bndr scrut con@(LitAlt lit) alt_bndrs
282 = ASSERT( null alt_bndrs ) extendAlt env case_bndr scrut (CV con []) []
284 extendCaseBndrs env case_bndr scrut con@(DataAlt data_con) alt_bndrs
285 | isVanillaDataCon data_con
286 = extendAlt env case_bndr scrut (CV con vanilla_args) alt_bndrs
289 = extendAlt env1 case_bndr scrut (CV con gadt_args) alt_bndrs
291 vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
292 map varToCoreExpr alt_bndrs
294 gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
295 -- This call generates some bogus warnings from substExpr,
296 -- because it's inconvenient to put all the Ids in scope
297 -- Will be fixed when we move to FC
299 (alt_tvs, _) = span isTyVar alt_bndrs
300 Just (tv_subst, is_local) = coreRefineTys data_con alt_tvs (idType case_bndr)
301 subst = mkSubst in_scope tv_subst emptyVarEnv -- No Id substitition
302 in_scope = mkInScopeSet (tyVarsOfTypes (varEnvElts tv_subst))
304 env1 | is_local = env
305 | otherwise = env { cons = refineConstrEnv subst (cons env) }
309 extendAlt :: ScEnv -> Id -> CoreExpr -> ConValue -> [Var] -> ScEnv
310 extendAlt env case_bndr scrut val alt_bndrs
312 env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
313 cons = extendVarEnv (cons env) case_bndr val }
316 Var v -> -- Bind the scrutinee in the ConstrEnv if it's a variable
317 -- Also forget if the scrutinee is a RecArg, because we're
318 -- now in the branch of a case, and we don't want to
319 -- record a non-scrutinee use of v if we have
320 -- case v of { (a,b) -> ...(f v)... }
321 SCE { scope = extendVarEnv (scope env1) v Other,
322 cons = extendVarEnv (cons env1) v val }
325 -- When we encounter a recursive function binding
327 -- we want to extend the scope env with bindings
328 -- that record that f is a RecFn and x,y are RecArgs
329 extendRecBndr env fn bndrs
330 = env { scope = scope env `extendVarEnvList`
331 ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
335 %************************************************************************
337 \subsection{Usage information: flows upwards}
339 %************************************************************************
344 calls :: !(IdEnv ([Call])), -- Calls
345 -- The functions are a subset of the
346 -- RecFuns in the ScEnv
348 occs :: !(IdEnv ArgOcc) -- Information on argument occurrences
349 } -- The variables are a subset of the
350 -- RecArg in the ScEnv
352 type Call = (ConstrEnv, [CoreArg])
353 -- The arguments of the call, together with the
354 -- env giving the constructor bindings at the call site
356 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
358 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
359 occs = plusVarEnv_C combineOcc (occs u1) (occs u2) }
361 combineUsages [] = nullUsage
362 combineUsages us = foldr1 combineUsage us
364 data ArgOcc = CaseScrut
368 instance Outputable ArgOcc where
369 ppr CaseScrut = ptext SLIT("case-scrut")
370 ppr OtherOcc = ptext SLIT("other-occ")
371 ppr Both = ptext SLIT("case-scrut and other")
373 combineOcc CaseScrut CaseScrut = CaseScrut
374 combineOcc OtherOcc OtherOcc = OtherOcc
375 combineOcc _ _ = Both
379 %************************************************************************
381 \subsection{The main recursive function}
383 %************************************************************************
385 The main recursive function gathers up usage information, and
386 creates specialised versions of functions.
389 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
390 -- The unique supply is needed when we invent
391 -- a new name for the specialised function and its args
393 scExpr env e@(Type t) = returnUs (nullUsage, e)
394 scExpr env e@(Lit l) = returnUs (nullUsage, e)
395 scExpr env e@(Var v) = returnUs (varUsage env v OtherOcc, e)
396 scExpr env (Note n e) = scExpr env e `thenUs` \ (usg,e') ->
397 returnUs (usg, Note n e')
398 scExpr env (Lam b e) = scExpr (extendBndr env b) e `thenUs` \ (usg,e') ->
399 returnUs (usg, Lam b e')
401 scExpr env (Case scrut b ty alts)
402 = sc_scrut scrut `thenUs` \ (scrut_usg, scrut') ->
403 mapAndUnzipUs sc_alt alts `thenUs` \ (alts_usgs, alts') ->
404 returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
405 Case scrut' b ty alts')
407 sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
408 sc_scrut e = scExpr env e
410 sc_alt (con,bs,rhs) = scExpr env1 rhs `thenUs` \ (usg,rhs') ->
411 returnUs (usg, (con,bs,rhs'))
413 env1 = extendCaseBndrs env b scrut con bs
415 scExpr env (Let bind body)
416 = scBind env bind `thenUs` \ (env', bind_usg, bind') ->
417 scExpr env' body `thenUs` \ (body_usg, body') ->
418 returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
420 scExpr env e@(App _ _)
422 (fn, args) = collectArgs e
424 mapAndUnzipUs (scExpr env) (fn:args) `thenUs` \ (usgs, (fn':args')) ->
425 -- Process the function too. It's almost always a variable,
426 -- but not always. In particular, if this pass follows float-in,
427 -- which it may, we can get
428 -- (let f = ...f... in f) arg1 arg2
430 call_usg = case fn of
431 Var f | Just RecFun <- lookupScopeEnv env f
432 -> SCU { calls = unitVarEnv f [(cons env, args)],
436 returnUs (combineUsages usgs `combineUsage` call_usg, mkApps fn' args')
439 ----------------------
440 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
441 scBind env (Rec [(fn,rhs)])
443 = scExpr env_fn_body body `thenUs` \ (usg, body') ->
444 specialise env fn bndrs body usg `thenUs` \ (rules, spec_prs) ->
446 SCU { calls = calls, occs = occs } = usg
448 returnUs (extendBndr env fn, -- For the body of the letrec, just
449 -- extend the env with Other to record
450 -- that it's in scope; no funny RecFun business
451 SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
452 Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
454 (bndrs,body) = collectBinders rhs
455 val_bndrs = filter isId bndrs
456 env_fn_body = extendRecBndr env fn bndrs
459 = mapAndUnzipUs do_one prs `thenUs` \ (usgs, prs') ->
460 returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
462 do_one (bndr,rhs) = scExpr env rhs `thenUs` \ (usg, rhs') ->
463 returnUs (usg, (bndr,rhs'))
465 scBind env (NonRec bndr rhs)
466 = scExpr env rhs `thenUs` \ (usg, rhs') ->
467 returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
469 ----------------------
471 | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
472 occs = unitVarEnv v use }
473 | otherwise = nullUsage
477 %************************************************************************
479 \subsection{The specialiser}
481 %************************************************************************
486 -> [CoreBndr] -> CoreExpr -- Its RHS
487 -> ScUsage -- Info on usage
488 -> UniqSM ([CoreRule], -- Rules
489 [(Id,CoreExpr)]) -- Bindings
491 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
492 = getUs `thenUs` \ us ->
494 all_calls = lookupVarEnv calls fn `orElse` []
496 good_calls :: [[CoreArg]]
498 | (con_env, call_args) <- all_calls,
499 call_args `lengthAtLeast` n_bndrs, -- App is saturated
500 let call = bndrs `zip` call_args,
501 any (good_arg con_env occs) call, -- At least one arg is a constr app
502 let (_, pats) = argsToPats con_env us call_args
505 mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
506 (nubBy same_call good_calls `zip` [1..])
508 n_bndrs = length bndrs
509 same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
511 ---------------------
512 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
513 good_arg con_env arg_occs (bndr, arg)
514 = case is_con_app_maybe con_env arg of
515 Just _ -> bndr_usg_ok arg_occs bndr arg
518 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
519 bndr_usg_ok arg_occs bndr arg
520 = case lookupVarEnv arg_occs bndr of
521 Just CaseScrut -> True -- Used only by case scrutiny
522 Just Both -> case arg of -- Used by case and elsewhere
523 App _ _ -> True -- so the arg should be an explicit con app
525 other -> False -- Not used, or used wonkily
528 ---------------------
531 -> CoreExpr -- Rhs of the original function
533 -> UniqSM (CoreRule, (Id,CoreExpr)) -- Rule and binding
535 -- spec_one creates a specialised copy of the function, together
536 -- with a rule for using it. I'm very proud of how short this
537 -- function is, considering what it does :-).
543 f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
544 [c::*, v::(b,c) are presumably bound by the (...) part]
546 f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
547 (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
549 RULE: forall b::* c::*, -- Note, *not* forall a, x
553 f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
556 spec_one env fn rhs (pats, rule_number)
557 = getUniqueUs `thenUs` \ spec_uniq ->
560 fn_loc = nameSrcLoc fn_name
561 spec_occ = mkSpecOcc (nameOccName fn_name)
562 pat_fvs = varSetElems (exprsFreeVars pats)
563 vars_to_bind = filter not_avail pat_fvs
564 -- See Note [Shadowing] at the top
566 not_avail v = not (v `elemVarEnv` scope env)
567 -- Put the type variables first; the type of a term
568 -- variable may mention a type variable
569 (tvs, ids) = partition isTyVar vars_to_bind
571 spec_body = mkApps rhs pats
572 body_ty = exprType spec_body
574 (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
575 -- Usual w/w hack to avoid generating
576 -- a spec_rhs of unlifted type and no args
578 rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
579 spec_rhs = mkLams spec_lam_args spec_body
580 spec_id = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
581 rule_rhs = mkVarApps (Var spec_id) spec_call_args
582 rule = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
584 returnUs (rule, (spec_id, spec_rhs))
586 -- In which phase should the specialise-constructor rules be active?
587 -- Originally I made them always-active, but Manuel found that
588 -- this defeated some clever user-written rules. So Plan B
589 -- is to make them active only in Phase 0; after all, currently,
590 -- the specConstr transformation is only run after the simplifier
591 -- has reached Phase 0. In general one would want it to be
592 -- flag-controllable, but for now I'm leaving it baked in
594 specConstrActivation :: Activation
595 specConstrActivation = ActiveAfter 0 -- Baked in; see comments above
598 %************************************************************************
600 \subsection{Argument analysis}
602 %************************************************************************
604 This code deals with analysing call-site arguments to see whether
605 they are constructor applications.
608 -- argToPat takes an actual argument, and returns an abstracted
609 -- version, consisting of just the "constructor skeleton" of the
610 -- argument, with non-constructor sub-expression replaced by new
611 -- placeholder variables. For example:
612 -- C a (D (f x) (g y)) ==> C p1 (D p2 p3)
614 argToPat :: ConstrEnv -> UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
615 argToPat env us (Type ty)
619 | Just (CV dc args) <- is_con_app_maybe env arg
621 (us',args') = argsToPats env us args
623 (us', mk_con_app dc args')
625 argToPat env us (Var v) -- Don't uniqify existing vars,
626 = (us, Var v) -- so that we can spot when we pass them twice
629 = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
631 (us1,us2) = splitUniqSupply us
633 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
634 argsToPats env us args = mapAccumL (argToPat env) us args
639 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
640 is_con_app_maybe env (Var v)
642 -- You might think we could look in the idUnfolding here
643 -- but that doesn't take account of which branch of a
644 -- case we are in, which is the whole point
646 is_con_app_maybe env (Lit lit)
647 = Just (CV (LitAlt lit) [])
649 is_con_app_maybe env expr
650 = case collectArgs expr of
651 (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
652 args `lengthAtLeast` dataConRepArity con
653 -- Might be > because the arity excludes type args
654 -> Just (CV (DataAlt con) args)
658 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
659 mk_con_app (LitAlt lit) [] = Lit lit
660 mk_con_app (DataAlt con) args = mkConApp con args