[project @ 2000-07-07 09:37:39 by simonmar]
[ghc-hetmet.git] / ghc / compiler / specialise / Rules.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 \section[CoreRules]{Transformation rules}
5
6 \begin{code}
7 module Rules (
8         RuleBase, prepareLocalRuleBase, prepareOrphanRuleBase,
9         unionRuleBase, lookupRule, addRule, addIdSpecialisations,
10         ProtoCoreRule(..), pprProtoCoreRule, pprRuleBase,
11         localRule, orphanRule
12     ) where
13
14 #include "HsVersions.h"
15
16 import CoreSyn          -- All of it
17 import OccurAnal        ( occurAnalyseRule )
18 import BinderInfo       ( markMany )
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, substBndr, lookupSubst, extendSubst,
24                           mkSubst, substEnv, setSubstEnv, emptySubst, isInScope,
25                           unBindSubst, bindSubstList, unBindSubstList, substInScope
26                         )
27 import Id               ( Id, idUnfolding, zapLamIdInfo, 
28                           idSpecialisation, setIdSpecialisation,
29                           setIdNoDiscard, maybeModifyIdInfo, modifyIdInfo
30                         ) 
31 import Name             ( Name, isLocallyDefined )
32 import Var              ( isTyVar, isId )
33 import VarSet
34 import VarEnv
35 import Type             ( mkTyVarTy, getTyVar_maybe )
36 import qualified Unify  ( match )
37 import CmdLineOpts      ( opt_D_dump_simpl, opt_D_verbose_core2core )
38
39 import UniqFM
40 import ErrUtils         ( dumpIfSet )
41 import Outputable
42 import Maybes           ( maybeToBool )
43 import List             ( partition )
44 import Util             ( sortLt )
45 \end{code}
46
47
48 %************************************************************************
49 %*                                                                      *
50 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
51 %*                                                                      *
52 %************************************************************************
53
54 A @CoreRule@ holds details of one rule for an @Id@, which
55 includes its specialisations.
56
57 For example, if a rule for @f@ contains the mapping:
58 \begin{verbatim}
59         forall a b d. [Type (List a), Type b, Var d]  ===>  f' a b
60 \end{verbatim}
61 then when we find an application of f to matching types, we simply replace
62 it by the matching RHS:
63 \begin{verbatim}
64         f (List Int) Bool dict ===>  f' Int Bool
65 \end{verbatim}
66 All the stuff about how many dictionaries to discard, and what types
67 to apply the specialised function to, are handled by the fact that the
68 Rule contains a template for the result of the specialisation.
69
70 There is one more exciting case, which is dealt with in exactly the same
71 way.  If the specialised value is unboxed then it is lifted at its
72 definition site and unlifted at its uses.  For example:
73
74         pi :: forall a. Num a => a
75
76 might have a specialisation
77
78         [Int#] ===>  (case pi' of Lift pi# -> pi#)
79
80 where pi' :: Lift Int# is the specialised version of pi.
81
82
83 %************************************************************************
84 %*                                                                      *
85 \subsection{Matching}
86 %*                                                                      *
87 %************************************************************************
88
89 \begin{code}
90 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
91 -- See comments on matchRule
92 matchRules in_scope [] args = Nothing
93 matchRules in_scope (rule:rules) args
94   = case matchRule in_scope rule args of
95         Just result -> Just result
96         Nothing     -> matchRules in_scope rules args
97
98
99 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
100
101 -- If (matchRule rule args) returns Just (name,rhs)
102 -- then (f args) matches the rule, and the corresponding
103 -- rewritten RHS is rhs
104 --
105 -- The bndrs and rhs is occurrence-analysed
106 --
107 --      Example
108 --
109 -- The rule
110 --      forall f g x. map f (map g x) ==> map (f . g) x
111 -- is stored
112 --      CoreRule "map/map" 
113 --               [f,g,x]                -- tpl_vars
114 --               [f,map g x]            -- tpl_args
115 --               map (f.g) x)           -- rhs
116 --        
117 -- Then the call: matchRule the_rule [e1,map e2 e3]
118 --        = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
119 --
120 -- Any 'surplus' arguments in the input are simply put on the end
121 -- of the output.
122 --
123 -- ASSUMPTION (A):
124 --      A1. No top-level variable is bound in the target
125 --      A2. No template variable  is bound in the target
126 --      A3. No lambda bound template variable is free in any subexpression of the target
127 --
128 -- To see why A1 is necessary, consider matching
129 --      \x->f      against    \f->f
130 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
131 -- and then erroneously succeed in matching f against f.
132 --
133 -- To see why A2 is needed consider matching 
134 --      forall a. \b->b    against   \a->3
135 -- When we meet the lambdas we substitute [a/b] in the template, and then
136 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
137 --
138 -- A3 is needed to validate the rule that says
139 --      (\x->E) matches F
140 -- if
141 --      (\x->E) matches (\x->F x)
142
143
144 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
145
146 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
147   = go tpl_args args emptySubst
148         -- We used to use the in_scope set, but I don't think that's necessary
149         -- After all, the result is going to be simplified again with that in_scope set
150  where
151    tpl_var_set = mkVarSet tpl_vars
152
153    -----------------------
154         -- Do the business
155    go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
156
157         -- Two easy ways to terminate
158    go [] []         subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
159    go [] args       subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
160
161         -- One tiresome way to terminate: check for excess unmatched
162         -- template arguments
163    go tpl_args []   subst = Nothing     -- Failure
164
165
166    -----------------------
167    app_match subst fn vs = foldl go fn vs
168         where   
169           senv    = substEnv subst
170           go fn v = case lookupSubstEnv senv v of
171                         Just (DoneEx ex)  -> fn `App` ex 
172                         Just (DoneTy ty)  -> fn `App` Type ty
173                         -- Substitution should bind them all!
174
175
176    -----------------------
177 {-      The code below tries to match even if there are more 
178         template args than real args.
179
180         I now think this is probably a bad idea.
181         Should the template (map f xs) match (map g)?  I think not.
182         For a start, in general eta expansion wastes work.
183         SLPJ July 99
184
185       = case eta_complete tpl_args (mkVarSet leftovers) of
186             Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs), 
187                                      mk_result_args subst done)
188             Nothing         -> Nothing  -- Failure
189       where
190         (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
191                                       (map zapOccInfo tpl_vars)
192                 -- Zap the occ info 
193         subst_env = substEnv subst
194                                                 
195    -----------------------
196    eta_complete [] vars = ASSERT( isEmptyVarSet vars )
197                           Just []
198    eta_complete (Type ty:tpl_args) vars
199         = case getTyVar_maybe ty of
200                 Just tv |  tv `elemVarSet` vars
201                         -> case eta_complete tpl_args (vars `delVarSet` tv) of
202                                 Just vars' -> Just (tv:vars')
203                                 Nothing    -> Nothing
204                 other   -> Nothing
205
206    eta_complete (Var v:tpl_args) vars
207         | v `elemVarSet` vars
208         = case eta_complete tpl_args (vars `delVarSet` v) of
209                 Just vars' -> Just (v:vars')
210                 Nothing    -> Nothing
211
212    eta_complete other vars = Nothing
213 -}
214
215
216 zapOccInfo bndr | isTyVar bndr = bndr
217                 | otherwise    = zapLamIdInfo bndr
218 \end{code}
219
220 \begin{code}
221 type Matcher result =  VarSet                   -- Template variables
222                     -> (Subst -> Maybe result)  -- Continuation if success
223                     -> Subst  -> Maybe result   -- Substitution so far -> result
224 -- The *SubstEnv* in these Substs apply to the TEMPLATE only 
225
226 -- The *InScopeSet* in these Substs gives variables bound so far in the
227 --      target term.  So when matching forall a. (\x. a x) against (\y. y y)
228 --      while processing the body of the lambdas, the in-scope set will be {y}.
229 --      That lets us do the occurs-check when matching 'a' against 'y'
230
231 match :: CoreExpr               -- Template
232       -> CoreExpr               -- Target
233       -> Matcher result
234
235 match_fail = Nothing
236
237 match (Var v1) e2 tpl_vars kont subst
238   = case lookupSubst subst v1 of
239         Nothing | v1 `elemVarSet` tpl_vars      -- v1 is a template variable
240                 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
241                          match_fail             -- Occurs check failure
242                                                 -- e.g. match forall a. (\x-> a x) against (\y. y y)
243                    else
244                          kont (extendSubst subst v1 (DoneEx e2))
245
246
247                 | eqExpr (Var v1) e2             -> kont subst
248                         -- v1 is not a template variable, so it must be a global constant
249
250         Just (DoneEx e2')  | eqExpr e2'       e2 -> kont subst
251
252         other -> match_fail
253
254 match (Lit lit1) (Lit lit2) tpl_vars kont subst
255   | lit1 == lit2
256   = kont subst
257
258 match (App f1 a1) (App f2 a2) tpl_vars kont subst
259   = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
260
261 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
262   = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
263
264 -- This rule does eta expansion
265 --              (\x.M)  ~  N    iff     M  ~  N x
266 -- See assumption A3
267 match (Lam x1 e1) e2 tpl_vars kont subst
268   = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
269
270 -- Eta expansion the other way
271 --      M  ~  (\y.N)    iff   \y.M y  ~  \y.N
272 --                      iff   M y     ~  N
273 -- Remembering that by (A), y can't be free in M, we get this
274 match e1 (Lam x2 e2) tpl_vars kont subst
275   = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
276   where
277     new_id = uniqAway (substInScope subst) x2
278         -- This uniqAway is actually needed.  Here's the example:
279         --  rule:       foldr (mapFB (:) f) [] = mapList
280         --  target:     foldr (\x. mapFB k f x) []
281         --            where
282         --              k = \x. mapFB ... x
283         -- The first \x is ok, but when we inline k, hoping it might
284         -- match (:) we find a second \x.
285
286 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
287   = match e1 e2 tpl_vars case_kont subst
288   where
289     case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
290                                      tpl_vars kont subst
291
292 match (Type ty1) (Type ty2) tpl_vars kont subst
293   = match_ty ty1 ty2 tpl_vars kont subst
294
295 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
296       tpl_vars kont subst
297   = (match_ty to1   to2   tpl_vars $
298      match_ty from1 from2 tpl_vars $
299      match e1 e2 tpl_vars kont) subst
300
301
302 {-      I don't buy this let-rule any more
303         The let rule fails on matching
304                 forall f,x,xs. f (x:xs)
305         against
306                 f (let y = e in (y:[]))
307         because we just get x->y, which is bogus.
308
309 -- This is an interesting rule: we simply ignore lets in the 
310 -- term being matched against!  The unfolding inside it is (by assumption)
311 -- already inside any occurrences of the bound variables, so we'll expand
312 -- them when we encounter them.  Meanwhile, we can't get false matches because
313 -- (also by assumption) the term being matched has no shadowing.
314 match e1 (Let bind e2) tpl_vars kont subst
315   = match e1 e2 tpl_vars kont subst
316 -}
317
318 -- Here is another important rule: if the term being matched is a
319 -- variable, we expand it so long as its unfolding is a WHNF
320 -- (Its occurrence information is not necessarily up to date,
321 --  so we don't use it.)
322 match e1 (Var v2) tpl_vars kont subst
323   | isCheapUnfolding unfolding
324   = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
325   where
326     unfolding = idUnfolding v2
327
328
329 -- We can't cope with lets in the template
330
331 match e1 e2 tpl_vars kont subst = match_fail
332
333
334 ------------------------------------------
335 match_alts [] [] tpl_vars kont subst
336   = kont subst
337 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
338   | c1 == c2
339   = bind vs1 vs2 (match r1 r2) tpl_vars
340                  (match_alts alts1 alts2 tpl_vars kont)
341                  subst
342 match_alts alts1 alts2 tpl_vars kont subst = match_fail
343
344 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
345
346 ----------------------------------------
347 bind :: [CoreBndr]      -- Template binders
348      -> [CoreBndr]      -- Target binders
349      -> Matcher result
350      -> Matcher result
351 -- This makes uses of assumption (A) above.  For example,
352 -- this would fail:
353 --      Template: (\x.y)        (y is free)
354 --      Target  : (\y.y)        (y is bound)
355 -- We rename x to y in the template... but then erroneously
356 -- match y against y.  But this can't happen because of (A)
357 bind vs1 vs2 matcher tpl_vars kont subst
358   = WARN( not (all not_in_subst vs1), bug_msg )
359     matcher tpl_vars kont' subst'
360   where
361     kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
362     subst'        = bindSubstList subst vs1 vs2
363
364         -- The unBindSubst relies on no shadowing in the template
365     not_in_subst v = not (maybeToBool (lookupSubst subst v))
366     bug_msg = sep [ppr vs1, ppr vs2]
367
368 ----------------------------------------
369 match_ty ty1 ty2 tpl_vars kont subst
370   = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
371         Nothing    -> match_fail
372         Just senv' -> kont (setSubstEnv subst senv') 
373
374 ----------------------------------------
375 matches [] [] tpl_vars kont subst 
376   = kont subst
377 matches (e:es) (e':es') tpl_vars kont subst
378   = match e e' tpl_vars (matches es es' tpl_vars kont) subst
379 matches es es' tpl_vars kont subst 
380   = match_fail
381
382 ----------------------------------------
383 mkVarArg :: CoreBndr -> CoreArg
384 mkVarArg v | isId v    = Var v
385            | otherwise = Type (mkTyVarTy v)
386 \end{code}
387
388 %************************************************************************
389 %*                                                                      *
390 \subsection{Adding a new rule}
391 %*                                                                      *
392 %************************************************************************
393
394 \begin{code}
395 addRule :: Id -> CoreRules -> CoreRule -> CoreRules
396
397 -- Insert the new rule just before a rule that is *less specific*
398 -- than the new one; or at the end if there isn't such a one.
399 -- In this way we make sure that when looking up, the first match
400 -- is the most specific.
401 --
402 -- We make no check for rules that unify without one dominating
403 -- the other.   Arguably this would be a bug.
404
405 addRule id (Rules rules rhs_fvs) rule@(BuiltinRule _)
406   = Rules (rule:rules) rhs_fvs
407         -- Put it at the start for lack of anything better
408
409 addRule id (Rules rules rhs_fvs) rule
410   = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
411   where
412     new_rule    = occurAnalyseRule rule
413     new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
414         -- Hack alert!
415         -- Don't include the Id in its own rhs free-var set.
416         -- Otherwise the occurrence analyser makes bindings recursive
417         -- that shoudn't be.  E.g.
418         --      RULE:  f (f x y) z  ==>  f x (f y z)
419
420 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
421   = go rules
422   where
423     tpl_var_set = mkVarSet tpl_vars
424         -- Actually we should probably include the free vars of tpl_args,
425         -- but I can't be bothered
426
427     go []                                       = [new_rule]
428     go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
429                     | otherwise                 = rule : go rules
430
431     new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
432
433 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
434 addIdSpecialisations id spec_stuff
435   = setIdSpecialisation id new_rules
436   where
437     rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
438     new_rules = foldr add (idSpecialisation id) spec_stuff
439     add (vars, args, rhs) rules = addRule id rules (Rule rule_name vars args rhs)
440 \end{code}
441
442
443 %************************************************************************
444 %*                                                                      *
445 \subsection{Preparing the rule base
446 %*                                                                      *
447 %************************************************************************
448
449 \begin{code}
450 data ProtoCoreRule 
451   = ProtoCoreRule 
452         Bool            -- True <=> this rule was defined in this module,
453         Id              -- What Id is it for
454         CoreRule        -- The rule itself
455         
456
457 pprProtoCoreRule (ProtoCoreRule _ fn rule) = pprCoreRule (ppr fn) rule
458
459 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
460 lookupRule in_scope fn args
461   = case idSpecialisation fn of
462         Rules rules _ -> matchRules in_scope rules args
463
464 localRule :: ProtoCoreRule -> Bool
465 localRule (ProtoCoreRule local _ _) = local
466
467 orphanRule :: ProtoCoreRule -> Bool
468 -- An "orphan rule" is one that is defined in this 
469 -- module, but for an *imported* function.  We need
470 -- to track these separately when generating the interface file
471 orphanRule (ProtoCoreRule local fn _)
472   = local && not (isLocallyDefined fn)
473 \end{code}
474
475
476 %************************************************************************
477 %*                                                                      *
478 \subsection{Getting the rules ready}
479 %*                                                                      *
480 %************************************************************************
481
482 \begin{code}
483 type RuleBase = (IdSet,         -- Imported Ids that have rules attached
484                  IdSet)         -- Ids (whether local or imported) mentioned on 
485                                 -- LHS of some rule; these should be black listed
486
487 unionRuleBase (rule_ids1, black_ids1) (rule_ids2, black_ids2)
488   = (plusUFM_C merge_rules rule_ids1 rule_ids2,
489      unionVarSet black_ids1 black_ids2)
490   where
491     merge_rules id1 id2 = let rules1 = idSpecialisation id1
492                               rules2 = idSpecialisation id2
493                               new_rules = foldl (addRule id1) rules1 (rulesRules rules2)
494                           in
495                           setIdSpecialisation id1 new_rules
496
497 pprRuleBase :: RuleBase -> SDoc
498 pprRuleBase (rules,_) = vcat [ pprCoreRule (ppr id) rs
499                              | id <- varSetElems rules,
500                                rs <- rulesRules $ idSpecialisation id ]
501
502 -- prepareLocalRuleBase takes the CoreBinds and rules defined in this module.
503 -- It attaches those rules that are for local Ids to their binders, and
504 -- returns the remainder attached to Ids in an IdSet.  It also returns
505 -- Ids mentioned on LHS of some rule; these should be blacklisted.
506
507 -- The rule Ids and LHS Ids are black-listed; that is, they aren't inlined
508 -- so that the opportunity to apply the rule isn't lost too soon
509
510 prepareLocalRuleBase :: [CoreBind] -> [ProtoCoreRule] -> ([CoreBind], RuleBase)
511 prepareLocalRuleBase binds local_rules
512   = (map zap_bind binds, (imported_id_rule_ids, rule_lhs_fvs))
513   where
514     (rule_ids, rule_lhs_fvs) = foldr add_rule (emptyVarSet, emptyVarSet) local_rules
515     imported_id_rule_ids = filterVarSet (not . isLocallyDefined) rule_ids
516
517         -- rule_fvs is the set of all variables mentioned in this module's rules
518     rule_fvs = foldVarSet (unionVarSet . idRuleVars) rule_lhs_fvs rule_ids
519
520         -- Attach the rules for each locally-defined Id to that Id.
521         --      - This makes the rules easier to look up
522         --      - It means that transformation rules and specialisations for
523         --        locally defined Ids are handled uniformly
524         --      - It keeps alive things that are referred to only from a rule
525         --        (the occurrence analyser knows about rules attached to Ids)
526         --      - It makes sure that, when we apply a rule, the free vars
527         --        of the RHS are more likely to be in scope
528         --
529         -- The LHS and RHS Ids are marked 'no-discard'. 
530         -- This means that the binding won't be discarded EVEN if the binding
531         -- ends up being trivial (v = w) -- the simplifier would usually just 
532         -- substitute w for v throughout, but we don't apply the substitution to
533         -- the rules (maybe we should?), so this substitution would make the rule
534         -- bogus.
535     zap_bind (NonRec b r) = NonRec (zap_bndr b) r
536     zap_bind (Rec prs)    = Rec [(zap_bndr b, r) | (b,r) <- prs]
537
538     zap_bndr bndr = case lookupVarSet rule_ids bndr of
539                           Just bndr'                           -> setIdNoDiscard bndr'
540                           Nothing | bndr `elemVarSet` rule_fvs -> setIdNoDiscard bndr
541                                   | otherwise                  -> bndr
542                   
543 add_rule (ProtoCoreRule _ id rule)
544          (rule_id_set, rule_fvs)
545   = (rule_id_set `extendVarSet` new_id,
546      rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
547   where
548     new_id = case lookupVarSet rule_id_set id of
549                 Just id' -> addRuleToId id' rule
550                 Nothing  -> addRuleToId id  rule
551     lhs_fvs = ruleSomeLhsFreeVars isId rule
552         -- Find *all* the free Ids of the LHS, not just
553         -- locally defined ones!!
554
555 addRuleToId id rule = setIdSpecialisation id (addRule id (idSpecialisation id) rule)
556
557 -- prepareOrphanRuleBase does exactly the same as prepareLocalRuleBase, except that
558 -- it assumes that none of the rules can be attached to local Ids.
559
560 prepareOrphanRuleBase :: [ProtoCoreRule] -> RuleBase
561 prepareOrphanRuleBase imported_rules
562   = foldr add_rule (emptyVarSet, emptyVarSet) imported_rules
563 \end{code}