2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[CoreRules]{Transformation rules}
8 RuleBase, emptyRuleBase, extendRuleBase, extendRuleBaseList,
9 prepareLocalRuleBase, prepareOrphanRuleBase,
10 unionRuleBase, lookupRule, addRule, addIdSpecialisations,
11 ProtoCoreRule(..), pprProtoCoreRule, pprRuleBase,
15 #include "HsVersions.h"
17 import CoreSyn -- All of it
18 import OccurAnal ( occurAnalyseRule )
19 import CoreFVs ( exprFreeVars, idRuleVars, ruleRhsFreeVars, ruleSomeLhsFreeVars )
20 import CoreUnfold ( isCheapUnfolding, unfoldingTemplate )
21 import CoreUtils ( eqExpr )
22 import PprCore ( pprCoreRule )
23 import Subst ( Subst, InScopeSet, mkInScopeSet, lookupSubst, extendSubst,
24 substEnv, setSubstEnv, emptySubst, isInScope,
25 bindSubstList, unBindSubstList, substInScope, uniqAway
27 import Id ( Id, idUnfolding, zapLamIdInfo,
28 idSpecialisation, setIdSpecialisation,
31 import Name ( isLocallyDefined )
32 import Var ( isTyVar, isId )
35 import Type ( mkTyVarTy )
36 import qualified Unify ( match )
40 import Maybes ( maybeToBool )
41 import Util ( sortLt )
45 %************************************************************************
47 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
49 %************************************************************************
51 A @CoreRule@ holds details of one rule for an @Id@, which
52 includes its specialisations.
54 For example, if a rule for @f@ contains the mapping:
56 forall a b d. [Type (List a), Type b, Var d] ===> f' a b
58 then when we find an application of f to matching types, we simply replace
59 it by the matching RHS:
61 f (List Int) Bool dict ===> f' Int Bool
63 All the stuff about how many dictionaries to discard, and what types
64 to apply the specialised function to, are handled by the fact that the
65 Rule contains a template for the result of the specialisation.
67 There is one more exciting case, which is dealt with in exactly the same
68 way. If the specialised value is unboxed then it is lifted at its
69 definition site and unlifted at its uses. For example:
71 pi :: forall a. Num a => a
73 might have a specialisation
75 [Int#] ===> (case pi' of Lift pi# -> pi#)
77 where pi' :: Lift Int# is the specialised version of pi.
80 %************************************************************************
84 %************************************************************************
87 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
88 -- See comments on matchRule
89 matchRules in_scope [] args = Nothing
90 matchRules in_scope (rule:rules) args
91 = case matchRule in_scope rule args of
92 Just result -> Just result
93 Nothing -> matchRules in_scope rules args
96 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
98 -- If (matchRule rule args) returns Just (name,rhs)
99 -- then (f args) matches the rule, and the corresponding
100 -- rewritten RHS is rhs
102 -- The bndrs and rhs is occurrence-analysed
107 -- forall f g x. map f (map g x) ==> map (f . g) x
109 -- CoreRule "map/map"
110 -- [f,g,x] -- tpl_vars
111 -- [f,map g x] -- tpl_args
112 -- map (f.g) x) -- rhs
114 -- Then the call: matchRule the_rule [e1,map e2 e3]
115 -- = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
117 -- Any 'surplus' arguments in the input are simply put on the end
121 -- A1. No top-level variable is bound in the target
122 -- A2. No template variable is bound in the target
123 -- A3. No lambda bound template variable is free in any subexpression of the target
125 -- To see why A1 is necessary, consider matching
126 -- \x->f against \f->f
127 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
128 -- and then erroneously succeed in matching f against f.
130 -- To see why A2 is needed consider matching
131 -- forall a. \b->b against \a->3
132 -- When we meet the lambdas we substitute [a/b] in the template, and then
133 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
135 -- A3 is needed to validate the rule that says
138 -- (\x->E) matches (\x->F x)
141 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
143 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
144 = go tpl_args args emptySubst
145 -- We used to use the in_scope set, but I don't think that's necessary
146 -- After all, the result is going to be simplified again with that in_scope set
148 tpl_var_set = mkVarSet tpl_vars
150 -----------------------
152 go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
154 -- Two easy ways to terminate
155 go [] [] subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
156 go [] args subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
158 -- One tiresome way to terminate: check for excess unmatched
159 -- template arguments
160 go tpl_args [] subst = Nothing -- Failure
163 -----------------------
164 app_match subst fn vs = foldl go fn vs
166 senv = substEnv subst
167 go fn v = case lookupSubstEnv senv v of
168 Just (DoneEx ex) -> fn `App` ex
169 Just (DoneTy ty) -> fn `App` Type ty
170 -- Substitution should bind them all!
173 -----------------------
174 {- The code below tries to match even if there are more
175 template args than real args.
177 I now think this is probably a bad idea.
178 Should the template (map f xs) match (map g)? I think not.
179 For a start, in general eta expansion wastes work.
182 = case eta_complete tpl_args (mkVarSet leftovers) of
183 Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
184 mk_result_args subst done)
185 Nothing -> Nothing -- Failure
187 (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
188 (map zapOccInfo tpl_vars)
190 subst_env = substEnv subst
192 -----------------------
193 eta_complete [] vars = ASSERT( isEmptyVarSet vars )
195 eta_complete (Type ty:tpl_args) vars
196 = case getTyVar_maybe ty of
197 Just tv | tv `elemVarSet` vars
198 -> case eta_complete tpl_args (vars `delVarSet` tv) of
199 Just vars' -> Just (tv:vars')
203 eta_complete (Var v:tpl_args) vars
204 | v `elemVarSet` vars
205 = case eta_complete tpl_args (vars `delVarSet` v) of
206 Just vars' -> Just (v:vars')
209 eta_complete other vars = Nothing
213 zapOccInfo bndr | isTyVar bndr = bndr
214 | otherwise = zapLamIdInfo bndr
218 type Matcher result = VarSet -- Template variables
219 -> (Subst -> Maybe result) -- Continuation if success
220 -> Subst -> Maybe result -- Substitution so far -> result
221 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
223 -- The *InScopeSet* in these Substs gives variables bound so far in the
224 -- target term. So when matching forall a. (\x. a x) against (\y. y y)
225 -- while processing the body of the lambdas, the in-scope set will be {y}.
226 -- That lets us do the occurs-check when matching 'a' against 'y'
228 match :: CoreExpr -- Template
229 -> CoreExpr -- Target
234 match (Var v1) e2 tpl_vars kont subst
235 = case lookupSubst subst v1 of
236 Nothing | v1 `elemVarSet` tpl_vars -- v1 is a template variable
237 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
238 match_fail -- Occurs check failure
239 -- e.g. match forall a. (\x-> a x) against (\y. y y)
241 kont (extendSubst subst v1 (DoneEx e2))
244 | eqExpr (Var v1) e2 -> kont subst
245 -- v1 is not a template variable, so it must be a global constant
247 Just (DoneEx e2') | eqExpr e2' e2 -> kont subst
251 match (Lit lit1) (Lit lit2) tpl_vars kont subst
255 match (App f1 a1) (App f2 a2) tpl_vars kont subst
256 = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
258 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
259 = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
261 -- This rule does eta expansion
262 -- (\x.M) ~ N iff M ~ N x
264 match (Lam x1 e1) e2 tpl_vars kont subst
265 = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
267 -- Eta expansion the other way
268 -- M ~ (\y.N) iff \y.M y ~ \y.N
270 -- Remembering that by (A), y can't be free in M, we get this
271 match e1 (Lam x2 e2) tpl_vars kont subst
272 = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
274 new_id = uniqAway (substInScope subst) x2
275 -- This uniqAway is actually needed. Here's the example:
276 -- rule: foldr (mapFB (:) f) [] = mapList
277 -- target: foldr (\x. mapFB k f x) []
279 -- k = \x. mapFB ... x
280 -- The first \x is ok, but when we inline k, hoping it might
281 -- match (:) we find a second \x.
283 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
284 = match e1 e2 tpl_vars case_kont subst
286 case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
289 match (Type ty1) (Type ty2) tpl_vars kont subst
290 = match_ty ty1 ty2 tpl_vars kont subst
292 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
294 = (match_ty to1 to2 tpl_vars $
295 match_ty from1 from2 tpl_vars $
296 match e1 e2 tpl_vars kont) subst
299 {- I don't buy this let-rule any more
300 The let rule fails on matching
301 forall f,x,xs. f (x:xs)
303 f (let y = e in (y:[]))
304 because we just get x->y, which is bogus.
306 -- This is an interesting rule: we simply ignore lets in the
307 -- term being matched against! The unfolding inside it is (by assumption)
308 -- already inside any occurrences of the bound variables, so we'll expand
309 -- them when we encounter them. Meanwhile, we can't get false matches because
310 -- (also by assumption) the term being matched has no shadowing.
311 match e1 (Let bind e2) tpl_vars kont subst
312 = match e1 e2 tpl_vars kont subst
315 -- Here is another important rule: if the term being matched is a
316 -- variable, we expand it so long as its unfolding is a WHNF
317 -- (Its occurrence information is not necessarily up to date,
318 -- so we don't use it.)
319 match e1 (Var v2) tpl_vars kont subst
320 | isCheapUnfolding unfolding
321 = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
323 unfolding = idUnfolding v2
326 -- We can't cope with lets in the template
328 match e1 e2 tpl_vars kont subst = match_fail
331 ------------------------------------------
332 match_alts [] [] tpl_vars kont subst
334 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
336 = bind vs1 vs2 (match r1 r2) tpl_vars
337 (match_alts alts1 alts2 tpl_vars kont)
339 match_alts alts1 alts2 tpl_vars kont subst = match_fail
341 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
343 ----------------------------------------
344 bind :: [CoreBndr] -- Template binders
345 -> [CoreBndr] -- Target binders
348 -- This makes uses of assumption (A) above. For example,
350 -- Template: (\x.y) (y is free)
351 -- Target : (\y.y) (y is bound)
352 -- We rename x to y in the template... but then erroneously
353 -- match y against y. But this can't happen because of (A)
354 bind vs1 vs2 matcher tpl_vars kont subst
355 = WARN( not (all not_in_subst vs1), bug_msg )
356 matcher tpl_vars kont' subst'
358 kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
359 subst' = bindSubstList subst vs1 vs2
361 -- The unBindSubst relies on no shadowing in the template
362 not_in_subst v = not (maybeToBool (lookupSubst subst v))
363 bug_msg = sep [ppr vs1, ppr vs2]
365 ----------------------------------------
366 match_ty ty1 ty2 tpl_vars kont subst
367 = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
368 Nothing -> match_fail
369 Just senv' -> kont (setSubstEnv subst senv')
371 ----------------------------------------
372 matches [] [] tpl_vars kont subst
374 matches (e:es) (e':es') tpl_vars kont subst
375 = match e e' tpl_vars (matches es es' tpl_vars kont) subst
376 matches es es' tpl_vars kont subst
379 ----------------------------------------
380 mkVarArg :: CoreBndr -> CoreArg
381 mkVarArg v | isId v = Var v
382 | otherwise = Type (mkTyVarTy v)
385 %************************************************************************
387 \subsection{Adding a new rule}
389 %************************************************************************
392 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
394 -- Insert the new rule just before a rule that is *less specific*
395 -- than the new one; or at the end if there isn't such a one.
396 -- In this way we make sure that when looking up, the first match
397 -- is the most specific.
399 -- We make no check for rules that unify without one dominating
400 -- the other. Arguably this would be a bug.
402 addRule (Rules rules rhs_fvs) id rule@(BuiltinRule _)
403 = Rules (rule:rules) rhs_fvs
404 -- Put it at the start for lack of anything better
406 addRule (Rules rules rhs_fvs) id rule
407 = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
409 new_rule = occurAnalyseRule rule
410 new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
412 -- Don't include the Id in its own rhs free-var set.
413 -- Otherwise the occurrence analyser makes bindings recursive
414 -- that shoudn't be. E.g.
415 -- RULE: f (f x y) z ==> f x (f y z)
417 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
420 tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
421 -- Actually we should probably include the free vars of tpl_args,
422 -- but I can't be bothered
425 go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
426 | otherwise = rule : go rules
428 new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
430 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
431 addIdSpecialisations id spec_stuff
432 = setIdSpecialisation id new_rules
434 rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
435 new_rules = foldr add (idSpecialisation id) spec_stuff
436 add (vars, args, rhs) rules = addRule rules id (Rule rule_name vars args rhs)
440 %************************************************************************
442 \subsection{Preparing the rule base
444 %************************************************************************
449 Bool -- True <=> this rule was defined in this module,
450 Id -- What Id is it for
451 CoreRule -- The rule itself
454 pprProtoCoreRule (ProtoCoreRule _ fn rule) = pprCoreRule (ppr fn) rule
456 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
457 lookupRule in_scope fn args
458 = case idSpecialisation fn of
459 Rules rules _ -> matchRules in_scope rules args
461 localRule :: ProtoCoreRule -> Bool
462 localRule (ProtoCoreRule local _ _) = local
464 orphanRule :: ProtoCoreRule -> Bool
465 -- An "orphan rule" is one that is defined in this
466 -- module, but for an *imported* function. We need
467 -- to track these separately when generating the interface file
468 orphanRule (ProtoCoreRule local fn _)
469 = local && not (isLocallyDefined fn)
473 %************************************************************************
475 \subsection{Getting the rules ready}
477 %************************************************************************
480 data RuleBase = RuleBase
481 IdSet -- Ids with their rules in their specialisations
482 -- Held as a set, so that it can simply be the initial
483 -- in-scope set in the simplifier
485 IdSet -- Ids (whether local or imported) mentioned on
486 -- LHS of some rule; these should be black listed
488 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
490 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
491 extendRuleBaseList rule_base new_guys
492 = foldl extendRuleBase rule_base new_guys
494 extendRuleBase :: RuleBase -> (Id,CoreRule) -> RuleBase
495 extendRuleBase (RuleBase rule_ids rule_fvs) (id, rule)
496 = RuleBase (extendVarSet rule_ids new_id)
497 (rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
499 new_id = setIdSpecialisation id (addRule old_rules id rule)
500 old_rules = case lookupVarSet rule_ids id of
501 Nothing -> emptyCoreRules
502 Just id' -> idSpecialisation id'
504 lhs_fvs = ruleSomeLhsFreeVars isId rule
505 -- Find *all* the free Ids of the LHS, not just
506 -- locally defined ones!!
508 unionRuleBase (RuleBase rule_ids1 black_ids1) (RuleBase rule_ids2 black_ids2)
509 = RuleBase (plusUFM_C merge_rules rule_ids1 rule_ids2)
510 (unionVarSet black_ids1 black_ids2)
513 merge_rules id1 id2 = let rules1 = idSpecialisation id1
514 rules2 = idSpecialisation id2
515 new_rules = foldl (addRule id1) rules1 (rulesRules rules2)
517 setIdSpecialisation id1 new_rules
519 pprRuleBase :: RuleBase -> SDoc
520 pprRuleBase (RuleBase rules _) = vcat [ pprCoreRule (ppr id) rs
521 | id <- varSetElems rules,
522 rs <- rulesRules $ idSpecialisation id ]
524 -- prepareLocalRuleBase takes the CoreBinds and rules defined in this module.
525 -- It attaches those rules that are for local Ids to their binders, and
526 -- returns the remainder attached to Ids in an IdSet. It also returns
527 -- Ids mentioned on LHS of some rule; these should be blacklisted.
529 -- The rule Ids and LHS Ids are black-listed; that is, they aren't inlined
530 -- so that the opportunity to apply the rule isn't lost too soon
532 prepareLocalRuleBase :: [CoreBind] -> [(Id,CoreRule)] -> ([CoreBind], RuleBase)
533 prepareLocalRuleBase binds local_rules
536 = (map zap_bind binds, RuleBase imported_id_rule_ids rule_lhs_fvs)
538 RuleBase rule_ids rule_lhs_fvs = extendRuleBaseList emptyRuleBase local_rules
539 imported_id_rule_ids = filterVarSet (not . isLocallyDefined) rule_ids
541 -- rule_fvs is the set of all variables mentioned in this module's rules
542 rule_fvs = foldVarSet (unionVarSet . idRuleVars) rule_lhs_fvs rule_ids
544 -- Attach the rules for each locally-defined Id to that Id.
545 -- - This makes the rules easier to look up
546 -- - It means that transformation rules and specialisations for
547 -- locally defined Ids are handled uniformly
548 -- - It keeps alive things that are referred to only from a rule
549 -- (the occurrence analyser knows about rules attached to Ids)
550 -- - It makes sure that, when we apply a rule, the free vars
551 -- of the RHS are more likely to be in scope
553 -- The LHS and RHS Ids are marked 'no-discard'.
554 -- This means that the binding won't be discarded EVEN if the binding
555 -- ends up being trivial (v = w) -- the simplifier would usually just
556 -- substitute w for v throughout, but we don't apply the substitution to
557 -- the rules (maybe we should?), so this substitution would make the rule
559 zap_bind (NonRec b r) = NonRec (zap_bndr b) r
560 zap_bind (Rec prs) = Rec [(zap_bndr b, r) | (b,r) <- prs]
562 zap_bndr bndr = case lookupVarSet rule_ids bndr of
563 Just bndr' -> setIdNoDiscard bndr'
564 Nothing | bndr `elemVarSet` rule_fvs -> setIdNoDiscard bndr
568 addRuleToId id rule = setIdSpecialisation id (addRule (idSpecialisation id) id rule)
570 -- prepareOrphanRuleBase does exactly the same as prepareLocalRuleBase, except that
571 -- it assumes that none of the rules can be attached to local Ids.
573 prepareOrphanRuleBase :: [ProtoCoreRule] -> RuleBase
574 prepareOrphanRuleBase imported_rules
577 = foldr add_rule emptyRuleBase imported_rules