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 TcType ( mkTyVarTy )
33 import qualified TcType ( match )
34 import TypeRep ( Type(..) ) -- Can see type representation for matching
37 import Maybe ( isJust, isNothing, fromMaybe )
38 import Util ( sortLt )
42 %************************************************************************
44 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
46 %************************************************************************
48 A @CoreRule@ holds details of one rule for an @Id@, which
49 includes its specialisations.
51 For example, if a rule for @f@ contains the mapping:
53 forall a b d. [Type (List a), Type b, Var d] ===> f' a b
55 then when we find an application of f to matching types, we simply replace
56 it by the matching RHS:
58 f (List Int) Bool dict ===> f' Int Bool
60 All the stuff about how many dictionaries to discard, and what types
61 to apply the specialised function to, are handled by the fact that the
62 Rule contains a template for the result of the specialisation.
64 There is one more exciting case, which is dealt with in exactly the same
65 way. If the specialised value is unboxed then it is lifted at its
66 definition site and unlifted at its uses. For example:
68 pi :: forall a. Num a => a
70 might have a specialisation
72 [Int#] ===> (case pi' of Lift pi# -> pi#)
74 where pi' :: Lift Int# is the specialised version of pi.
77 %************************************************************************
81 %************************************************************************
84 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
85 -- See comments on matchRule
86 matchRules in_scope [] args = Nothing
87 matchRules in_scope (rule:rules) args
88 = case matchRule in_scope rule args of
89 Just result -> Just result
90 Nothing -> matchRules in_scope rules args
93 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
95 -- If (matchRule rule args) returns Just (name,rhs)
96 -- then (f args) matches the rule, and the corresponding
97 -- rewritten RHS is rhs
99 -- The bndrs and rhs is occurrence-analysed
104 -- forall f g x. map f (map g x) ==> map (f . g) x
106 -- CoreRule "map/map"
107 -- [f,g,x] -- tpl_vars
108 -- [f,map g x] -- tpl_args
109 -- map (f.g) x) -- rhs
111 -- Then the call: matchRule the_rule [e1,map e2 e3]
112 -- = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
114 -- Any 'surplus' arguments in the input are simply put on the end
118 -- A1. No top-level variable is bound in the target
119 -- A2. No template variable is bound in the target
120 -- A3. No lambda bound template variable is free in any subexpression of the target
122 -- To see why A1 is necessary, consider matching
123 -- \x->f against \f->f
124 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
125 -- and then erroneously succeed in matching f against f.
127 -- To see why A2 is needed consider matching
128 -- forall a. \b->b against \a->3
129 -- When we meet the lambdas we substitute [a/b] in the template, and then
130 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
132 -- A3 is needed to validate the rule that says
135 -- (\x->E) matches (\x->F x)
138 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
140 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
141 = go tpl_args args emptySubst
142 -- We used to use the in_scope set, but I don't think that's necessary
143 -- After all, the result is going to be simplified again with that in_scope set
145 tpl_var_set = mkVarSet tpl_vars
147 -----------------------
149 go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
151 -- Two easy ways to terminate
152 go [] [] subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
153 go [] args subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
155 -- One tiresome way to terminate: check for excess unmatched
156 -- template arguments
157 go tpl_args [] subst = Nothing -- Failure
160 -----------------------
161 app_match subst fn vs = foldl go fn vs
163 senv = substEnv subst
164 go fn v = case lookupSubstEnv senv v of
165 Just (DoneEx ex) -> fn `App` ex
166 Just (DoneTy ty) -> fn `App` Type ty
167 -- Substitution should bind them all!
170 -----------------------
171 {- The code below tries to match even if there are more
172 template args than real args.
174 I now think this is probably a bad idea.
175 Should the template (map f xs) match (map g)? I think not.
176 For a start, in general eta expansion wastes work.
179 = case eta_complete tpl_args (mkVarSet leftovers) of
180 Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
181 mk_result_args subst done)
182 Nothing -> Nothing -- Failure
184 (done, leftovers) = partition (\v -> isJust (lookupSubstEnv subst_env v))
185 (map zapOccInfo tpl_vars)
187 subst_env = substEnv subst
189 -----------------------
190 eta_complete [] vars = ASSERT( isEmptyVarSet vars )
192 eta_complete (Type ty:tpl_args) vars
193 = case getTyVar_maybe ty of
194 Just tv | tv `elemVarSet` vars
195 -> case eta_complete tpl_args (vars `delVarSet` tv) of
196 Just vars' -> Just (tv:vars')
200 eta_complete (Var v:tpl_args) vars
201 | v `elemVarSet` vars
202 = case eta_complete tpl_args (vars `delVarSet` v) of
203 Just vars' -> Just (v:vars')
206 eta_complete other vars = Nothing
209 zapOccInfo bndr | isTyVar bndr = bndr
210 | otherwise = zapLamIdInfo bndr
215 type Matcher result = VarSet -- Template variables
216 -> (Subst -> Maybe result) -- Continuation if success
217 -> Subst -> Maybe result -- Substitution so far -> result
218 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
220 -- The *InScopeSet* in these Substs gives variables bound so far in the
221 -- target term. So when matching forall a. (\x. a x) against (\y. y y)
222 -- while processing the body of the lambdas, the in-scope set will be {y}.
223 -- That lets us do the occurs-check when matching 'a' against 'y'
225 match :: CoreExpr -- Template
226 -> CoreExpr -- Target
231 match (Var v1) e2 tpl_vars kont subst
232 = case lookupSubst subst v1 of
233 Nothing | v1 `elemVarSet` tpl_vars -- v1 is a template variable
234 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
235 match_fail -- Occurs check failure
236 -- e.g. match forall a. (\x-> a x) against (\y. y y)
238 kont (extendSubst subst v1 (DoneEx e2))
241 | eqExpr (Var v1) e2 -> kont subst
242 -- v1 is not a template variable, so it must be a global constant
244 Just (DoneEx e2') | eqExpr e2' e2 -> kont subst
248 match (Lit lit1) (Lit lit2) tpl_vars kont subst
252 match (App f1 a1) (App f2 a2) tpl_vars kont subst
253 = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
255 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
256 = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
258 -- This rule does eta expansion
259 -- (\x.M) ~ N iff M ~ N x
261 match (Lam x1 e1) e2 tpl_vars kont subst
262 = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
264 -- Eta expansion the other way
265 -- M ~ (\y.N) iff \y.M y ~ \y.N
267 -- Remembering that by (A), y can't be free in M, we get this
268 match e1 (Lam x2 e2) tpl_vars kont subst
269 = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
271 new_id = uniqAway (substInScope subst) x2
272 -- This uniqAway is actually needed. Here's the example:
273 -- rule: foldr (mapFB (:) f) [] = mapList
274 -- target: foldr (\x. mapFB k f x) []
276 -- k = \x. mapFB ... x
277 -- The first \x is ok, but when we inline k, hoping it might
278 -- match (:) we find a second \x.
280 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
281 = match e1 e2 tpl_vars case_kont subst
283 case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
286 match (Type ty1) (Type ty2) tpl_vars kont subst
287 = match_ty ty1 ty2 tpl_vars kont subst
289 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
291 = (match_ty to1 to2 tpl_vars $
292 match_ty from1 from2 tpl_vars $
293 match e1 e2 tpl_vars kont) subst
296 {- I don't buy this let-rule any more
297 The let rule fails on matching
298 forall f,x,xs. f (x:xs)
300 f (let y = e in (y:[]))
301 because we just get x->y, which is bogus.
303 -- This is an interesting rule: we simply ignore lets in the
304 -- term being matched against! The unfolding inside it is (by assumption)
305 -- already inside any occurrences of the bound variables, so we'll expand
306 -- them when we encounter them. Meanwhile, we can't get false matches because
307 -- (also by assumption) the term being matched has no shadowing.
308 match e1 (Let bind e2) tpl_vars kont subst
309 = match e1 e2 tpl_vars kont subst
312 -- Here is another important rule: if the term being matched is a
313 -- variable, we expand it so long as its unfolding is a WHNF
314 -- (Its occurrence information is not necessarily up to date,
315 -- so we don't use it.)
316 match e1 (Var v2) tpl_vars kont subst
317 | isCheapUnfolding unfolding
318 = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
320 unfolding = idUnfolding v2
323 -- We can't cope with lets in the template
325 match e1 e2 tpl_vars kont subst = match_fail
328 ------------------------------------------
329 match_alts [] [] tpl_vars kont subst
331 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
333 = bind vs1 vs2 (match r1 r2) tpl_vars
334 (match_alts alts1 alts2 tpl_vars kont)
336 match_alts alts1 alts2 tpl_vars kont subst = match_fail
338 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
340 ----------------------------------------
341 bind :: [CoreBndr] -- Template binders
342 -> [CoreBndr] -- Target binders
345 -- This makes uses of assumption (A) above. For example,
347 -- Template: (\x.y) (y is free)
348 -- Target : (\y.y) (y is bound)
349 -- We rename x to y in the template... but then erroneously
350 -- match y against y. But this can't happen because of (A)
351 bind vs1 vs2 matcher tpl_vars kont subst
352 = WARN( not (all not_in_subst vs1), bug_msg )
353 matcher tpl_vars kont' subst'
355 kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
356 subst' = bindSubstList subst vs1 vs2
358 -- The unBindSubst relies on no shadowing in the template
359 not_in_subst v = isNothing (lookupSubst subst v)
360 bug_msg = sep [ppr vs1, ppr vs2]
362 ----------------------------------------
363 matches [] [] tpl_vars kont subst
365 matches (e:es) (e':es') tpl_vars kont subst
366 = match e e' tpl_vars (matches es es' tpl_vars kont) subst
367 matches es es' tpl_vars kont subst
370 ----------------------------------------
371 mkVarArg :: CoreBndr -> CoreArg
372 mkVarArg v | isId v = Var v
373 | otherwise = Type (mkTyVarTy v)
376 Matching Core types: use the matcher in TcType.
377 Notice that we treat newtypes as opaque. For example, suppose
378 we have a specialised version of a function at a newtype, say
380 We only want to replace (f T) with f', not (f Int).
383 ----------------------------------------
384 match_ty ty1 ty2 tpl_vars kont subst
385 = TcType.match ty1 ty2 tpl_vars kont' (substEnv subst)
387 kont' senv = kont (setSubstEnv subst senv)
392 %************************************************************************
394 \subsection{Adding a new rule}
396 %************************************************************************
399 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
401 -- Insert the new rule just before a rule that is *less specific*
402 -- than the new one; or at the end if there isn't such a one.
403 -- In this way we make sure that when looking up, the first match
404 -- is the most specific.
406 -- We make no check for rules that unify without one dominating
407 -- the other. Arguably this would be a bug.
409 addRule (Rules rules rhs_fvs) id rule@(BuiltinRule _)
410 = Rules (rule:rules) rhs_fvs
411 -- Put it at the start for lack of anything better
413 addRule (Rules rules rhs_fvs) id rule
414 = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
416 new_rule = occurAnalyseRule rule
417 new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
419 -- Don't include the Id in its own rhs free-var set.
420 -- Otherwise the occurrence analyser makes bindings recursive
421 -- that shoudn't be. E.g.
422 -- RULE: f (f x y) z ==> f x (f y z)
424 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
427 tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
428 -- Actually we should probably include the free vars of tpl_args,
429 -- but I can't be bothered
432 go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
433 | otherwise = rule : go rules
435 new_is_more_specific rule = isJust (matchRule tpl_var_set rule tpl_args)
437 addIdSpecialisations :: Id -> [CoreRule] -> Id
438 addIdSpecialisations id rules
439 = setIdSpecialisation id new_specs
441 new_specs = foldr add (idSpecialisation id) rules
442 add rule rules = addRule rules id rule
446 %************************************************************************
448 \subsection{Preparing the rule base
450 %************************************************************************
453 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
454 lookupRule in_scope fn args
455 = case idSpecialisation fn of
456 Rules rules _ -> matchRules in_scope rules args
460 %************************************************************************
462 \subsection{Getting the rules ready}
464 %************************************************************************
467 data RuleBase = RuleBase
468 IdSet -- Ids with their rules in their specialisations
469 -- Held as a set, so that it can simply be the initial
470 -- in-scope set in the simplifier
472 IdSet -- Ids (whether local or imported) mentioned on
473 -- LHS of some rule; these should be black listed
475 -- This representation is a bit cute, and I wonder if we should
476 -- change it to use (IdEnv CoreRule) which seems a bit more natural
478 ruleBaseIds (RuleBase ids _) = ids
479 ruleBaseFVs (RuleBase _ fvs) = fvs
481 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
483 addRuleBaseFVs :: RuleBase -> IdSet -> RuleBase
484 addRuleBaseFVs (RuleBase rules fvs) extra_fvs
485 = RuleBase rules (fvs `unionVarSet` extra_fvs)
487 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
488 extendRuleBaseList rule_base new_guys
489 = foldl extendRuleBase rule_base new_guys
491 extendRuleBase :: RuleBase -> (Id,CoreRule) -> RuleBase
492 extendRuleBase (RuleBase rule_ids rule_fvs) (id, rule)
493 = RuleBase (extendVarSet rule_ids new_id)
494 (rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
496 new_id = setIdSpecialisation id (addRule old_rules id rule)
498 old_rules = idSpecialisation (fromMaybe id (lookupVarSet rule_ids id))
499 -- Get the old rules from rule_ids if the Id is already there, but
500 -- if not, use the Id from the incoming rule. If may be a PrimOpId,
501 -- in which case it may have rules in its belly already. Seems
502 -- dreadfully hackoid.
504 lhs_fvs = ruleLhsFreeIds rule
505 -- Finds *all* the free Ids of the LHS, not just
506 -- locally defined ones!!
508 pprRuleBase :: RuleBase -> SDoc
509 pprRuleBase (RuleBase rules _) = vcat [ pprCoreRule (ppr id) rs
510 | id <- varSetElems rules,
511 rs <- rulesRules $ idSpecialisation id ]