2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[CoreRules]{Transformation rules}
8 RuleBase, prepareLocalRuleBase, prepareOrphanRuleBase,
9 unionRuleBase, lookupRule, addRule, addIdSpecialisations,
10 ProtoCoreRule(..), pprProtoCoreRule, pprRuleBase,
14 #include "HsVersions.h"
16 import CoreSyn -- All of it
17 import OccurAnal ( occurAnalyseRule )
18 import CoreFVs ( exprFreeVars, idRuleVars, ruleRhsFreeVars, ruleSomeLhsFreeVars )
19 import CoreUnfold ( isCheapUnfolding, unfoldingTemplate )
20 import CoreUtils ( eqExpr )
21 import PprCore ( pprCoreRule )
22 import Subst ( Subst, InScopeSet, mkInScopeSet, lookupSubst, extendSubst,
23 substEnv, setSubstEnv, emptySubst, isInScope,
24 bindSubstList, unBindSubstList, substInScope, uniqAway
26 import Id ( Id, idUnfolding, zapLamIdInfo,
27 idSpecialisation, setIdSpecialisation,
30 import Name ( isLocallyDefined )
31 import Var ( isTyVar, isId )
34 import Type ( mkTyVarTy )
35 import qualified Unify ( match )
39 import Maybes ( maybeToBool )
40 import Util ( sortLt )
44 %************************************************************************
46 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
48 %************************************************************************
50 A @CoreRule@ holds details of one rule for an @Id@, which
51 includes its specialisations.
53 For example, if a rule for @f@ contains the mapping:
55 forall a b d. [Type (List a), Type b, Var d] ===> f' a b
57 then when we find an application of f to matching types, we simply replace
58 it by the matching RHS:
60 f (List Int) Bool dict ===> f' Int Bool
62 All the stuff about how many dictionaries to discard, and what types
63 to apply the specialised function to, are handled by the fact that the
64 Rule contains a template for the result of the specialisation.
66 There is one more exciting case, which is dealt with in exactly the same
67 way. If the specialised value is unboxed then it is lifted at its
68 definition site and unlifted at its uses. For example:
70 pi :: forall a. Num a => a
72 might have a specialisation
74 [Int#] ===> (case pi' of Lift pi# -> pi#)
76 where pi' :: Lift Int# is the specialised version of pi.
79 %************************************************************************
83 %************************************************************************
86 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
87 -- See comments on matchRule
88 matchRules in_scope [] args = Nothing
89 matchRules in_scope (rule:rules) args
90 = case matchRule in_scope rule args of
91 Just result -> Just result
92 Nothing -> matchRules in_scope rules args
95 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
97 -- If (matchRule rule args) returns Just (name,rhs)
98 -- then (f args) matches the rule, and the corresponding
99 -- rewritten RHS is rhs
101 -- The bndrs and rhs is occurrence-analysed
106 -- forall f g x. map f (map g x) ==> map (f . g) x
108 -- CoreRule "map/map"
109 -- [f,g,x] -- tpl_vars
110 -- [f,map g x] -- tpl_args
111 -- map (f.g) x) -- rhs
113 -- Then the call: matchRule the_rule [e1,map e2 e3]
114 -- = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
116 -- Any 'surplus' arguments in the input are simply put on the end
120 -- A1. No top-level variable is bound in the target
121 -- A2. No template variable is bound in the target
122 -- A3. No lambda bound template variable is free in any subexpression of the target
124 -- To see why A1 is necessary, consider matching
125 -- \x->f against \f->f
126 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
127 -- and then erroneously succeed in matching f against f.
129 -- To see why A2 is needed consider matching
130 -- forall a. \b->b against \a->3
131 -- When we meet the lambdas we substitute [a/b] in the template, and then
132 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
134 -- A3 is needed to validate the rule that says
137 -- (\x->E) matches (\x->F x)
140 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
142 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
143 = go tpl_args args emptySubst
144 -- We used to use the in_scope set, but I don't think that's necessary
145 -- After all, the result is going to be simplified again with that in_scope set
147 tpl_var_set = mkVarSet tpl_vars
149 -----------------------
151 go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
153 -- Two easy ways to terminate
154 go [] [] subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
155 go [] args subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
157 -- One tiresome way to terminate: check for excess unmatched
158 -- template arguments
159 go tpl_args [] subst = Nothing -- Failure
162 -----------------------
163 app_match subst fn vs = foldl go fn vs
165 senv = substEnv subst
166 go fn v = case lookupSubstEnv senv v of
167 Just (DoneEx ex) -> fn `App` ex
168 Just (DoneTy ty) -> fn `App` Type ty
169 -- Substitution should bind them all!
172 -----------------------
173 {- The code below tries to match even if there are more
174 template args than real args.
176 I now think this is probably a bad idea.
177 Should the template (map f xs) match (map g)? I think not.
178 For a start, in general eta expansion wastes work.
181 = case eta_complete tpl_args (mkVarSet leftovers) of
182 Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
183 mk_result_args subst done)
184 Nothing -> Nothing -- Failure
186 (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
187 (map zapOccInfo tpl_vars)
189 subst_env = substEnv subst
191 -----------------------
192 eta_complete [] vars = ASSERT( isEmptyVarSet vars )
194 eta_complete (Type ty:tpl_args) vars
195 = case getTyVar_maybe ty of
196 Just tv | tv `elemVarSet` vars
197 -> case eta_complete tpl_args (vars `delVarSet` tv) of
198 Just vars' -> Just (tv:vars')
202 eta_complete (Var v:tpl_args) vars
203 | v `elemVarSet` vars
204 = case eta_complete tpl_args (vars `delVarSet` v) of
205 Just vars' -> Just (v:vars')
208 eta_complete other vars = Nothing
212 zapOccInfo bndr | isTyVar bndr = bndr
213 | otherwise = zapLamIdInfo bndr
217 type Matcher result = VarSet -- Template variables
218 -> (Subst -> Maybe result) -- Continuation if success
219 -> Subst -> Maybe result -- Substitution so far -> result
220 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
222 -- The *InScopeSet* in these Substs gives variables bound so far in the
223 -- target term. So when matching forall a. (\x. a x) against (\y. y y)
224 -- while processing the body of the lambdas, the in-scope set will be {y}.
225 -- That lets us do the occurs-check when matching 'a' against 'y'
227 match :: CoreExpr -- Template
228 -> CoreExpr -- Target
233 match (Var v1) e2 tpl_vars kont subst
234 = case lookupSubst subst v1 of
235 Nothing | v1 `elemVarSet` tpl_vars -- v1 is a template variable
236 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
237 match_fail -- Occurs check failure
238 -- e.g. match forall a. (\x-> a x) against (\y. y y)
240 kont (extendSubst subst v1 (DoneEx e2))
243 | eqExpr (Var v1) e2 -> kont subst
244 -- v1 is not a template variable, so it must be a global constant
246 Just (DoneEx e2') | eqExpr e2' e2 -> kont subst
250 match (Lit lit1) (Lit lit2) tpl_vars kont subst
254 match (App f1 a1) (App f2 a2) tpl_vars kont subst
255 = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
257 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
258 = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
260 -- This rule does eta expansion
261 -- (\x.M) ~ N iff M ~ N x
263 match (Lam x1 e1) e2 tpl_vars kont subst
264 = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
266 -- Eta expansion the other way
267 -- M ~ (\y.N) iff \y.M y ~ \y.N
269 -- Remembering that by (A), y can't be free in M, we get this
270 match e1 (Lam x2 e2) tpl_vars kont subst
271 = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
273 new_id = uniqAway (substInScope subst) x2
274 -- This uniqAway is actually needed. Here's the example:
275 -- rule: foldr (mapFB (:) f) [] = mapList
276 -- target: foldr (\x. mapFB k f x) []
278 -- k = \x. mapFB ... x
279 -- The first \x is ok, but when we inline k, hoping it might
280 -- match (:) we find a second \x.
282 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
283 = match e1 e2 tpl_vars case_kont subst
285 case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
288 match (Type ty1) (Type ty2) tpl_vars kont subst
289 = match_ty ty1 ty2 tpl_vars kont subst
291 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
293 = (match_ty to1 to2 tpl_vars $
294 match_ty from1 from2 tpl_vars $
295 match e1 e2 tpl_vars kont) subst
298 {- I don't buy this let-rule any more
299 The let rule fails on matching
300 forall f,x,xs. f (x:xs)
302 f (let y = e in (y:[]))
303 because we just get x->y, which is bogus.
305 -- This is an interesting rule: we simply ignore lets in the
306 -- term being matched against! The unfolding inside it is (by assumption)
307 -- already inside any occurrences of the bound variables, so we'll expand
308 -- them when we encounter them. Meanwhile, we can't get false matches because
309 -- (also by assumption) the term being matched has no shadowing.
310 match e1 (Let bind e2) tpl_vars kont subst
311 = match e1 e2 tpl_vars kont subst
314 -- Here is another important rule: if the term being matched is a
315 -- variable, we expand it so long as its unfolding is a WHNF
316 -- (Its occurrence information is not necessarily up to date,
317 -- so we don't use it.)
318 match e1 (Var v2) tpl_vars kont subst
319 | isCheapUnfolding unfolding
320 = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
322 unfolding = idUnfolding v2
325 -- We can't cope with lets in the template
327 match e1 e2 tpl_vars kont subst = match_fail
330 ------------------------------------------
331 match_alts [] [] tpl_vars kont subst
333 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
335 = bind vs1 vs2 (match r1 r2) tpl_vars
336 (match_alts alts1 alts2 tpl_vars kont)
338 match_alts alts1 alts2 tpl_vars kont subst = match_fail
340 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
342 ----------------------------------------
343 bind :: [CoreBndr] -- Template binders
344 -> [CoreBndr] -- Target binders
347 -- This makes uses of assumption (A) above. For example,
349 -- Template: (\x.y) (y is free)
350 -- Target : (\y.y) (y is bound)
351 -- We rename x to y in the template... but then erroneously
352 -- match y against y. But this can't happen because of (A)
353 bind vs1 vs2 matcher tpl_vars kont subst
354 = WARN( not (all not_in_subst vs1), bug_msg )
355 matcher tpl_vars kont' subst'
357 kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
358 subst' = bindSubstList subst vs1 vs2
360 -- The unBindSubst relies on no shadowing in the template
361 not_in_subst v = not (maybeToBool (lookupSubst subst v))
362 bug_msg = sep [ppr vs1, ppr vs2]
364 ----------------------------------------
365 match_ty ty1 ty2 tpl_vars kont subst
366 = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
367 Nothing -> match_fail
368 Just senv' -> kont (setSubstEnv subst senv')
370 ----------------------------------------
371 matches [] [] tpl_vars kont subst
373 matches (e:es) (e':es') tpl_vars kont subst
374 = match e e' tpl_vars (matches es es' tpl_vars kont) subst
375 matches es es' tpl_vars kont subst
378 ----------------------------------------
379 mkVarArg :: CoreBndr -> CoreArg
380 mkVarArg v | isId v = Var v
381 | otherwise = Type (mkTyVarTy v)
384 %************************************************************************
386 \subsection{Adding a new rule}
388 %************************************************************************
391 addRule :: Id -> CoreRules -> CoreRule -> CoreRules
393 -- Insert the new rule just before a rule that is *less specific*
394 -- than the new one; or at the end if there isn't such a one.
395 -- In this way we make sure that when looking up, the first match
396 -- is the most specific.
398 -- We make no check for rules that unify without one dominating
399 -- the other. Arguably this would be a bug.
401 addRule id (Rules rules rhs_fvs) rule@(BuiltinRule _)
402 = Rules (rule:rules) rhs_fvs
403 -- Put it at the start for lack of anything better
405 addRule id (Rules rules rhs_fvs) rule
406 = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
408 new_rule = occurAnalyseRule rule
409 new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
411 -- Don't include the Id in its own rhs free-var set.
412 -- Otherwise the occurrence analyser makes bindings recursive
413 -- that shoudn't be. E.g.
414 -- RULE: f (f x y) z ==> f x (f y z)
416 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
419 tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
420 -- Actually we should probably include the free vars of tpl_args,
421 -- but I can't be bothered
424 go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
425 | otherwise = rule : go rules
427 new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
429 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
430 addIdSpecialisations id spec_stuff
431 = setIdSpecialisation id new_rules
433 rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
434 new_rules = foldr add (idSpecialisation id) spec_stuff
435 add (vars, args, rhs) rules = addRule id rules (Rule rule_name vars args rhs)
439 %************************************************************************
441 \subsection{Preparing the rule base
443 %************************************************************************
448 Bool -- True <=> this rule was defined in this module,
449 Id -- What Id is it for
450 CoreRule -- The rule itself
453 pprProtoCoreRule (ProtoCoreRule _ fn rule) = pprCoreRule (ppr fn) rule
455 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
456 lookupRule in_scope fn args
457 = case idSpecialisation fn of
458 Rules rules _ -> matchRules in_scope rules args
460 localRule :: ProtoCoreRule -> Bool
461 localRule (ProtoCoreRule local _ _) = local
463 orphanRule :: ProtoCoreRule -> Bool
464 -- An "orphan rule" is one that is defined in this
465 -- module, but for an *imported* function. We need
466 -- to track these separately when generating the interface file
467 orphanRule (ProtoCoreRule local fn _)
468 = local && not (isLocallyDefined fn)
472 %************************************************************************
474 \subsection{Getting the rules ready}
476 %************************************************************************
479 type RuleBase = (IdSet, -- Imported Ids that have rules attached
480 IdSet) -- Ids (whether local or imported) mentioned on
481 -- LHS of some rule; these should be black listed
483 unionRuleBase (rule_ids1, black_ids1) (rule_ids2, black_ids2)
484 = (plusUFM_C merge_rules rule_ids1 rule_ids2,
485 unionVarSet black_ids1 black_ids2)
487 merge_rules id1 id2 = let rules1 = idSpecialisation id1
488 rules2 = idSpecialisation id2
489 new_rules = foldl (addRule id1) rules1 (rulesRules rules2)
491 setIdSpecialisation id1 new_rules
493 pprRuleBase :: RuleBase -> SDoc
494 pprRuleBase (rules,_) = vcat [ pprCoreRule (ppr id) rs
495 | id <- varSetElems rules,
496 rs <- rulesRules $ idSpecialisation id ]
498 -- prepareLocalRuleBase takes the CoreBinds and rules defined in this module.
499 -- It attaches those rules that are for local Ids to their binders, and
500 -- returns the remainder attached to Ids in an IdSet. It also returns
501 -- Ids mentioned on LHS of some rule; these should be blacklisted.
503 -- The rule Ids and LHS Ids are black-listed; that is, they aren't inlined
504 -- so that the opportunity to apply the rule isn't lost too soon
506 prepareLocalRuleBase :: [CoreBind] -> [ProtoCoreRule] -> ([CoreBind], RuleBase)
507 prepareLocalRuleBase binds local_rules
508 = (map zap_bind binds, (imported_id_rule_ids, rule_lhs_fvs))
510 (rule_ids, rule_lhs_fvs) = foldr add_rule (emptyVarSet, emptyVarSet) local_rules
511 imported_id_rule_ids = filterVarSet (not . isLocallyDefined) rule_ids
513 -- rule_fvs is the set of all variables mentioned in this module's rules
514 rule_fvs = foldVarSet (unionVarSet . idRuleVars) rule_lhs_fvs rule_ids
516 -- Attach the rules for each locally-defined Id to that Id.
517 -- - This makes the rules easier to look up
518 -- - It means that transformation rules and specialisations for
519 -- locally defined Ids are handled uniformly
520 -- - It keeps alive things that are referred to only from a rule
521 -- (the occurrence analyser knows about rules attached to Ids)
522 -- - It makes sure that, when we apply a rule, the free vars
523 -- of the RHS are more likely to be in scope
525 -- The LHS and RHS Ids are marked 'no-discard'.
526 -- This means that the binding won't be discarded EVEN if the binding
527 -- ends up being trivial (v = w) -- the simplifier would usually just
528 -- substitute w for v throughout, but we don't apply the substitution to
529 -- the rules (maybe we should?), so this substitution would make the rule
531 zap_bind (NonRec b r) = NonRec (zap_bndr b) r
532 zap_bind (Rec prs) = Rec [(zap_bndr b, r) | (b,r) <- prs]
534 zap_bndr bndr = case lookupVarSet rule_ids bndr of
535 Just bndr' -> setIdNoDiscard bndr'
536 Nothing | bndr `elemVarSet` rule_fvs -> setIdNoDiscard bndr
539 add_rule (ProtoCoreRule _ id rule)
540 (rule_id_set, rule_fvs)
541 = (rule_id_set `extendVarSet` new_id,
542 rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
544 new_id = case lookupVarSet rule_id_set id of
545 Just id' -> addRuleToId id' rule
546 Nothing -> addRuleToId id rule
547 lhs_fvs = ruleSomeLhsFreeVars isId rule
548 -- Find *all* the free Ids of the LHS, not just
549 -- locally defined ones!!
551 addRuleToId id rule = setIdSpecialisation id (addRule id (idSpecialisation id) rule)
553 -- prepareOrphanRuleBase does exactly the same as prepareLocalRuleBase, except that
554 -- it assumes that none of the rules can be attached to local Ids.
556 prepareOrphanRuleBase :: [ProtoCoreRule] -> RuleBase
557 prepareOrphanRuleBase imported_rules
558 = foldr add_rule (emptyVarSet, emptyVarSet) imported_rules