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