2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[CoreRules]{Transformation rules}
8 RuleBase, emptyRuleBase,
9 extendRuleBase, extendRuleBaseList, addRuleBaseFVs,
10 ruleBaseIds, ruleBaseFVs,
13 lookupRule, addRule, addIdSpecialisations
16 #include "HsVersions.h"
18 import CoreSyn -- All of it
19 import OccurAnal ( occurAnalyseRule )
20 import CoreFVs ( exprFreeVars, ruleRhsFreeVars, ruleSomeLhsFreeVars )
21 import CoreUnfold ( isCheapUnfolding, unfoldingTemplate )
22 import CoreUtils ( eqExpr )
23 import PprCore ( pprCoreRule )
24 import Subst ( Subst, InScopeSet, mkInScopeSet, lookupSubst, extendSubst,
25 substEnv, setSubstEnv, emptySubst, isInScope,
26 bindSubstList, unBindSubstList, substInScope, uniqAway
28 import Id ( Id, idUnfolding, zapLamIdInfo,
29 idSpecialisation, setIdSpecialisation
31 import Var ( isTyVar, isId )
34 import Type ( mkTyVarTy )
35 import qualified Unify ( match )
38 import Maybes ( maybeToBool )
39 import Util ( sortLt )
43 %************************************************************************
45 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
47 %************************************************************************
49 A @CoreRule@ holds details of one rule for an @Id@, which
50 includes its specialisations.
52 For example, if a rule for @f@ contains the mapping:
54 forall a b d. [Type (List a), Type b, Var d] ===> f' a b
56 then when we find an application of f to matching types, we simply replace
57 it by the matching RHS:
59 f (List Int) Bool dict ===> f' Int Bool
61 All the stuff about how many dictionaries to discard, and what types
62 to apply the specialised function to, are handled by the fact that the
63 Rule contains a template for the result of the specialisation.
65 There is one more exciting case, which is dealt with in exactly the same
66 way. If the specialised value is unboxed then it is lifted at its
67 definition site and unlifted at its uses. For example:
69 pi :: forall a. Num a => a
71 might have a specialisation
73 [Int#] ===> (case pi' of Lift pi# -> pi#)
75 where pi' :: Lift Int# is the specialised version of pi.
78 %************************************************************************
82 %************************************************************************
85 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
86 -- See comments on matchRule
87 matchRules in_scope [] args = Nothing
88 matchRules in_scope (rule:rules) args
89 = case matchRule in_scope rule args of
90 Just result -> Just result
91 Nothing -> matchRules in_scope rules args
94 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
96 -- If (matchRule rule args) returns Just (name,rhs)
97 -- then (f args) matches the rule, and the corresponding
98 -- rewritten RHS is rhs
100 -- The bndrs and rhs is occurrence-analysed
105 -- forall f g x. map f (map g x) ==> map (f . g) x
107 -- CoreRule "map/map"
108 -- [f,g,x] -- tpl_vars
109 -- [f,map g x] -- tpl_args
110 -- map (f.g) x) -- rhs
112 -- Then the call: matchRule the_rule [e1,map e2 e3]
113 -- = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
115 -- Any 'surplus' arguments in the input are simply put on the end
119 -- A1. No top-level variable is bound in the target
120 -- A2. No template variable is bound in the target
121 -- A3. No lambda bound template variable is free in any subexpression of the target
123 -- To see why A1 is necessary, consider matching
124 -- \x->f against \f->f
125 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
126 -- and then erroneously succeed in matching f against f.
128 -- To see why A2 is needed consider matching
129 -- forall a. \b->b against \a->3
130 -- When we meet the lambdas we substitute [a/b] in the template, and then
131 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
133 -- A3 is needed to validate the rule that says
136 -- (\x->E) matches (\x->F x)
139 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
141 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
142 = go tpl_args args emptySubst
143 -- We used to use the in_scope set, but I don't think that's necessary
144 -- After all, the result is going to be simplified again with that in_scope set
146 tpl_var_set = mkVarSet tpl_vars
148 -----------------------
150 go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
152 -- Two easy ways to terminate
153 go [] [] subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
154 go [] args subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
156 -- One tiresome way to terminate: check for excess unmatched
157 -- template arguments
158 go tpl_args [] subst = Nothing -- Failure
161 -----------------------
162 app_match subst fn vs = foldl go fn vs
164 senv = substEnv subst
165 go fn v = case lookupSubstEnv senv v of
166 Just (DoneEx ex) -> fn `App` ex
167 Just (DoneTy ty) -> fn `App` Type ty
168 -- Substitution should bind them all!
171 -----------------------
172 {- The code below tries to match even if there are more
173 template args than real args.
175 I now think this is probably a bad idea.
176 Should the template (map f xs) match (map g)? I think not.
177 For a start, in general eta expansion wastes work.
180 = case eta_complete tpl_args (mkVarSet leftovers) of
181 Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
182 mk_result_args subst done)
183 Nothing -> Nothing -- Failure
185 (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
186 (map zapOccInfo tpl_vars)
188 subst_env = substEnv subst
190 -----------------------
191 eta_complete [] vars = ASSERT( isEmptyVarSet vars )
193 eta_complete (Type ty:tpl_args) vars
194 = case getTyVar_maybe ty of
195 Just tv | tv `elemVarSet` vars
196 -> case eta_complete tpl_args (vars `delVarSet` tv) of
197 Just vars' -> Just (tv:vars')
201 eta_complete (Var v:tpl_args) vars
202 | v `elemVarSet` vars
203 = case eta_complete tpl_args (vars `delVarSet` v) of
204 Just vars' -> Just (v:vars')
207 eta_complete other vars = Nothing
210 zapOccInfo bndr | isTyVar bndr = bndr
211 | otherwise = zapLamIdInfo bndr
216 type Matcher result = VarSet -- Template variables
217 -> (Subst -> Maybe result) -- Continuation if success
218 -> Subst -> Maybe result -- Substitution so far -> result
219 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
221 -- The *InScopeSet* in these Substs gives variables bound so far in the
222 -- target term. So when matching forall a. (\x. a x) against (\y. y y)
223 -- while processing the body of the lambdas, the in-scope set will be {y}.
224 -- That lets us do the occurs-check when matching 'a' against 'y'
226 match :: CoreExpr -- Template
227 -> CoreExpr -- Target
232 match (Var v1) e2 tpl_vars kont subst
233 = case lookupSubst subst v1 of
234 Nothing | v1 `elemVarSet` tpl_vars -- v1 is a template variable
235 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
236 match_fail -- Occurs check failure
237 -- e.g. match forall a. (\x-> a x) against (\y. y y)
239 kont (extendSubst subst v1 (DoneEx e2))
242 | eqExpr (Var v1) e2 -> kont subst
243 -- v1 is not a template variable, so it must be a global constant
245 Just (DoneEx e2') | eqExpr e2' e2 -> kont subst
249 match (Lit lit1) (Lit lit2) tpl_vars kont subst
253 match (App f1 a1) (App f2 a2) tpl_vars kont subst
254 = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
256 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
257 = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
259 -- This rule does eta expansion
260 -- (\x.M) ~ N iff M ~ N x
262 match (Lam x1 e1) e2 tpl_vars kont subst
263 = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
265 -- Eta expansion the other way
266 -- M ~ (\y.N) iff \y.M y ~ \y.N
268 -- Remembering that by (A), y can't be free in M, we get this
269 match e1 (Lam x2 e2) tpl_vars kont subst
270 = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
272 new_id = uniqAway (substInScope subst) x2
273 -- This uniqAway is actually needed. Here's the example:
274 -- rule: foldr (mapFB (:) f) [] = mapList
275 -- target: foldr (\x. mapFB k f x) []
277 -- k = \x. mapFB ... x
278 -- The first \x is ok, but when we inline k, hoping it might
279 -- match (:) we find a second \x.
281 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
282 = match e1 e2 tpl_vars case_kont subst
284 case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
287 match (Type ty1) (Type ty2) tpl_vars kont subst
288 = match_ty ty1 ty2 tpl_vars kont subst
290 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
292 = (match_ty to1 to2 tpl_vars $
293 match_ty from1 from2 tpl_vars $
294 match e1 e2 tpl_vars kont) subst
297 {- I don't buy this let-rule any more
298 The let rule fails on matching
299 forall f,x,xs. f (x:xs)
301 f (let y = e in (y:[]))
302 because we just get x->y, which is bogus.
304 -- This is an interesting rule: we simply ignore lets in the
305 -- term being matched against! The unfolding inside it is (by assumption)
306 -- already inside any occurrences of the bound variables, so we'll expand
307 -- them when we encounter them. Meanwhile, we can't get false matches because
308 -- (also by assumption) the term being matched has no shadowing.
309 match e1 (Let bind e2) tpl_vars kont subst
310 = match e1 e2 tpl_vars kont subst
313 -- Here is another important rule: if the term being matched is a
314 -- variable, we expand it so long as its unfolding is a WHNF
315 -- (Its occurrence information is not necessarily up to date,
316 -- so we don't use it.)
317 match e1 (Var v2) tpl_vars kont subst
318 | isCheapUnfolding unfolding
319 = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
321 unfolding = idUnfolding v2
324 -- We can't cope with lets in the template
326 match e1 e2 tpl_vars kont subst = match_fail
329 ------------------------------------------
330 match_alts [] [] tpl_vars kont subst
332 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
334 = bind vs1 vs2 (match r1 r2) tpl_vars
335 (match_alts alts1 alts2 tpl_vars kont)
337 match_alts alts1 alts2 tpl_vars kont subst = match_fail
339 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
341 ----------------------------------------
342 bind :: [CoreBndr] -- Template binders
343 -> [CoreBndr] -- Target binders
346 -- This makes uses of assumption (A) above. For example,
348 -- Template: (\x.y) (y is free)
349 -- Target : (\y.y) (y is bound)
350 -- We rename x to y in the template... but then erroneously
351 -- match y against y. But this can't happen because of (A)
352 bind vs1 vs2 matcher tpl_vars kont subst
353 = WARN( not (all not_in_subst vs1), bug_msg )
354 matcher tpl_vars kont' subst'
356 kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
357 subst' = bindSubstList subst vs1 vs2
359 -- The unBindSubst relies on no shadowing in the template
360 not_in_subst v = not (maybeToBool (lookupSubst subst v))
361 bug_msg = sep [ppr vs1, ppr vs2]
363 ----------------------------------------
364 match_ty ty1 ty2 tpl_vars kont subst
365 = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
366 Nothing -> match_fail
367 Just senv' -> kont (setSubstEnv subst senv')
369 ----------------------------------------
370 matches [] [] tpl_vars kont subst
372 matches (e:es) (e':es') tpl_vars kont subst
373 = match e e' tpl_vars (matches es es' tpl_vars kont) subst
374 matches es es' tpl_vars kont subst
377 ----------------------------------------
378 mkVarArg :: CoreBndr -> CoreArg
379 mkVarArg v | isId v = Var v
380 | otherwise = Type (mkTyVarTy v)
383 %************************************************************************
385 \subsection{Adding a new rule}
387 %************************************************************************
390 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
392 -- Insert the new rule just before a rule that is *less specific*
393 -- than the new one; or at the end if there isn't such a one.
394 -- In this way we make sure that when looking up, the first match
395 -- is the most specific.
397 -- We make no check for rules that unify without one dominating
398 -- the other. Arguably this would be a bug.
400 addRule (Rules rules rhs_fvs) id rule@(BuiltinRule _)
401 = Rules (rule:rules) rhs_fvs
402 -- Put it at the start for lack of anything better
404 addRule (Rules rules rhs_fvs) id rule
405 = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
407 new_rule = occurAnalyseRule rule
408 new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
410 -- Don't include the Id in its own rhs free-var set.
411 -- Otherwise the occurrence analyser makes bindings recursive
412 -- that shoudn't be. E.g.
413 -- RULE: f (f x y) z ==> f x (f y z)
415 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
418 tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
419 -- Actually we should probably include the free vars of tpl_args,
420 -- but I can't be bothered
423 go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
424 | otherwise = rule : go rules
426 new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
428 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
429 addIdSpecialisations id spec_stuff
430 = setIdSpecialisation id new_rules
432 rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
433 new_rules = foldr add (idSpecialisation id) spec_stuff
434 add (vars, args, rhs) rules = addRule rules id (Rule rule_name vars args rhs)
438 %************************************************************************
440 \subsection{Preparing the rule base
442 %************************************************************************
445 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
446 lookupRule in_scope fn args
447 = case idSpecialisation fn of
448 Rules rules _ -> matchRules in_scope rules args
452 %************************************************************************
454 \subsection{Getting the rules ready}
456 %************************************************************************
459 data RuleBase = RuleBase
460 IdSet -- Ids with their rules in their specialisations
461 -- Held as a set, so that it can simply be the initial
462 -- in-scope set in the simplifier
464 IdSet -- Ids (whether local or imported) mentioned on
465 -- LHS of some rule; these should be black listed
467 -- This representation is a bit cute, and I wonder if we should
468 -- change it to use (IdEnv CoreRule) which seems a bit more natural
470 ruleBaseIds (RuleBase ids _) = ids
471 ruleBaseFVs (RuleBase _ fvs) = fvs
473 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
475 addRuleBaseFVs :: RuleBase -> IdSet -> RuleBase
476 addRuleBaseFVs (RuleBase rules fvs) extra_fvs
477 = RuleBase rules (fvs `unionVarSet` extra_fvs)
479 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
480 extendRuleBaseList rule_base new_guys
481 = foldl extendRuleBase rule_base new_guys
483 extendRuleBase :: RuleBase -> (Id,CoreRule) -> RuleBase
484 extendRuleBase (RuleBase rule_ids rule_fvs) (id, rule)
485 = RuleBase (extendVarSet rule_ids new_id)
486 (rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
488 new_id = setIdSpecialisation id (addRule old_rules id rule)
489 old_rules = case lookupVarSet rule_ids id of
490 Nothing -> emptyCoreRules
491 Just id' -> idSpecialisation id'
493 lhs_fvs = ruleSomeLhsFreeVars isId rule
494 -- Find *all* the free Ids of the LHS, not just
495 -- locally defined ones!!
497 pprRuleBase :: RuleBase -> SDoc
498 pprRuleBase (RuleBase rules _) = vcat [ pprCoreRule (ppr id) rs
499 | id <- varSetElems rules,
500 rs <- rulesRules $ idSpecialisation id ]