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, prepareRuleBase, lookupRule, addRule,
   9         addIdSpecialisations,
  10         ProtoCoreRule(..), pprProtoCoreRule,
  11         orphanRule
  12     ) where
  13
  14 #include "HsVersions.h"
  15
  16 import CoreSyn          -- All of it
  17 import Const            ( Con(..), Literal(..) )
  18 import OccurAnal        ( occurAnalyseExpr, tagBinders, UsageDetails )
  19 import BinderInfo       ( markMany )
  20 import CoreFVs          ( exprFreeVars, idRuleVars, ruleSomeLhsFreeVars )
  21 import CoreUnfold       ( isCheapUnfolding, unfoldingTemplate )
  22 import CoreUtils        ( eqExpr, cheapEqExpr )
  23 import PprCore          ( pprCoreRule )
  24 import Subst            ( Subst, InScopeSet, substBndr, lookupSubst, extendSubst,
  25                           mkSubst, substEnv, setSubstEnv, emptySubst, isInScope,
  26                           unBindSubst, bindSubstList, unBindSubstList, substInScope
  27                         )
  28 import Id               ( Id, getIdUnfolding, zapLamIdInfo,
  29                           getIdSpecialisation, setIdSpecialisation,
  30                           setIdNoDiscard, maybeModifyIdInfo, modifyIdInfo
  31                         )
  32 import IdInfo           ( setSpecInfo, specInfo )
  33 import Name             ( Name, isLocallyDefined )
  34 import Var              ( isTyVar, isId )
  35 import VarSet
  36 import VarEnv
  37 import Type             ( mkTyVarTy, getTyVar_maybe )
  38 import qualified Unify  ( match )
  39 import CmdLineOpts      ( opt_D_dump_simpl, opt_D_verbose_core2core )
  40
  41 import UniqFM
  42 import ErrUtils         ( dumpIfSet )
  43 import Outputable
  44 import Maybes           ( maybeToBool )
  45 import List             ( partition )
  46 import Util             ( sortLt )
  47 \end{code}
  48
  49
  50 %************************************************************************
  51 %*                                                                      *
  52 \subsection[specialisation-IdInfo]{Specialisation info about an @Id@}
  53 %*                                                                      *
  54 %************************************************************************
  55
  56 A @CoreRule@ holds details of one rule for an @Id@, which
  57 includes its specialisations.
  58
  59 For example, if a rule for @f@ contains the mapping:
  60 \begin{verbatim}
  61         forall a b d. [Type (List a), Type b, Var d]  ===>  f' a b
  62 \end{verbatim}
  63 then when we find an application of f to matching types, we simply replace
  64 it by the matching RHS:
  65 \begin{verbatim}
  66         f (List Int) Bool dict ===>  f' Int Bool
  67 \end{verbatim}
  68 All the stuff about how many dictionaries to discard, and what types
  69 to apply the specialised function to, are handled by the fact that the
  70 Rule contains a template for the result of the specialisation.
  71
  72 There is one more exciting case, which is dealt with in exactly the same
  73 way.  If the specialised value is unboxed then it is lifted at its
  74 definition site and unlifted at its uses.  For example:
  75
  76         pi :: forall a. Num a => a
  77
  78 might have a specialisation
  79
  80         [Int#] ===>  (case pi' of Lift pi# -> pi#)
  81
  82 where pi' :: Lift Int# is the specialised version of pi.
  83
  84
  85 %************************************************************************
  86 %*                                                                      *
  87 \subsection{Matching}
  88 %*                                                                      *
  89 %************************************************************************
  90
  91 \begin{code}
  92 matchRules :: InScopeSet -> [CoreRule] -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
  93 -- See comments on matchRule
  94 matchRules in_scope [] args = Nothing
  95 matchRules in_scope (rule:rules) args
  96   = case matchRule in_scope rule args of
  97         Just result -> Just result
  98         Nothing     -> matchRules in_scope rules args
  99
 100
 101 matchRule :: InScopeSet -> CoreRule -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
 102
 103 -- If (matchRule rule args) returns Just (name,rhs)
 104 -- then (f args) matches the rule, and the corresponding
 105 -- rewritten RHS is rhs
 106 --
 107 -- The bndrs and rhs is occurrence-analysed
 108 --
 109 --      Example
 110 --
 111 -- The rule
 112 --      forall f g x. map f (map g x) ==> map (f . g) x
 113 -- is stored
 114 --      CoreRule "map/map"
 115 --               [f,g,x]                -- tpl_vars
 116 --               [f,map g x]            -- tpl_args
 117 --               map (f.g) x)           -- rhs
 118 --
 119 -- Then the call: matchRule the_rule [e1,map e2 e3]
 120 --        = Just ("map/map", (\f,g,x -> rhs) e1 e2 e3)
 121 --
 122 -- Any 'surplus' arguments in the input are simply put on the end
 123 -- of the output.
 124 --
 125 -- ASSUMPTION (A):
 126 --      A1. No top-level variable is bound in the target
 127 --      A2. No template variable  is bound in the target
 128 --      A3. No lambda bound template variable is free in any subexpression of the target
 129 --
 130 -- To see why A1 is necessary, consider matching
 131 --      \x->f      against    \f->f
 132 -- When we meet the lambdas we substitute [f/x] in the template (a no-op),
 133 -- and then erroneously succeed in matching f against f.
 134 --
 135 -- To see why A2 is needed consider matching
 136 --      forall a. \b->b    against   \a->3
 137 -- When we meet the lambdas we substitute [a/b] in the template, and then
 138 -- erroneously succeed in matching what looks like the template variable 'a' against 3.
 139 --
 140 -- A3 is needed to validate the rule that says
 141 --      (\x->E) matches F
 142 -- if
 143 --      (\x->E) matches (\x->F x)
 144
 145
 146 matchRule in_scope rule@(BuiltinRule match_fn) args = match_fn args
 147
 148 matchRule in_scope rule@(Rule rn tpl_vars tpl_args rhs) args
 149   = go tpl_args args emptySubst
 150         -- We used to use the in_scope set, but I don't think that's necessary
 151         -- After all, the result is going to be simplified again with that in_scope set
 152  where
 153    tpl_var_set = mkVarSet tpl_vars
 154
 155    -----------------------
 156         -- Do the business
 157    go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
 158
 159         -- Two easy ways to terminate
 160    go [] []         subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars)
 161    go [] args       subst = Just (rn, app_match subst (mkLams tpl_vars rhs) tpl_vars `mkApps` args)
 162
 163         -- One tiresome way to terminate: check for excess unmatched
 164         -- template arguments
 165    go tpl_args []   subst = Nothing     -- Failure
 166
 167
 168    -----------------------
 169    app_match subst fn vs = foldl go fn vs
 170         where
 171           senv    = substEnv subst
 172           go fn v = case lookupSubstEnv senv v of
 173                         Just (DoneEx ex)  -> fn `App` ex
 174                         Just (DoneTy ty)  -> fn `App` Type ty
 175                         -- Substitution should bind them all!
 176
 177
 178    -----------------------
 179 {-      The code below tries to match even if there are more
 180         template args than real args.
 181
 182         I now think this is probably a bad idea.
 183         Should the template (map f xs) match (map g)?  I think not.
 184         For a start, in general eta expansion wastes work.
 185         SLPJ July 99
 186
 187       = case eta_complete tpl_args (mkVarSet leftovers) of
 188             Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
 189                                      mk_result_args subst done)
 190             Nothing         -> Nothing  -- Failure
 191       where
 192         (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
 193                                       (map zapOccInfo tpl_vars)
 194                 -- Zap the occ info
 195         subst_env = substEnv subst
 196
 197    -----------------------
 198    eta_complete [] vars = ASSERT( isEmptyVarSet vars )
 199                           Just []
 200    eta_complete (Type ty:tpl_args) vars
 201         = case getTyVar_maybe ty of
 202                 Just tv |  tv `elemVarSet` vars
 203                         -> case eta_complete tpl_args (vars `delVarSet` tv) of
 204                                 Just vars' -> Just (tv:vars')
 205                                 Nothing    -> Nothing
 206                 other   -> Nothing
 207
 208    eta_complete (Var v:tpl_args) vars
 209         | v `elemVarSet` vars
 210         = case eta_complete tpl_args (vars `delVarSet` v) of
 211                 Just vars' -> Just (v:vars')
 212                 Nothing    -> Nothing
 213
 214    eta_complete other vars = Nothing
 215 -}
 216
 217
 218 zapOccInfo bndr | isTyVar bndr = bndr
 219                 | otherwise    = zapLamIdInfo bndr
 220 \end{code}
 221
 222 \begin{code}
 223 type Matcher result =  IdOrTyVarSet             -- Template variables
 224                     -> (Subst -> Maybe result)  -- Continuation if success
 225                     -> Subst  -> Maybe result   -- Substitution so far -> result
 226 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
 227
 228 -- The *InScopeSet* in these Substs gives variables bound so far in the
 229 --      target term.  So when matching forall a. (\x. a x) against (\y. y y)
 230 --      while processing the body of the lambdas, the in-scope set will be {y}.
 231 --      That lets us do the occurs-check when matching 'a' against 'y'
 232
 233 match :: CoreExpr               -- Template
 234       -> CoreExpr               -- Target
 235       -> Matcher result
 236
 237 match_fail = Nothing
 238
 239 match (Var v1) e2 tpl_vars kont subst
 240   = case lookupSubst subst v1 of
 241         Nothing | v1 `elemVarSet` tpl_vars      -- v1 is a template variable
 242                 -> if (any (`isInScope` subst) (varSetElems (exprFreeVars e2))) then
 243                          match_fail             -- Occurs check failure
 244                                                 -- e.g. match forall a. (\x-> a x) against (\y. y y)
 245                    else
 246                          kont (extendSubst subst v1 (DoneEx e2))
 247
 248
 249                 | eqExpr (Var v1) e2             -> kont subst
 250                         -- v1 is not a template variable, so it must be a global constant
 251
 252         Just (DoneEx e2')  | eqExpr e2'       e2 -> kont subst
 253
 254         other -> match_fail
 255
 256 match (Con c1 es1) (Con c2 es2) tpl_vars kont subst
 257   | c1 == c2
 258   = matches es1 es2 tpl_vars kont subst
 259
 260 match (App f1 a1) (App f2 a2) tpl_vars kont subst
 261   = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
 262
 263 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
 264   = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
 265
 266 -- This rule does eta expansion
 267 --              (\x.M)  ~  N    iff     M  ~  N x
 268 -- See assumption A3
 269 match (Lam x1 e1) e2 tpl_vars kont subst
 270   = bind [x1] [x1] (match e1 (App e2 (mkVarArg x1))) tpl_vars kont subst
 271
 272 -- Eta expansion the other way
 273 --      M  ~  (\y.N)    iff   \y.M y  ~  \y.N
 274 --                      iff   M y     ~  N
 275 -- Remembering that by (A), y can't be free in M, we get this
 276 match e1 (Lam x2 e2) tpl_vars kont subst
 277   = bind [new_id] [x2] (match (App e1 (mkVarArg new_id)) e2) tpl_vars kont subst
 278   where
 279     new_id = uniqAway (substInScope subst) x2
 280         -- This uniqAway is actually needed.  Here's the example:
 281         --  rule:       foldr (mapFB (:) f) [] = mapList
 282         --  target:     foldr (\x. mapFB k f x) []
 283         --            where
 284         --              k = \x. mapFB ... x
 285         -- The first \x is ok, but when we inline k, hoping it might
 286         -- match (:) we find a second \x.
 287
 288 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
 289   = match e1 e2 tpl_vars case_kont subst
 290   where
 291     case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
 292                                      tpl_vars kont subst
 293
 294 match (Type ty1) (Type ty2) tpl_vars kont subst
 295   = match_ty ty1 ty2 tpl_vars kont subst
 296
 297 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
 298       tpl_vars kont subst
 299   = (match_ty to1   to2   tpl_vars $
 300      match_ty from1 from2 tpl_vars $
 301      match e1 e2 tpl_vars kont) subst
 302
 303
 304 {-      I don't buy this let-rule any more
 305         The let rule fails on matching
 306                 forall f,x,xs. f (x:xs)
 307         against
 308                 f (let y = e in (y:[]))
 309         because we just get x->y, which is bogus.
 310
 311 -- This is an interesting rule: we simply ignore lets in the
 312 -- term being matched against!  The unfolding inside it is (by assumption)
 313 -- already inside any occurrences of the bound variables, so we'll expand
 314 -- them when we encounter them.  Meanwhile, we can't get false matches because
 315 -- (also by assumption) the term being matched has no shadowing.
 316 match e1 (Let bind e2) tpl_vars kont subst
 317   = match e1 e2 tpl_vars kont subst
 318 -}
 319
 320 -- Here is another important rule: if the term being matched is a
 321 -- variable, we expand it so long as its unfolding is a WHNF
 322 -- (Its occurrence information is not necessarily up to date,
 323 --  so we don't use it.)
 324 match e1 (Var v2) tpl_vars kont subst
 325   | isCheapUnfolding unfolding
 326   = match e1 (unfoldingTemplate unfolding) tpl_vars kont subst
 327   where
 328     unfolding = getIdUnfolding v2
 329
 330
 331 -- We can't cope with lets in the template
 332
 333 match e1 e2 tpl_vars kont subst = match_fail
 334
 335
 336 ------------------------------------------
 337 match_alts [] [] tpl_vars kont subst
 338   = kont subst
 339 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
 340   | c1 == c2
 341   = bind vs1 vs2 (match r1 r2) tpl_vars
 342                  (match_alts alts1 alts2 tpl_vars kont)
 343                  subst
 344 match_alts alts1 alts2 tpl_vars kont subst = match_fail
 345
 346 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
 347
 348 ----------------------------------------
 349 bind :: [CoreBndr]      -- Template binders
 350      -> [CoreBndr]      -- Target binders
 351      -> Matcher result
 352      -> Matcher result
 353 -- This makes uses of assumption (A) above.  For example,
 354 -- this would fail:
 355 --      Template: (\x.y)        (y is free)
 356 --      Target  : (\y.y)        (y is bound)
 357 -- We rename x to y in the template... but then erroneously
 358 -- match y against y.  But this can't happen because of (A)
 359 bind vs1 vs2 matcher tpl_vars kont subst
 360   = WARN( not (all not_in_subst vs1), bug_msg )
 361     matcher tpl_vars kont' subst'
 362   where
 363     kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
 364     subst'        = bindSubstList subst vs1 vs2
 365
 366         -- The unBindSubst relies on no shadowing in the template
 367     not_in_subst v = not (maybeToBool (lookupSubst subst v))
 368     bug_msg = sep [ppr vs1, ppr vs2]
 369
 370 ----------------------------------------
 371 match_ty ty1 ty2 tpl_vars kont subst
 372   = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
 373         Nothing    -> match_fail
 374         Just senv' -> kont (setSubstEnv subst senv')
 375
 376 ----------------------------------------
 377 matches [] [] tpl_vars kont subst
 378   = kont subst
 379 matches (e:es) (e':es') tpl_vars kont subst
 380   = match e e' tpl_vars (matches es es' tpl_vars kont) subst
 381 matches es es' tpl_vars kont subst
 382   = match_fail
 383
 384 ----------------------------------------
 385 mkVarArg :: CoreBndr -> CoreArg
 386 mkVarArg v | isId v    = Var v
 387            | otherwise = Type (mkTyVarTy v)
 388 \end{code}
 389
 390 %************************************************************************
 391 %*                                                                      *
 392 \subsection{Adding a new rule}
 393 %*                                                                      *
 394 %************************************************************************
 395
 396 \begin{code}
 397 addRule :: Id -> CoreRules -> CoreRule -> CoreRules
 398
 399 -- Insert the new rule just before a rule that is *less specific*
 400 -- than the new one; or at the end if there isn't such a one.
 401 -- In this way we make sure that when looking up, the first match
 402 -- is the most specific.
 403 --
 404 -- We make no check for rules that unify without one dominating
 405 -- the other.   Arguably this would be a bug.
 406
 407 addRule id (Rules rules rhs_fvs) rule@(BuiltinRule _)
 408   = Rules (rule:rules) rhs_fvs
 409         -- Put it at the start for lack of anything better
 410
 411 addRule id (Rules rules rhs_fvs) (Rule str tpl_vars tpl_args rhs)
 412   = Rules (insert rules) (rhs_fvs `unionVarSet` new_rhs_fvs)
 413   where
 414     new_rule = Rule str tpl_vars' tpl_args rhs'
 415                 -- Add occ info to tpl_vars, rhs
 416
 417     (rhs_uds, rhs')       = occurAnalyseExpr isLocallyDefined rhs
 418     (rhs_uds1, tpl_vars') = tagBinders rhs_uds tpl_vars
 419
 420     insert []                                       = [new_rule]
 421     insert (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
 422                         | otherwise                 = rule : insert rules
 423
 424     new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
 425
 426     tpl_var_set = mkVarSet tpl_vars'
 427         -- Actually we should probably include the free vars of tpl_args,
 428         -- but I can't be bothered
 429
 430     new_rhs_fvs = (exprFreeVars rhs' `minusVarSet` tpl_var_set) `delVarSet` id
 431         -- Hack alert!
 432         -- Don't include the Id in its own rhs free-var set.
 433         -- Otherwise the occurrence analyser makes bindings recursive
 434         -- that shoudn't be.  E.g.
 435         --      RULE:  f (f x y) z  ==>  f x (f y z)
 436
 437 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
 438 addIdSpecialisations id spec_stuff
 439   = setIdSpecialisation id new_rules
 440   where
 441     rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
 442     new_rules = foldr add (getIdSpecialisation id) spec_stuff
 443     add (vars, args, rhs) rules = addRule id rules (Rule rule_name vars args rhs)
 444 \end{code}
 445
 446
 447 %************************************************************************
 448 %*                                                                      *
 449 \subsection{Preparing the rule base
 450 %*                                                                      *
 451 %************************************************************************
 452
 453 \begin{code}
 454 data ProtoCoreRule
 455   = ProtoCoreRule
 456         Bool            -- True <=> this rule was defined in this module,
 457         Id              -- What Id is it for
 458         CoreRule        -- The rule itself
 459
 460
 461 pprProtoCoreRule (ProtoCoreRule _ fn rule) = pprCoreRule (Just fn) rule
 462
 463 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (RuleName, CoreExpr)
 464 lookupRule in_scope fn args
 465   = case getIdSpecialisation fn of
 466         Rules rules _ -> matchRules in_scope rules args
 467
 468 orphanRule :: ProtoCoreRule -> Bool
 469 -- An "orphan rule" is one that is defined in this
 470 -- module, but of ran *imported* function.  We need
 471 -- to track these separately when generating the interface file
 472 orphanRule (ProtoCoreRule local fn _)
 473   = local && not (isLocallyDefined fn)
 474 \end{code}
 475
 476
 477 %************************************************************************
 478 %*                                                                      *
 479 \subsection{Getting the rules ready}
 480 %*                                                                      *
 481 %************************************************************************
 482
 483 \begin{code}
 484 type RuleBase = (IdSet,         -- Imported Ids that have rules attached
 485                  IdSet)         -- Ids (whether local or imported) mentioned on
 486                                 -- LHS of some rule; these should be black listed
 487
 488 -- The rule Ids and LHS Ids are black-listed; that is, they aren't inlined
 489 -- so that the opportunity to apply the rule isn't lost too soon
 490
 491 prepareRuleBase :: [CoreBind] -> [ProtoCoreRule] -> ([CoreBind], RuleBase)
 492 prepareRuleBase binds all_rules
 493   = (map zap_bind binds, (imported_rule_ids, rule_lhs_fvs))
 494   where
 495     (rule_ids, rule_lhs_fvs) = foldr add_rule (emptyVarSet, emptyVarSet) all_rules
 496     imported_rule_ids = filterVarSet (not . isLocallyDefined) rule_ids
 497
 498         -- rule_fvs is the set of all variables mentioned in rules
 499     rule_fvs = foldVarSet (unionVarSet . idRuleVars) rule_lhs_fvs rule_ids
 500
 501         -- Attach the rules for each locally-defined Id to that Id.
 502         --      - This makes the rules easier to look up
 503         --      - It means that transformation rules and specialisations for
 504         --        locally defined Ids are handled uniformly
 505         --      - It keeps alive things that are referred to only from a rule
 506         --        (the occurrence analyser knows about rules attached to Ids)
 507         --      - It makes sure that, when we apply a rule, the free vars
 508         --        of the RHS are more likely to be in scope
 509         --
 510         -- The LHS and RHS Ids are marked 'no-discard'.
 511         -- This means that the binding won't be discarded EVEN if the binding
 512         -- ends up being trivial (v = w) -- the simplifier would usually just
 513         -- substitute w for v throughout, but we don't apply the substitution to
 514         -- the rules (maybe we should?), so this substitution would make the rule
 515         -- bogus.
 516     zap_bind (NonRec b r) = NonRec (zap_bndr b) r
 517     zap_bind (Rec prs)    = Rec [(zap_bndr b, r) | (b,r) <- prs]
 518
 519     zap_bndr bndr = case lookupVarSet rule_ids bndr of
 520                           Just bndr'                           -> setIdNoDiscard bndr'
 521                           Nothing | bndr `elemVarSet` rule_fvs -> setIdNoDiscard bndr
 522                                   | otherwise                  -> bndr
 523
 524 add_rule (ProtoCoreRule _ id rule)
 525          (rule_id_set, rule_fvs)
 526   = (rule_id_set `extendVarSet` new_id,
 527      rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
 528   where
 529     new_id = case lookupVarSet rule_id_set id of
 530                 Just id' -> addRuleToId id' rule
 531                 Nothing  -> addRuleToId id  rule
 532     lhs_fvs = ruleSomeLhsFreeVars isId rule
 533         -- Find *all* the free Ids of the LHS, not just
 534         -- locally defined ones!!
 535
 536 addRuleToId id rule = setIdSpecialisation id (addRule id (getIdSpecialisation id) rule)
 537 \end{code}