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, ruleLhsFreeIds )
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, idSpecialisation, setIdSpecialisation )
32 import Type ( mkTyVarTy )
33 import qualified Unify ( match )
36 import Maybe ( isJust, isNothing, fromMaybe )
37 import Util ( sortLt )
41 %************************************************************************
43 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
45 %************************************************************************
47 A @CoreRule@ holds details of one rule for an @Id@, which
48 includes its specialisations.
50 For example, if a rule for @f@ contains the mapping:
52 forall a b d. [Type (List a), Type b, Var d] ===> f' a b
54 then when we find an application of f to matching types, we simply replace
55 it by the matching RHS:
57 f (List Int) Bool dict ===> f' Int Bool
59 All the stuff about how many dictionaries to discard, and what types
60 to apply the specialised function to, are handled by the fact that the
61 Rule contains a template for the result of the specialisation.
63 There is one more exciting case, which is dealt with in exactly the same
64 way. If the specialised value is unboxed then it is lifted at its
65 definition site and unlifted at its uses. For example:
67 pi :: forall a. Num a => a
69 might have a specialisation
71 [Int#] ===> (case pi' of Lift pi# -> pi#)
73 where pi' :: Lift Int# is the specialised version of pi.
76 %************************************************************************
80 %************************************************************************
83 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
84 -- See comments on matchRule
85 matchRules in_scope [] args = Nothing
86 matchRules in_scope (rule:rules) args
87 = case matchRule in_scope rule args of
88 Just result -> Just result
89 Nothing -> matchRules in_scope rules args
92 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
94 -- If (matchRule rule args) returns Just (name,rhs)
95 -- then (f args) matches the rule, and the corresponding
96 -- rewritten RHS is rhs
98 -- The bndrs and rhs is occurrence-analysed
103 -- forall f g x. map f (map g x) ==> map (f . g) x
105 -- CoreRule "map/map"
106 -- [f,g,x] -- tpl_vars
107 -- [f,map g x] -- tpl_args
108 -- map (f.g) x) -- rhs
110 -- Then the call: matchRule the_rule [e1,map e2 e3]
111 -- = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
113 -- Any 'surplus' arguments in the input are simply put on the end
117 -- A1. No top-level variable is bound in the target
118 -- A2. No template variable is bound in the target
119 -- A3. No lambda bound template variable is free in any subexpression of the target
121 -- To see why A1 is necessary, consider matching
122 -- \x->f against \f->f
123 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
124 -- and then erroneously succeed in matching f against f.
126 -- To see why A2 is needed consider matching
127 -- forall a. \b->b against \a->3
128 -- When we meet the lambdas we substitute [a/b] in the template, and then
129 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
131 -- A3 is needed to validate the rule that says
134 -- (\x->E) matches (\x->F x)
137 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
139 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
140 = go tpl_args args emptySubst
141 -- We used to use the in_scope set, but I don't think that's necessary
142 -- After all, the result is going to be simplified again with that in_scope set
144 tpl_var_set = mkVarSet tpl_vars
146 -----------------------
148 go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
150 -- Two easy ways to terminate
151 go [] [] subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
152 go [] args subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
154 -- One tiresome way to terminate: check for excess unmatched
155 -- template arguments
156 go tpl_args [] subst = Nothing -- Failure
159 -----------------------
160 app_match subst fn vs = foldl go fn vs
162 senv = substEnv subst
163 go fn v = case lookupSubstEnv senv v of
164 Just (DoneEx ex) -> fn `App` ex
165 Just (DoneTy ty) -> fn `App` Type ty
166 -- Substitution should bind them all!
169 -----------------------
170 {- The code below tries to match even if there are more
171 template args than real args.
173 I now think this is probably a bad idea.
174 Should the template (map f xs) match (map g)? I think not.
175 For a start, in general eta expansion wastes work.
178 = case eta_complete tpl_args (mkVarSet leftovers) of
179 Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
180 mk_result_args subst done)
181 Nothing -> Nothing -- Failure
183 (done, leftovers) = partition (\v -> isJust (lookupSubstEnv subst_env v))
184 (map zapOccInfo tpl_vars)
186 subst_env = substEnv subst
188 -----------------------
189 eta_complete [] vars = ASSERT( isEmptyVarSet vars )
191 eta_complete (Type ty:tpl_args) vars
192 = case getTyVar_maybe ty of
193 Just tv | tv `elemVarSet` vars
194 -> case eta_complete tpl_args (vars `delVarSet` tv) of
195 Just vars' -> Just (tv:vars')
199 eta_complete (Var v:tpl_args) vars
200 | v `elemVarSet` vars
201 = case eta_complete tpl_args (vars `delVarSet` v) of
202 Just vars' -> Just (v:vars')
205 eta_complete other vars = Nothing
208 zapOccInfo bndr | isTyVar bndr = bndr
209 | otherwise = zapLamIdInfo bndr
214 type Matcher result = VarSet -- Template variables
215 -> (Subst -> Maybe result) -- Continuation if success
216 -> Subst -> Maybe result -- Substitution so far -> result
217 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
219 -- The *InScopeSet* in these Substs gives variables bound so far in the
220 -- target term. So when matching forall a. (\x. a x) against (\y. y y)
221 -- while processing the body of the lambdas, the in-scope set will be {y}.
222 -- That lets us do the occurs-check when matching 'a' against 'y'
224 match :: CoreExpr -- Template
225 -> CoreExpr -- Target
230 match (Var v1) e2 tpl_vars kont subst
231 = case lookupSubst subst v1 of
232 Nothing | v1 `elemVarSet` tpl_vars -- v1 is a template variable
233 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
234 match_fail -- Occurs check failure
235 -- e.g. match forall a. (\x-> a x) against (\y. y y)
237 kont (extendSubst subst v1 (DoneEx e2))
240 | eqExpr (Var v1) e2 -> kont subst
241 -- v1 is not a template variable, so it must be a global constant
243 Just (DoneEx e2') | eqExpr e2' e2 -> kont subst
247 match (Lit lit1) (Lit lit2) tpl_vars kont subst
251 match (App f1 a1) (App f2 a2) tpl_vars kont subst
252 = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
254 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
255 = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
257 -- This rule does eta expansion
258 -- (\x.M) ~ N iff M ~ N x
260 match (Lam x1 e1) e2 tpl_vars kont subst
261 = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
263 -- Eta expansion the other way
264 -- M ~ (\y.N) iff \y.M y ~ \y.N
266 -- Remembering that by (A), y can't be free in M, we get this
267 match e1 (Lam x2 e2) tpl_vars kont subst
268 = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
270 new_id = uniqAway (substInScope subst) x2
271 -- This uniqAway is actually needed. Here's the example:
272 -- rule: foldr (mapFB (:) f) [] = mapList
273 -- target: foldr (\x. mapFB k f x) []
275 -- k = \x. mapFB ... x
276 -- The first \x is ok, but when we inline k, hoping it might
277 -- match (:) we find a second \x.
279 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
280 = match e1 e2 tpl_vars case_kont subst
282 case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
285 match (Type ty1) (Type ty2) tpl_vars kont subst
286 = match_ty ty1 ty2 tpl_vars kont subst
288 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
290 = (match_ty to1 to2 tpl_vars $
291 match_ty from1 from2 tpl_vars $
292 match e1 e2 tpl_vars kont) subst
295 {- I don't buy this let-rule any more
296 The let rule fails on matching
297 forall f,x,xs. f (x:xs)
299 f (let y = e in (y:[]))
300 because we just get x->y, which is bogus.
302 -- This is an interesting rule: we simply ignore lets in the
303 -- term being matched against! The unfolding inside it is (by assumption)
304 -- already inside any occurrences of the bound variables, so we'll expand
305 -- them when we encounter them. Meanwhile, we can't get false matches because
306 -- (also by assumption) the term being matched has no shadowing.
307 match e1 (Let bind e2) tpl_vars kont subst
308 = match e1 e2 tpl_vars kont subst
311 -- Here is another important rule: if the term being matched is a
312 -- variable, we expand it so long as its unfolding is a WHNF
313 -- (Its occurrence information is not necessarily up to date,
314 -- so we don't use it.)
315 match e1 (Var v2) tpl_vars kont subst
316 | isCheapUnfolding unfolding
317 = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
319 unfolding = idUnfolding v2
322 -- We can't cope with lets in the template
324 match e1 e2 tpl_vars kont subst = match_fail
327 ------------------------------------------
328 match_alts [] [] tpl_vars kont subst
330 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
332 = bind vs1 vs2 (match r1 r2) tpl_vars
333 (match_alts alts1 alts2 tpl_vars kont)
335 match_alts alts1 alts2 tpl_vars kont subst = match_fail
337 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
339 ----------------------------------------
340 bind :: [CoreBndr] -- Template binders
341 -> [CoreBndr] -- Target binders
344 -- This makes uses of assumption (A) above. For example,
346 -- Template: (\x.y) (y is free)
347 -- Target : (\y.y) (y is bound)
348 -- We rename x to y in the template... but then erroneously
349 -- match y against y. But this can't happen because of (A)
350 bind vs1 vs2 matcher tpl_vars kont subst
351 = WARN( not (all not_in_subst vs1), bug_msg )
352 matcher tpl_vars kont' subst'
354 kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
355 subst' = bindSubstList subst vs1 vs2
357 -- The unBindSubst relies on no shadowing in the template
358 not_in_subst v = isNothing (lookupSubst subst v)
359 bug_msg = sep [ppr vs1, ppr vs2]
361 ----------------------------------------
362 match_ty ty1 ty2 tpl_vars kont subst
363 = case Unify.match False {- for now: KSW 2000-10 -} ty1 ty2 tpl_vars Just (substEnv subst) of
364 Nothing -> match_fail
365 Just senv' -> kont (setSubstEnv subst senv')
367 ----------------------------------------
368 matches [] [] tpl_vars kont subst
370 matches (e:es) (e':es') tpl_vars kont subst
371 = match e e' tpl_vars (matches es es' tpl_vars kont) subst
372 matches es es' tpl_vars kont subst
375 ----------------------------------------
376 mkVarArg :: CoreBndr -> CoreArg
377 mkVarArg v | isId v = Var v
378 | otherwise = Type (mkTyVarTy v)
381 %************************************************************************
383 \subsection{Adding a new rule}
385 %************************************************************************
388 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
390 -- Insert the new rule just before a rule that is *less specific*
391 -- than the new one; or at the end if there isn't such a one.
392 -- In this way we make sure that when looking up, the first match
393 -- is the most specific.
395 -- We make no check for rules that unify without one dominating
396 -- the other. Arguably this would be a bug.
398 addRule (Rules rules rhs_fvs) id rule@(BuiltinRule _)
399 = Rules (rule:rules) rhs_fvs
400 -- Put it at the start for lack of anything better
402 addRule (Rules rules rhs_fvs) id rule
403 = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
405 new_rule = occurAnalyseRule rule
406 new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
408 -- Don't include the Id in its own rhs free-var set.
409 -- Otherwise the occurrence analyser makes bindings recursive
410 -- that shoudn't be. E.g.
411 -- RULE: f (f x y) z ==> f x (f y z)
413 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
416 tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
417 -- Actually we should probably include the free vars of tpl_args,
418 -- but I can't be bothered
421 go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
422 | otherwise = rule : go rules
424 new_is_more_specific rule = isJust (matchRule tpl_var_set rule tpl_args)
426 addIdSpecialisations :: Id -> [CoreRule] -> Id
427 addIdSpecialisations id rules
428 = setIdSpecialisation id new_specs
430 new_specs = foldr add (idSpecialisation id) rules
431 add rule rules = addRule rules id rule
435 %************************************************************************
437 \subsection{Preparing the rule base
439 %************************************************************************
442 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
443 lookupRule in_scope fn args
444 = case idSpecialisation fn of
445 Rules rules _ -> matchRules in_scope rules args
449 %************************************************************************
451 \subsection{Getting the rules ready}
453 %************************************************************************
456 data RuleBase = RuleBase
457 IdSet -- Ids with their rules in their specialisations
458 -- Held as a set, so that it can simply be the initial
459 -- in-scope set in the simplifier
461 IdSet -- Ids (whether local or imported) mentioned on
462 -- LHS of some rule; these should be black listed
464 -- This representation is a bit cute, and I wonder if we should
465 -- change it to use (IdEnv CoreRule) which seems a bit more natural
467 ruleBaseIds (RuleBase ids _) = ids
468 ruleBaseFVs (RuleBase _ fvs) = fvs
470 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
472 addRuleBaseFVs :: RuleBase -> IdSet -> RuleBase
473 addRuleBaseFVs (RuleBase rules fvs) extra_fvs
474 = RuleBase rules (fvs `unionVarSet` extra_fvs)
476 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
477 extendRuleBaseList rule_base new_guys
478 = foldl extendRuleBase rule_base new_guys
480 extendRuleBase :: RuleBase -> (Id,CoreRule) -> RuleBase
481 extendRuleBase (RuleBase rule_ids rule_fvs) (id, rule)
482 = RuleBase (extendVarSet rule_ids new_id)
483 (rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
485 new_id = setIdSpecialisation id (addRule old_rules id rule)
487 old_rules = idSpecialisation (fromMaybe id (lookupVarSet rule_ids id))
488 -- Get the old rules from rule_ids if the Id is already there, but
489 -- if not, use the Id from the incoming rule. If may be a PrimOpId,
490 -- in which case it may have rules in its belly already. Seems
491 -- dreadfully hackoid.
493 lhs_fvs = ruleLhsFreeIds rule
494 -- Finds *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 ]