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