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