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, emptyRuleBase,
   9         extendRuleBase, extendRuleBaseList, addRuleBaseFVs,
  10         ruleBaseIds, ruleBaseFVs,
  11         pprRuleBase,
  12
  13         lookupRule, addRule, addIdSpecialisations
  14     ) where
  15
  16 #include "HsVersions.h"
  17
  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
  27                         )
  28 import Id               ( Id, idUnfolding, zapLamIdInfo,
  29                           idSpecialisation, setIdSpecialisation
  30                         )
  31 import Var              ( isTyVar, isId )
  32 import VarSet
  33 import VarEnv
  34 import Type             ( mkTyVarTy )
  35 import qualified Unify  ( match )
  36
  37 import Outputable
  38 import Maybes           ( maybeToBool )
  39 import Util             ( sortLt )
  40 \end{code}
  41
  42
  43 %************************************************************************
  44 %*                                                                      *
  45 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
  46 %*                                                                      *
  47 %************************************************************************
  48
  49 A @CoreRule@ holds details of one rule for an @Id@, which
  50 includes its specialisations.
  51
  52 For example, if a rule for @f@ contains the mapping:
  53 \begin{verbatim}
  54         forall a b d. [Type (List a), Type b, Var d]  ===>  f' a b
  55 \end{verbatim}
  56 then when we find an application of f to matching types, we simply replace
  57 it by the matching RHS:
  58 \begin{verbatim}
  59         f (List Int) Bool dict ===>  f' Int Bool
  60 \end{verbatim}
  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.
  64
  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:
  68
  69         pi :: forall a. Num a => a
  70
  71 might have a specialisation
  72
  73         [Int#] ===>  (case pi' of Lift pi# -> pi#)
  74
  75 where pi' :: Lift Int# is the specialised version of pi.
  76
  77
  78 %************************************************************************
  79 %*                                                                      *
  80 \subsection{Matching}
  81 %*                                                                      *
  82 %************************************************************************
  83
  84 \begin{code}
  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
  92
  93
  94 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
  95
  96 -- If (matchRule rule args) returns Just (name,rhs)
  97 -- then (f args) matches the rule, and the corresponding
  98 -- rewritten RHS is rhs
  99 --
 100 -- The bndrs and rhs is occurrence-analysed
 101 --
 102 --      Example
 103 --
 104 -- The rule
 105 --      forall f g x. map f (map g x) ==> map (f . g) x
 106 -- is stored
 107 --      CoreRule "map/map"
 108 --               [f,g,x]                -- tpl_vars
 109 --               [f,map g x]            -- tpl_args
 110 --               map (f.g) x)           -- rhs
 111 --
 112 -- Then the call: matchRule the_rule [e1,map e2 e3]
 113 --        = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
 114 --
 115 -- Any 'surplus' arguments in the input are simply put on the end
 116 -- of the output.
 117 --
 118 -- ASSUMPTION (A):
 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
 122 --
 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.
 127 --
 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.
 132 --
 133 -- A3 is needed to validate the rule that says
 134 --      (\x->E) matches F
 135 -- if
 136 --      (\x->E) matches (\x->F x)
 137
 138
 139 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
 140
 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
 145  where
 146    tpl_var_set = mkVarSet tpl_vars
 147
 148    -----------------------
 149         -- Do the business
 150    go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
 151
 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)
 155
 156         -- One tiresome way to terminate: check for excess unmatched
 157         -- template arguments
 158    go tpl_args []   subst = Nothing     -- Failure
 159
 160
 161    -----------------------
 162    app_match subst fn vs = foldl go fn vs
 163         where
 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!
 169
 170
 171    -----------------------
 172 {-      The code below tries to match even if there are more
 173         template args than real args.
 174
 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.
 178         SLPJ July 99
 179
 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
 184       where
 185         (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
 186                                       (map zapOccInfo tpl_vars)
 187                 -- Zap the occ info
 188         subst_env = substEnv subst
 189
 190    -----------------------
 191    eta_complete [] vars = ASSERT( isEmptyVarSet vars )
 192                           Just []
 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')
 198                                 Nothing    -> Nothing
 199                 other   -> Nothing
 200
 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')
 205                 Nothing    -> Nothing
 206
 207    eta_complete other vars = Nothing
 208
 209
 210 zapOccInfo bndr | isTyVar bndr = bndr
 211                 | otherwise    = zapLamIdInfo bndr
 212 -}
 213 \end{code}
 214
 215 \begin{code}
 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
 220
 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'
 225
 226 match :: CoreExpr               -- Template
 227       -> CoreExpr               -- Target
 228       -> Matcher result
 229
 230 match_fail = Nothing
 231
 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)
 238                    else
 239                          kont (extendSubst subst v1 (DoneEx e2))
 240
 241
 242                 | eqExpr (Var v1) e2             -> kont subst
 243                         -- v1 is not a template variable, so it must be a global constant
 244
 245         Just (DoneEx e2')  | eqExpr e2'       e2 -> kont subst
 246
 247         other -> match_fail
 248
 249 match (Lit lit1) (Lit lit2) tpl_vars kont subst
 250   | lit1 == lit2
 251   = kont subst
 252
 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
 255
 256 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
 257   = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
 258
 259 -- This rule does eta expansion
 260 --              (\x.M)  ~  N    iff     M  ~  N x
 261 -- See assumption A3
 262 match (Lam x1 e1) e2 tpl_vars kont subst
 263   = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
 264
 265 -- Eta expansion the other way
 266 --      M  ~  (\y.N)    iff   \y.M y  ~  \y.N
 267 --                      iff   M 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
 271   where
 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) []
 276         --            where
 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.
 280
 281 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
 282   = match e1 e2 tpl_vars case_kont subst
 283   where
 284     case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
 285                                      tpl_vars kont subst
 286
 287 match (Type ty1) (Type ty2) tpl_vars kont subst
 288   = match_ty ty1 ty2 tpl_vars kont subst
 289
 290 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
 291       tpl_vars kont subst
 292   = (match_ty to1   to2   tpl_vars $
 293      match_ty from1 from2 tpl_vars $
 294      match e1 e2 tpl_vars kont) subst
 295
 296
 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)
 300         against
 301                 f (let y = e in (y:[]))
 302         because we just get x->y, which is bogus.
 303
 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
 311 -}
 312
 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
 320   where
 321     unfolding = idUnfolding v2
 322
 323
 324 -- We can't cope with lets in the template
 325
 326 match e1 e2 tpl_vars kont subst = match_fail
 327
 328
 329 ------------------------------------------
 330 match_alts [] [] tpl_vars kont subst
 331   = kont subst
 332 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
 333   | c1 == c2
 334   = bind vs1 vs2 (match r1 r2) tpl_vars
 335                  (match_alts alts1 alts2 tpl_vars kont)
 336                  subst
 337 match_alts alts1 alts2 tpl_vars kont subst = match_fail
 338
 339 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
 340
 341 ----------------------------------------
 342 bind :: [CoreBndr]      -- Template binders
 343      -> [CoreBndr]      -- Target binders
 344      -> Matcher result
 345      -> Matcher result
 346 -- This makes uses of assumption (A) above.  For example,
 347 -- this would fail:
 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'
 355   where
 356     kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
 357     subst'        = bindSubstList subst vs1 vs2
 358
 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]
 362
 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')
 368
 369 ----------------------------------------
 370 matches [] [] tpl_vars kont subst
 371   = 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
 375   = match_fail
 376
 377 ----------------------------------------
 378 mkVarArg :: CoreBndr -> CoreArg
 379 mkVarArg v | isId v    = Var v
 380            | otherwise = Type (mkTyVarTy v)
 381 \end{code}
 382
 383 %************************************************************************
 384 %*                                                                      *
 385 \subsection{Adding a new rule}
 386 %*                                                                      *
 387 %************************************************************************
 388
 389 \begin{code}
 390 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
 391
 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.
 396 --
 397 -- We make no check for rules that unify without one dominating
 398 -- the other.   Arguably this would be a bug.
 399
 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
 403
 404 addRule (Rules rules rhs_fvs) id rule
 405   = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
 406   where
 407     new_rule    = occurAnalyseRule rule
 408     new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
 409         -- Hack alert!
 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)
 414
 415 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
 416   = go rules
 417   where
 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
 421
 422     go []                                       = [new_rule]
 423     go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
 424                     | otherwise                 = rule : go rules
 425
 426     new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
 427
 428 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
 429 addIdSpecialisations id spec_stuff
 430   = setIdSpecialisation id new_rules
 431   where
 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)
 435 \end{code}
 436
 437
 438 %************************************************************************
 439 %*                                                                      *
 440 \subsection{Preparing the rule base
 441 %*                                                                      *
 442 %************************************************************************
 443
 444 \begin{code}
 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
 449 \end{code}
 450
 451
 452 %************************************************************************
 453 %*                                                                      *
 454 \subsection{Getting the rules ready}
 455 %*                                                                      *
 456 %************************************************************************
 457
 458 \begin{code}
 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
 463
 464                      IdSet      -- Ids (whether local or imported) mentioned on
 465                                 -- LHS of some rule; these should be black listed
 466
 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
 469
 470 ruleBaseIds (RuleBase ids _) = ids
 471 ruleBaseFVs (RuleBase _ fvs) = fvs
 472
 473 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
 474
 475 addRuleBaseFVs :: RuleBase -> IdSet -> RuleBase
 476 addRuleBaseFVs (RuleBase rules fvs) extra_fvs
 477   = RuleBase rules (fvs `unionVarSet` extra_fvs)
 478
 479 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
 480 extendRuleBaseList rule_base new_guys
 481   = foldl extendRuleBase rule_base new_guys
 482
 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)
 487   where
 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'
 492
 493     lhs_fvs = ruleSomeLhsFreeVars isId rule
 494         -- Find *all* the free Ids of the LHS, not just
 495         -- locally defined ones!!
 496
 497 pprRuleBase :: RuleBase -> SDoc
 498 pprRuleBase (RuleBase rules _) = vcat [ pprCoreRule (ppr id) rs
 499                                       | id <- varSetElems rules,
 500                                         rs <- rulesRules $ idSpecialisation id ]
 501 \end{code}