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, 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
  27                         )
  28 import Id               ( Id, idUnfolding, idSpecialisation, setIdSpecialisation )
  29 import Var              ( isId )
  30 import VarSet
  31 import VarEnv
  32 import Type             ( mkTyVarTy )
  33 import qualified Unify  ( match )
  34
  35 import Outputable
  36 import Maybes           ( maybeToBool )
  37 import Util             ( sortLt )
  38 \end{code}
  39
  40
  41 %************************************************************************
  42 %*                                                                      *
  43 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
  44 %*                                                                      *
  45 %************************************************************************
  46
  47 A @CoreRule@ holds details of one rule for an @Id@, which
  48 includes its specialisations.
  49
  50 For example, if a rule for @f@ contains the mapping:
  51 \begin{verbatim}
  52         forall a b d. [Type (List a), Type b, Var d]  ===>  f' a b
  53 \end{verbatim}
  54 then when we find an application of f to matching types, we simply replace
  55 it by the matching RHS:
  56 \begin{verbatim}
  57         f (List Int) Bool dict ===>  f' Int Bool
  58 \end{verbatim}
  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.
  62
  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:
  66
  67         pi :: forall a. Num a => a
  68
  69 might have a specialisation
  70
  71         [Int#] ===>  (case pi' of Lift pi# -> pi#)
  72
  73 where pi' :: Lift Int# is the specialised version of pi.
  74
  75
  76 %************************************************************************
  77 %*                                                                      *
  78 \subsection{Matching}
  79 %*                                                                      *
  80 %************************************************************************
  81
  82 \begin{code}
  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
  90
  91
  92 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
  93
  94 -- If (matchRule rule args) returns Just (name,rhs)
  95 -- then (f args) matches the rule, and the corresponding
  96 -- rewritten RHS is rhs
  97 --
  98 -- The bndrs and rhs is occurrence-analysed
  99 --
 100 --      Example
 101 --
 102 -- The rule
 103 --      forall f g x. map f (map g x) ==> map (f . g) x
 104 -- is stored
 105 --      CoreRule "map/map"
 106 --               [f,g,x]                -- tpl_vars
 107 --               [f,map g x]            -- tpl_args
 108 --               map (f.g) x)           -- rhs
 109 --
 110 -- Then the call: matchRule the_rule [e1,map e2 e3]
 111 --        = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
 112 --
 113 -- Any 'surplus' arguments in the input are simply put on the end
 114 -- of the output.
 115 --
 116 -- ASSUMPTION (A):
 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
 120 --
 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.
 125 --
 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.
 130 --
 131 -- A3 is needed to validate the rule that says
 132 --      (\x->E) matches F
 133 -- if
 134 --      (\x->E) matches (\x->F x)
 135
 136
 137 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
 138
 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
 143  where
 144    tpl_var_set = mkVarSet tpl_vars
 145
 146    -----------------------
 147         -- Do the business
 148    go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
 149
 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)
 153
 154         -- One tiresome way to terminate: check for excess unmatched
 155         -- template arguments
 156    go tpl_args []   subst = Nothing     -- Failure
 157
 158
 159    -----------------------
 160    app_match subst fn vs = foldl go fn vs
 161         where
 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!
 167
 168
 169    -----------------------
 170 {-      The code below tries to match even if there are more
 171         template args than real args.
 172
 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.
 176         SLPJ July 99
 177
 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
 182       where
 183         (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
 184                                       (map zapOccInfo tpl_vars)
 185                 -- Zap the occ info
 186         subst_env = substEnv subst
 187
 188    -----------------------
 189    eta_complete [] vars = ASSERT( isEmptyVarSet vars )
 190                           Just []
 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')
 196                                 Nothing    -> Nothing
 197                 other   -> Nothing
 198
 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')
 203                 Nothing    -> Nothing
 204
 205    eta_complete other vars = Nothing
 206
 207
 208 zapOccInfo bndr | isTyVar bndr = bndr
 209                 | otherwise    = zapLamIdInfo bndr
 210 -}
 211 \end{code}
 212
 213 \begin{code}
 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
 218
 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'
 223
 224 match :: CoreExpr               -- Template
 225       -> CoreExpr               -- Target
 226       -> Matcher result
 227
 228 match_fail = Nothing
 229
 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)
 236                    else
 237                          kont (extendSubst subst v1 (DoneEx e2))
 238
 239
 240                 | eqExpr (Var v1) e2             -> kont subst
 241                         -- v1 is not a template variable, so it must be a global constant
 242
 243         Just (DoneEx e2')  | eqExpr e2'       e2 -> kont subst
 244
 245         other -> match_fail
 246
 247 match (Lit lit1) (Lit lit2) tpl_vars kont subst
 248   | lit1 == lit2
 249   = kont subst
 250
 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
 253
 254 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
 255   = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
 256
 257 -- This rule does eta expansion
 258 --              (\x.M)  ~  N    iff     M  ~  N x
 259 -- See assumption A3
 260 match (Lam x1 e1) e2 tpl_vars kont subst
 261   = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
 262
 263 -- Eta expansion the other way
 264 --      M  ~  (\y.N)    iff   \y.M y  ~  \y.N
 265 --                      iff   M 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
 269   where
 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) []
 274         --            where
 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.
 278
 279 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
 280   = match e1 e2 tpl_vars case_kont subst
 281   where
 282     case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
 283                                      tpl_vars kont subst
 284
 285 match (Type ty1) (Type ty2) tpl_vars kont subst
 286   = match_ty ty1 ty2 tpl_vars kont subst
 287
 288 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
 289       tpl_vars kont subst
 290   = (match_ty to1   to2   tpl_vars $
 291      match_ty from1 from2 tpl_vars $
 292      match e1 e2 tpl_vars kont) subst
 293
 294
 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)
 298         against
 299                 f (let y = e in (y:[]))
 300         because we just get x->y, which is bogus.
 301
 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
 309 -}
 310
 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
 318   where
 319     unfolding = idUnfolding v2
 320
 321
 322 -- We can't cope with lets in the template
 323
 324 match e1 e2 tpl_vars kont subst = match_fail
 325
 326
 327 ------------------------------------------
 328 match_alts [] [] tpl_vars kont subst
 329   = kont subst
 330 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
 331   | c1 == c2
 332   = bind vs1 vs2 (match r1 r2) tpl_vars
 333                  (match_alts alts1 alts2 tpl_vars kont)
 334                  subst
 335 match_alts alts1 alts2 tpl_vars kont subst = match_fail
 336
 337 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
 338
 339 ----------------------------------------
 340 bind :: [CoreBndr]      -- Template binders
 341      -> [CoreBndr]      -- Target binders
 342      -> Matcher result
 343      -> Matcher result
 344 -- This makes uses of assumption (A) above.  For example,
 345 -- this would fail:
 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'
 353   where
 354     kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
 355     subst'        = bindSubstList subst vs1 vs2
 356
 357         -- The unBindSubst relies on no shadowing in the template
 358     not_in_subst v = not (maybeToBool (lookupSubst subst v))
 359     bug_msg = sep [ppr vs1, ppr vs2]
 360
 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')
 366
 367 ----------------------------------------
 368 matches [] [] tpl_vars kont subst
 369   = 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
 373   = match_fail
 374
 375 ----------------------------------------
 376 mkVarArg :: CoreBndr -> CoreArg
 377 mkVarArg v | isId v    = Var v
 378            | otherwise = Type (mkTyVarTy v)
 379 \end{code}
 380
 381 %************************************************************************
 382 %*                                                                      *
 383 \subsection{Adding a new rule}
 384 %*                                                                      *
 385 %************************************************************************
 386
 387 \begin{code}
 388 addRule :: CoreRules -> Id -> CoreRule -> CoreRules
 389
 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.
 394 --
 395 -- We make no check for rules that unify without one dominating
 396 -- the other.   Arguably this would be a bug.
 397
 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
 401
 402 addRule (Rules rules rhs_fvs) id rule
 403   = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
 404   where
 405     new_rule    = occurAnalyseRule rule
 406     new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
 407         -- Hack alert!
 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)
 412
 413 insertRule rules new_rule@(Rule _ tpl_vars tpl_args _)
 414   = go rules
 415   where
 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
 419
 420     go []                                       = [new_rule]
 421     go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
 422                     | otherwise                 = rule : go rules
 423
 424     new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
 425
 426 addIdSpecialisations :: Id -> [CoreRule] -> Id
 427 addIdSpecialisations id rules
 428   = setIdSpecialisation id new_specs
 429   where
 430     new_specs = foldr add (idSpecialisation id) rules
 431     add rule rules = addRule rules id rule
 432 \end{code}
 433
 434
 435 %************************************************************************
 436 %*                                                                      *
 437 \subsection{Preparing the rule base
 438 %*                                                                      *
 439 %************************************************************************
 440
 441 \begin{code}
 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
 446 \end{code}
 447
 448
 449 %************************************************************************
 450 %*                                                                      *
 451 \subsection{Getting the rules ready}
 452 %*                                                                      *
 453 %************************************************************************
 454
 455 \begin{code}
 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
 460
 461                      IdSet      -- Ids (whether local or imported) mentioned on
 462                                 -- LHS of some rule; these should be black listed
 463
 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
 466
 467 ruleBaseIds (RuleBase ids _) = ids
 468 ruleBaseFVs (RuleBase _ fvs) = fvs
 469
 470 emptyRuleBase = RuleBase emptyVarSet emptyVarSet
 471
 472 addRuleBaseFVs :: RuleBase -> IdSet -> RuleBase
 473 addRuleBaseFVs (RuleBase rules fvs) extra_fvs
 474   = RuleBase rules (fvs `unionVarSet` extra_fvs)
 475
 476 extendRuleBaseList :: RuleBase -> [(Id,CoreRule)] -> RuleBase
 477 extendRuleBaseList rule_base new_guys
 478   = foldl extendRuleBase rule_base new_guys
 479
 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)
 484   where
 485     new_id = setIdSpecialisation id (addRule old_rules id rule)
 486     old_rules = case lookupVarSet rule_ids id of
 487                    Nothing  -> emptyCoreRules
 488                    Just id' -> idSpecialisation id'
 489
 490     lhs_fvs = ruleLhsFreeIds rule
 491         -- Finds *all* the free Ids of the LHS, not just
 492         -- locally defined ones!!
 493
 494 pprRuleBase :: RuleBase -> SDoc
 495 pprRuleBase (RuleBase rules _) = vcat [ pprCoreRule (ppr id) rs
 496                                       | id <- varSetElems rules,
 497                                         rs <- rulesRules $ idSpecialisation id ]
 498 \end{code}