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         extendRuleBaseList,
  10         ruleBaseIds, pprRuleBase, ruleCheckProgram,
  11
  12         lookupRule, addRule, addIdSpecialisations
  13     ) where
  14
  15 #include "HsVersions.h"
  16
  17 import CoreSyn          -- All of it
  18 import OccurAnal        ( occurAnalyseRule )
  19 import CoreFVs          ( exprFreeVars, exprsFreeVars, ruleRhsFreeVars )
  20 import CoreUnfold       ( isCheapUnfolding, unfoldingTemplate )
  21 import CoreUtils        ( tcEqExprX )
  22 import CoreTidy         ( pprTidyIdRules )
  23 import Subst            ( IdSubstEnv, SubstResult(..) )
  24 import Id               ( Id, idUnfolding, idSpecialisation, setIdSpecialisation )
  25 import Var              ( Var )
  26 import VarSet
  27 import VarEnv
  28 import TcType           ( TvSubstEnv )
  29 import Unify            ( matchTyX, MatchEnv(..) )
  30 import BasicTypes       ( Activation, CompilerPhase, isActive )
  31
  32 import Outputable
  33 import FastString
  34 import Maybe            ( isJust, fromMaybe )
  35 import Util             ( sortLe )
  36 import Bag
  37 import List             ( isPrefixOf )
  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 :: (Activation -> Bool) -> InScopeSet
  84            -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
  85 -- See comments on matchRule
  86 matchRules is_active in_scope [] args = Nothing
  87 matchRules is_active in_scope (rule:rules) args
  88   = case matchRule is_active in_scope rule args of
  89         Just result -> Just result
  90         Nothing     -> matchRules is_active in_scope rules args
  91
  92 noBlackList :: Activation -> Bool
  93 noBlackList act = False         -- Nothing is black listed
  94
  95 matchRule :: (Activation -> Bool) -> InScopeSet
  96           -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
  97
  98 -- If (matchRule rule args) returns Just (name,rhs)
  99 -- then (f args) matches the rule, and the corresponding
 100 -- rewritten RHS is rhs
 101 --
 102 -- The bndrs and rhs is occurrence-analysed
 103 --
 104 --      Example
 105 --
 106 -- The rule
 107 --      forall f g x. map f (map g x) ==> map (f . g) x
 108 -- is stored
 109 --      CoreRule "map/map"
 110 --               [f,g,x]                -- tpl_vars
 111 --               [f,map g x]            -- tpl_args
 112 --               map (f.g) x)           -- rhs
 113 --
 114 -- Then the call: matchRule the_rule [e1,map e2 e3]
 115 --        = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
 116 --
 117 -- Any 'surplus' arguments in the input are simply put on the end
 118 -- of the output.
 119
 120 matchRule is_active in_scope rule@(BuiltinRule name match_fn) args
 121   = case match_fn args of
 122         Just expr -> Just (name,expr)
 123         Nothing   -> Nothing
 124
 125 matchRule is_active in_scope rule@(Rule rn act tpl_vars tpl_args rhs) args
 126   | not (is_active act)
 127   = Nothing
 128   | otherwise
 129   = case matchN in_scope tpl_vars tpl_args args of
 130         Just (tpl_vals, leftovers) -> Just (rn, mkLams tpl_vars rhs `mkApps` tpl_vals `mkApps` leftovers)
 131         Nothing                    -> Nothing
 132 \end{code}
 133
 134 \begin{code}
 135 matchN  :: InScopeSet
 136         -> [Var]                -- Template tyvars
 137         -> [CoreExpr]           -- Template
 138         -> [CoreExpr]           -- Target; can have more elts than template
 139         -> Maybe ([CoreExpr],   -- What is substituted for each template var
 140                   [CoreExpr])   -- Leftover target exprs
 141
 142 matchN in_scope tmpl_vars tmpl_es target_es
 143   = do  { (subst, leftover_es) <- go init_menv emptySubstEnv tmpl_es target_es
 144         ; return (map (lookup_tmpl subst) tmpl_vars, leftover_es) }
 145   where
 146     init_menv = ME { me_tmpls = mkVarSet tmpl_vars, me_env = init_rn_env }
 147     init_rn_env = mkRnEnv2 (extendInScopeSetList in_scope tmpl_vars)
 148
 149     go menv subst []     es     = Just (subst, es)
 150     go menv subst ts     []     = Nothing       -- Fail if too few actual args
 151     go menv subst (t:ts) (e:es) = do { subst1 <- match menv subst t e
 152                                      ; go menv subst1 ts es }
 153
 154     lookup_tmpl :: (TvSubstEnv, IdSubstEnv) -> Var -> CoreExpr
 155     lookup_tmpl (tv_subst, id_subst) tmpl_var
 156         | isTyVar tmpl_var = case lookupVarEnv tv_subst tmpl_var of
 157                                 Just ty         -> Type ty
 158                                 Nothing         -> unbound tmpl_var
 159         | otherwise        = case lookupVarEnv id_subst tmpl_var of
 160                                 Just (DoneEx e) -> e
 161                                 other           -> unbound tmpl_var
 162
 163     unbound var = pprPanic "Template variable unbound in rewrite rule" (ppr var)
 164
 165 emptySubstEnv :: (TvSubstEnv, IdSubstEnv)
 166 emptySubstEnv = (emptyVarEnv, emptyVarEnv)
 167
 168
 169 --      At one stage I tried to match even if there are more
 170 --      template args than real args.
 171
 172 --      I now think this is probably a bad idea.
 173 --      Should the template (map f xs) match (map g)?  I think not.
 174 --      For a start, in general eta expansion wastes work.
 175 --      SLPJ July 99
 176
 177
 178 match :: MatchEnv
 179       -> (TvSubstEnv, IdSubstEnv)
 180       -> CoreExpr               -- Template
 181       -> CoreExpr               -- Target
 182       -> Maybe (TvSubstEnv, IdSubstEnv)
 183
 184 -- See the notes with Unify.match, which matches types
 185 -- Everything is very similar for terms
 186
 187 -- Interesting examples:
 188 -- Consider matching
 189 --      \x->f      against    \f->f
 190 -- When we meet the lambdas we must remember to rename f to f' in the
 191 -- second expresion.  The RnEnv2 does that.
 192 --
 193 -- Consider matching
 194 --      forall a. \b->b    against   \a->3
 195 -- We must rename the \a.  Otherwise when we meet the lambdas we
 196 -- might substitute [a/b] in the template, and then erroneously
 197 -- succeed in matching what looks like the template variable 'a' against 3.
 198
 199 -- The Var case follows closely what happens in Unify.match
 200 match menv subst@(tv_subst, id_subst) (Var v1) e2
 201   | v1 `elemVarSet` me_tmpls menv
 202   = case lookupVarEnv id_subst v1' of
 203         Nothing | any (inRnEnvR rn_env) (varSetElems (exprFreeVars e2))
 204                 -> Nothing      -- Occurs check failure
 205                 -- e.g. match forall a. (\x-> a x) against (\y. y y)
 206
 207                 | otherwise
 208                 -> Just (tv_subst, extendVarEnv id_subst v1 (DoneEx e2))
 209
 210         Just (DoneEx e2') | tcEqExprX (nukeRnEnvL rn_env) e2' e2
 211                           -> Just subst
 212
 213         other -> Nothing
 214
 215   | otherwise   -- v1 is not a template variable
 216   = case e2 of
 217         Var v2 | v1' == rnOccR rn_env v2 -> Just subst
 218         other                            -> Nothing
 219   where
 220     rn_env = me_env menv
 221     v1'    = rnOccL rn_env v1
 222
 223 -- Here is another important rule: if the term being matched is a
 224 -- variable, we expand it so long as its unfolding is a WHNF
 225 -- (Its occurrence information is not necessarily up to date,
 226 --  so we don't use it.)
 227 match menv subst e1 (Var v2)
 228   | isCheapUnfolding unfolding
 229   = match menv subst e1 (unfoldingTemplate unfolding)
 230   where
 231     unfolding = idUnfolding v2
 232
 233 match menv subst (Lit lit1) (Lit lit2)
 234   | lit1 == lit2
 235   = Just subst
 236
 237 match menv subst (App f1 a1) (App f2 a2)
 238   = do  { subst' <- match menv subst f1 f2
 239         ; match menv subst' a1 a2 }
 240
 241 match menv subst (Lam x1 e1) (Lam x2 e2)
 242   = match menv' subst e1 e2
 243   where
 244     menv' = menv { me_env = rnBndr2 (me_env menv) x1 x2 }
 245
 246 -- This rule does eta expansion
 247 --              (\x.M)  ~  N    iff     M  ~  N x
 248 match menv subst (Lam x1 e1) e2
 249   = match menv' subst e1 (App e2 (varToCoreExpr new_x))
 250   where
 251     (rn_env', new_x) = rnBndrL (me_env menv) x1
 252     menv' = menv { me_env = rn_env' }
 253
 254 -- Eta expansion the other way
 255 --      M  ~  (\y.N)    iff   M y     ~  N
 256 match menv subst e1 (Lam x2 e2)
 257   = match menv' subst (App e1 (varToCoreExpr new_x)) e2
 258   where
 259     (rn_env', new_x) = rnBndrR (me_env menv) x2
 260     menv' = menv { me_env = rn_env' }
 261
 262 match menv subst (Case e1 x1 ty1 alts1) (Case e2 x2 ty2 alts2)
 263   = do  { subst1 <- match_ty menv subst ty1 ty2
 264         ; subst2 <- match menv subst1 e1 e2
 265         ; let menv' = menv { me_env = rnBndr2 (me_env menv) x2 x2 }
 266         ; match_alts menv' subst2 (sortLe le_alt alts1) (sortLe le_alt alts2)
 267         }
 268
 269 match menv subst (Type ty1) (Type ty2)
 270   = match_ty menv subst ty1 ty2
 271
 272 match menv subst (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
 273   = do  { subst1 <- match_ty menv subst  to1   to2
 274         ; subst2 <- match_ty menv subst1 from1 from2
 275         ; match menv subst2 e1 e2 }
 276
 277 -- This is an interesting rule: we simply ignore lets in the
 278 -- term being matched against!  The unfolding inside it is (by assumption)
 279 -- already inside any occurrences of the bound variables, so we'll expand
 280 -- them when we encounter them.
 281 match menv subst e1 (Let (NonRec x2 r2) e2)
 282   = match menv' subst e1 e2
 283   where
 284     menv' = menv { me_env = fst (rnBndrR (me_env menv) x2) }
 285         -- It's important to do this renaming. For example:
 286         -- Matching
 287         --      forall f,x,xs. f (x:xs)
 288         --   against
 289         --      f (let y = e in (y:[]))
 290         -- We must not get success with x->y!  Instead, we
 291         -- need an occurs check.
 292
 293 -- Everything else fails
 294 match menv subst e1 e2 = Nothing
 295
 296 ------------------------------------------
 297 match_alts :: MatchEnv
 298       -> (TvSubstEnv, IdSubstEnv)
 299       -> [CoreAlt]              -- Template
 300       -> [CoreAlt]              -- Target
 301       -> Maybe (TvSubstEnv, IdSubstEnv)
 302 match_alts menv subst [] []
 303   = return subst
 304 match_alts menv subst ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2)
 305   | c1 == c2
 306   = do  { subst1 <- match menv' subst r1 r2
 307         ; match_alts menv subst1 alts1 alts2 }
 308   where
 309     menv' :: MatchEnv
 310     menv' = menv { me_env = rnBndrs2 (me_env menv) vs1 vs2 }
 311
 312 match_alts menv subst alts1 alts2
 313   = Nothing
 314
 315 le_alt (con1, _, _) (con2, _, _) = con1 <= con2
 316 \end{code}
 317
 318 Matching Core types: use the matcher in TcType.
 319 Notice that we treat newtypes as opaque.  For example, suppose
 320 we have a specialised version of a function at a newtype, say
 321         newtype T = MkT Int
 322 We only want to replace (f T) with f', not (f Int).
 323
 324 \begin{code}
 325 ------------------------------------------
 326 match_ty menv (tv_subst, id_subst) ty1 ty2
 327   = do  { tv_subst' <- Unify.matchTyX menv tv_subst ty1 ty2
 328         ; return (tv_subst', id_subst) }
 329 \end{code}
 330
 331
 332 %************************************************************************
 333 %*                                                                      *
 334 \subsection{Adding a new rule}
 335 %*                                                                      *
 336 %************************************************************************
 337
 338 \begin{code}
 339 addRule :: Id -> CoreRules -> CoreRule -> CoreRules
 340
 341 -- Add a new rule to an existing bunch of rules.
 342 -- The rules are for the given Id; the Id argument is needed only
 343 -- so that we can exclude the Id from its own RHS free-var set
 344
 345 -- Insert the new rule just before a rule that is *less specific*
 346 -- than the new one; or at the end if there isn't such a one.
 347 -- In this way we make sure that when looking up, the first match
 348 -- is the most specific.
 349 --
 350 -- We make no check for rules that unify without one dominating
 351 -- the other.   Arguably this would be a bug.
 352
 353 addRule id (Rules rules rhs_fvs) rule@(BuiltinRule _ _)
 354   = Rules (rule:rules) rhs_fvs
 355         -- Put it at the start for lack of anything better
 356
 357 addRule id (Rules rules rhs_fvs) rule
 358   = Rules (insertRule rules new_rule) (rhs_fvs `unionVarSet` new_rhs_fvs)
 359   where
 360     new_rule    = occurAnalyseRule rule
 361     new_rhs_fvs = ruleRhsFreeVars new_rule `delVarSet` id
 362         -- Hack alert!
 363         -- Don't include the Id in its own rhs free-var set.
 364         -- Otherwise the occurrence analyser makes bindings recursive
 365         -- that shoudn't be.  E.g.
 366         --      RULE:  f (f x y) z  ==>  f x (f y z)
 367
 368 insertRule rules new_rule@(Rule _ _ tpl_vars tpl_args _)
 369   = go rules
 370   where
 371     tpl_var_set = mkInScopeSet (mkVarSet tpl_vars)
 372         -- Actually we should probably include the free vars of tpl_args,
 373         -- but I can't be bothered
 374
 375     go []                                       = [new_rule]
 376     go (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
 377                     | otherwise                 = rule : go rules
 378
 379     new_is_more_specific rule = isJust (matchRule noBlackList tpl_var_set rule tpl_args)
 380
 381 addIdSpecialisations :: Id -> [CoreRule] -> Id
 382 addIdSpecialisations id rules
 383   = setIdSpecialisation id new_specs
 384   where
 385     new_specs = foldl (addRule id) (idSpecialisation id) rules
 386 \end{code}
 387
 388
 389 %************************************************************************
 390 %*                                                                      *
 391 \subsection{Looking up a rule}
 392 %*                                                                      *
 393 %************************************************************************
 394
 395 \begin{code}
 396 lookupRule :: (Activation -> Bool) -> InScopeSet
 397            -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
 398 lookupRule is_active in_scope fn args
 399   = case idSpecialisation fn of
 400         Rules rules _ -> matchRules is_active in_scope rules args
 401 \end{code}
 402
 403
 404 %************************************************************************
 405 %*                                                                      *
 406 \subsection{Checking a program for failing rule applications}
 407 %*                                                                      *
 408 %************************************************************************
 409
 410 -----------------------------------------------------
 411                         Game plan
 412 -----------------------------------------------------
 413
 414 We want to know what sites have rules that could have fired but didn't.
 415 This pass runs over the tree (without changing it) and reports such.
 416
 417 NB: we assume that this follows a run of the simplifier, so every Id
 418 occurrence (including occurrences of imported Ids) is decorated with
 419 all its (active) rules.  No need to construct a rule base or anything
 420 like that.
 421
 422 \begin{code}
 423 ruleCheckProgram :: CompilerPhase -> String -> [CoreBind] -> SDoc
 424 -- Report partial matches for rules beginning
 425 -- with the specified string
 426 ruleCheckProgram phase rule_pat binds
 427   | isEmptyBag results
 428   = text "Rule check results: no rule application sites"
 429   | otherwise
 430   = vcat [text "Rule check results:",
 431           line,
 432           vcat [ p $$ line | p <- bagToList results ]
 433          ]
 434   where
 435     results = unionManyBags (map (ruleCheckBind (phase, rule_pat)) binds)
 436     line = text (replicate 20 '-')
 437
 438 type RuleCheckEnv = (CompilerPhase, String)     -- Phase and Pattern
 439
 440 ruleCheckBind :: RuleCheckEnv -> CoreBind -> Bag SDoc
 441    -- The Bag returned has one SDoc for each call site found
 442 ruleCheckBind env (NonRec b r) = ruleCheck env r
 443 ruleCheckBind env (Rec prs)    = unionManyBags [ruleCheck env r | (b,r) <- prs]
 444
 445 ruleCheck :: RuleCheckEnv -> CoreExpr -> Bag SDoc
 446 ruleCheck env (Var v)       = emptyBag
 447 ruleCheck env (Lit l)       = emptyBag
 448 ruleCheck env (Type ty)     = emptyBag
 449 ruleCheck env (App f a)     = ruleCheckApp env (App f a) []
 450 ruleCheck env (Note n e)    = ruleCheck env e
 451 ruleCheck env (Let bd e)    = ruleCheckBind env bd `unionBags` ruleCheck env e
 452 ruleCheck env (Lam b e)     = ruleCheck env e
 453 -- gaw 2004
 454 ruleCheck env (Case e _ _ as) = ruleCheck env e `unionBags`
 455                                 unionManyBags [ruleCheck env r | (_,_,r) <- as]
 456
 457 ruleCheckApp env (App f a) as = ruleCheck env a `unionBags` ruleCheckApp env f (a:as)
 458 ruleCheckApp env (Var f) as   = ruleCheckFun env f as
 459 ruleCheckApp env other as     = ruleCheck env other
 460 \end{code}
 461
 462 \begin{code}
 463 ruleCheckFun :: RuleCheckEnv -> Id -> [CoreExpr] -> Bag SDoc
 464 -- Produce a report for all rules matching the predicate
 465 -- saying why it doesn't match the specified application
 466
 467 ruleCheckFun (phase, pat) fn args
 468   | null name_match_rules = emptyBag
 469   | otherwise             = unitBag (ruleAppCheck_help phase fn args name_match_rules)
 470   where
 471     name_match_rules = case idSpecialisation fn of
 472                           Rules rules _ -> filter match rules
 473     match rule = pat `isPrefixOf` unpackFS (ruleName rule)
 474
 475 ruleAppCheck_help :: CompilerPhase -> Id -> [CoreExpr] -> [CoreRule] -> SDoc
 476 ruleAppCheck_help phase fn args rules
 477   =     -- The rules match the pattern, so we want to print something
 478     vcat [text "Expression:" <+> ppr (mkApps (Var fn) args),
 479           vcat (map check_rule rules)]
 480   where
 481     n_args = length args
 482     i_args = args `zip` [1::Int ..]
 483
 484     check_rule rule = rule_herald rule <> colon <+> rule_info rule
 485
 486     rule_herald (BuiltinRule name _) =
 487         ptext SLIT("Builtin rule") <+> doubleQuotes (ftext name)
 488     rule_herald (Rule name _ _ _ _)  =
 489         ptext SLIT("Rule") <+> doubleQuotes (ftext name)
 490
 491     rule_info rule
 492         | Just (name,_) <- matchRule noBlackList emptyInScopeSet rule args
 493         = text "matches (which is very peculiar!)"
 494
 495     rule_info (BuiltinRule name fn) = text "does not match"
 496
 497     rule_info (Rule name act rule_bndrs rule_args _)
 498         | not (isActive phase act)    = text "active only in later phase"
 499         | n_args < n_rule_args        = text "too few arguments"
 500         | n_mismatches == n_rule_args = text "no arguments match"
 501         | n_mismatches == 0           = text "all arguments match (considered individually), but rule as a whole does not"
 502         | otherwise                   = text "arguments" <+> ppr mismatches <+> text "do not match (1-indexing)"
 503         where
 504           n_rule_args  = length rule_args
 505           n_mismatches = length mismatches
 506           mismatches   = [i | (rule_arg, (arg,i)) <- rule_args `zip` i_args,
 507                               not (isJust (match_fn rule_arg arg))]
 508
 509           lhs_fvs = exprsFreeVars rule_args     -- Includes template tyvars
 510           match_fn rule_arg arg = match menv emptySubstEnv rule_arg arg
 511                 where
 512                   in_scope = lhs_fvs `unionVarSet` exprFreeVars arg
 513                   menv = ME { me_env   = mkRnEnv2 (mkInScopeSet in_scope)
 514                             , me_tmpls = mkVarSet rule_bndrs }
 515 \end{code}
 516
 517
 518 %************************************************************************
 519 %*                                                                      *
 520 \subsection{Getting the rules ready}
 521 %*                                                                      *
 522 %************************************************************************
 523
 524 \begin{code}
 525 data RuleBase = RuleBase
 526                     IdSet       -- Ids with their rules in their specialisations
 527                                 -- Held as a set, so that it can simply be the initial
 528                                 -- in-scope set in the simplifier
 529         -- This representation is a bit cute, and I wonder if we should
 530         -- change it to use (IdEnv CoreRule) which seems a bit more natural
 531
 532 ruleBaseIds (RuleBase ids) = ids
 533 emptyRuleBase = RuleBase emptyVarSet
 534
 535 extendRuleBaseList :: RuleBase -> [IdCoreRule] -> RuleBase
 536 extendRuleBaseList rule_base new_guys
 537   = foldl extendRuleBase rule_base new_guys
 538
 539 extendRuleBase :: RuleBase -> IdCoreRule -> RuleBase
 540 extendRuleBase (RuleBase rule_ids) (IdCoreRule id _ rule)
 541   = RuleBase (extendVarSet rule_ids new_id)
 542   where
 543     new_id    = setIdSpecialisation id (addRule id old_rules rule)
 544     old_rules = idSpecialisation (fromMaybe id (lookupVarSet rule_ids id))
 545         -- Get the old rules from rule_ids if the Id is already there, but
 546         -- if not, use the Id from the incoming rule.  If may be a PrimOpId,
 547         -- in which case it may have rules in its belly already.  Seems
 548         -- dreadfully hackoid.
 549
 550 pprRuleBase :: RuleBase -> SDoc
 551 pprRuleBase (RuleBase rules) = vcat [ pprTidyIdRules id | id <- varSetElems rules ]
 552 \end{code}