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