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       ( Unfolding(..) )
  21 import CoreUtils        ( whnfOrBottom, eqExpr )
  22 import PprCore          ( pprCoreRule )
  23 import Subst            ( Subst, InScopeSet, substBndr, lookupSubst, extendSubst,
  24                           mkSubst, substEnv, setSubstEnv,
  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 --      No variable free in the template is bound in the target
 126
 127 matchRule in_scope (Rule rn tpl_vars tpl_args rhs) args
 128  = go tpl_args args (mkSubst in_scope emptySubstEnv)
 129  where
 130    tpl_var_set = mkVarSet tpl_vars
 131
 132    -----------------------
 133         -- Do the business
 134    go (tpl_arg:tpl_args) (arg:args) subst = match tpl_arg arg tpl_var_set (go tpl_args args) subst
 135
 136         -- Two easy ways to terminate
 137    go []                 []         subst = Just (rn, mkLams tpl_vars rhs, mk_result_args subst tpl_vars)
 138    go []                 args       subst = Just (rn, mkLams tpl_vars rhs, mk_result_args subst tpl_vars ++ args)
 139
 140         -- One tiresome way to terminate: check for excess unmatched
 141         -- template arguments
 142    go tpl_args           []         subst
 143       = case eta_complete tpl_args (mkVarSet leftovers) of
 144             Just leftovers' -> Just (rn, mkLams done (mkLams leftovers' rhs),
 145                                      mk_result_args subst done)
 146             Nothing         -> Nothing  -- Failure
 147       where
 148         (done, leftovers) = partition (\v -> maybeToBool (lookupSubstEnv subst_env v))
 149                                       (map zapOccInfo tpl_vars)
 150                 -- Zap the occ info
 151         subst_env = substEnv subst
 152
 153    -----------------------
 154    eta_complete [] vars = ASSERT( isEmptyVarSet vars )
 155                           Just []
 156    eta_complete (Type ty:tpl_args) vars
 157         = case getTyVar_maybe ty of
 158                 Just tv |  tv `elemVarSet` vars
 159                         -> case eta_complete tpl_args (vars `delVarSet` tv) of
 160                                 Just vars' -> Just (tv:vars')
 161                                 Nothing    -> Nothing
 162                 other   -> Nothing
 163
 164    eta_complete (Var v:tpl_args) vars
 165         | v `elemVarSet` vars
 166         = case eta_complete tpl_args (vars `delVarSet` v) of
 167                 Just vars' -> Just (v:vars')
 168                 Nothing    -> Nothing
 169
 170    eta_complete other vars = Nothing
 171
 172    -----------------------
 173    mk_result_args subst vs = map go vs
 174         where
 175           senv = substEnv subst
 176           go v = case lookupSubstEnv senv v of
 177                         Just (DoneEx ex) -> ex
 178                         Just (DoneTy ty) -> Type ty
 179                         -- Substitution should bind them all!
 180
 181 zapOccInfo bndr | isTyVar bndr = bndr
 182                 | otherwise    = maybeModifyIdInfo zapLamIdInfo bndr
 183 \end{code}
 184
 185 \begin{code}
 186 type Matcher result =  IdOrTyVarSet             -- Template variables
 187                     -> (Subst -> Maybe result)  -- Continuation if success
 188                     -> Subst  -> Maybe result   -- Substitution so far -> result
 189 -- The *SubstEnv* in these Substs apply to the TEMPLATE only
 190
 191 -- The *InScopeSet* in these Substs gives a superset of the free vars
 192 --      in the term being matched.  This set can get augmented, for example
 193 --      when matching against a lambda:
 194 --              (\x.M)  ~  N    iff     M  ~  N x
 195 --      but we must clone x if it's already free in N
 196
 197 match :: CoreExpr               -- Template
 198       -> CoreExpr               -- Target
 199       -> Matcher result
 200
 201 match_fail = Nothing
 202
 203 match (Var v1) e2 tpl_vars kont subst
 204   = case lookupSubst subst v1 of
 205         Nothing | v1 `elemVarSet` tpl_vars  -> kont (extendSubst subst v1 (DoneEx e2))
 206                         -- v1 is a template variables
 207
 208                 | eqExpr (Var v1) e2             -> kont subst
 209                         -- v1 is not a template variable, so it must be a global constant
 210
 211         Just (DoneEx e2')  | eqExpr e2'       e2 -> kont subst
 212
 213         other -> match_fail
 214
 215 match (Con c1 es1) (Con c2 es2) tpl_vars kont subst
 216   | c1 == c2
 217   = matches es1 es2 tpl_vars kont subst
 218
 219 match (App f1 a1) (App f2 a2) tpl_vars kont subst
 220   = match f1 f2 tpl_vars (match a1 a2 tpl_vars kont) subst
 221
 222 match (Lam x1 e1) (Lam x2 e2) tpl_vars kont subst
 223   = bind [x1] [x2] (match e1 e2) tpl_vars kont subst
 224
 225 -- This rule does eta expansion
 226 --              (\x.M)  ~  N    iff     M  ~  N x
 227 -- We must clone the binder in case it's already in scope in N
 228 match (Lam x1 e1) e2 tpl_vars kont subst
 229   = match e1 (App e2 (mkVarArg x1')) tpl_vars kont' subst'
 230   where
 231     (subst', x1') = substBndr subst x1
 232     kont' subst   = kont (unBindSubst subst x1 x1')
 233
 234 -- Eta expansion the other way
 235 --      M  ~  (\y.N)    iff   \y.M y  ~  \y.N
 236 --                      iff   M y     ~  N
 237 -- Remembering that by (A), y can't be free in M, we get this
 238 match e1 (Lam x2 e2) tpl_vars kont subst
 239   = match (App e1 (mkVarArg x2)) e2 tpl_vars kont subst
 240
 241 match (Case e1 x1 alts1) (Case e2 x2 alts2) tpl_vars kont subst
 242   = match e1 e2 tpl_vars case_kont subst
 243   where
 244     case_kont subst = bind [x1] [x2] (match_alts alts1 (sortLt lt_alt alts2))
 245                                      tpl_vars kont subst
 246
 247 match (Type ty1) (Type ty2) tpl_vars kont subst
 248   = match_ty ty1 ty2 tpl_vars kont subst
 249
 250 match (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
 251       tpl_vars kont subst
 252   = (match_ty to1   to2   tpl_vars $
 253      match_ty from1 from2 tpl_vars $
 254      match e1 e2 tpl_vars kont) subst
 255
 256
 257 {-      I don't buy this let-rule any more
 258         The let rule fails on matching
 259                 forall f,x,xs. f (x:xs)
 260         against
 261                 f (let y = e in (y:[]))
 262         because we just get x->y, which is bogus.
 263
 264 -- This is an interesting rule: we simply ignore lets in the
 265 -- term being matched against!  The unfolding inside it is (by assumption)
 266 -- already inside any occurrences of the bound variables, so we'll expand
 267 -- them when we encounter them.  Meanwhile, we can't get false matches because
 268 -- (also by assumption) the term being matched has no shadowing.
 269 match e1 (Let bind e2) tpl_vars kont subst
 270   = match e1 e2 tpl_vars kont subst
 271 -}
 272
 273 -- Here is another important rule: if the term being matched is a
 274 -- variable, we expand it so long as its unfolding is a WHNF
 275 -- (Its occurrence information is not necessarily up to date,
 276 --  so we don't use it.)
 277 match e1 (Var v2) tpl_vars kont subst
 278   = case getIdUnfolding v2 of
 279         CoreUnfolding form guidance unfolding
 280            |  whnfOrBottom form
 281            -> match e1 unfolding tpl_vars kont subst
 282
 283         other -> match_fail
 284
 285 -- We can't cope with lets in the template
 286
 287 match e1 e2 tpl_vars kont subst = match_fail
 288
 289
 290 ------------------------------------------
 291 match_alts [] [] tpl_vars kont subst
 292   = kont subst
 293 match_alts ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) tpl_vars kont subst
 294   | c1 == c2
 295   = bind vs1 vs2 (match r1 r2) tpl_vars
 296                  (match_alts alts1 alts2 tpl_vars kont)
 297                  subst
 298 match_alts alts1 alts2 tpl_vars kont subst = match_fail
 299
 300 lt_alt (con1, _, _) (con2, _, _) = con1 < con2
 301
 302 ----------------------------------------
 303 bind :: [CoreBndr]      -- Template binders
 304      -> [CoreBndr]      -- Target binders
 305      -> Matcher result
 306      -> Matcher result
 307 -- This makes uses of assumption (A) above.  For example,
 308 -- this would fail:
 309 --      Template: (\x.y)        (y is free)
 310 --      Target  : (\y.y)        (y is bound)
 311 -- We rename x to y in the template... but then erroneously
 312 -- match y against y.  But this can't happen because of (A)
 313 bind vs1 vs2 matcher tpl_vars kont subst
 314   = ASSERT( all not_in_subst vs1)
 315     matcher tpl_vars kont' subst'
 316   where
 317     kont' subst'' = kont (unBindSubstList subst'' vs1 vs2)
 318     subst'        = bindSubstList subst vs1 vs2
 319
 320         -- The unBindSubst relies on no shadowing in the template
 321     not_in_subst v = not (maybeToBool (lookupSubst subst v))
 322
 323 ----------------------------------------
 324 match_ty ty1 ty2 tpl_vars kont subst
 325   = case Unify.match ty1 ty2 tpl_vars Just (substEnv subst) of
 326         Nothing    -> match_fail
 327         Just senv' -> kont (setSubstEnv subst senv')
 328
 329 ----------------------------------------
 330 matches [] [] tpl_vars kont subst
 331   = kont subst
 332 matches (e:es) (e':es') tpl_vars kont subst
 333   = match e e' tpl_vars (matches es es' tpl_vars kont) subst
 334 matches es es' tpl_vars kont subst
 335   = match_fail
 336
 337 ----------------------------------------
 338 mkVarArg :: CoreBndr -> CoreArg
 339 mkVarArg v | isId v    = Var v
 340            | otherwise = Type (mkTyVarTy v)
 341 \end{code}
 342
 343 %************************************************************************
 344 %*                                                                      *
 345 \subsection{Adding a new rule}
 346 %*                                                                      *
 347 %************************************************************************
 348
 349 \begin{code}
 350 addRule :: Id -> CoreRules -> CoreRule -> CoreRules
 351
 352 -- Insert the new rule just before a rule that is *less specific*
 353 -- than the new one; or at the end if there isn't such a one.
 354 -- In this way we make sure that when looking up, the first match
 355 -- is the most specific.
 356 --
 357 -- We make no check for rules that unify without one dominating
 358 -- the other.   Arguably this would be a bug.
 359
 360 addRule id (Rules rules rhs_fvs) (Rule str tpl_vars tpl_args rhs)
 361   = Rules (insert rules) (rhs_fvs `unionVarSet` new_rhs_fvs)
 362   where
 363     new_rule = Rule str tpl_vars' tpl_args rhs'
 364                 -- Add occ info to tpl_vars, rhs
 365
 366     (rhs_uds, rhs')       = occurAnalyseExpr isLocallyDefined rhs
 367     (rhs_uds1, tpl_vars') = tagBinders rhs_uds tpl_vars
 368
 369     insert []                                       = [new_rule]
 370     insert (rule:rules) | new_is_more_specific rule = (new_rule:rule:rules)
 371                         | otherwise                 = rule : insert rules
 372
 373     new_is_more_specific rule = maybeToBool (matchRule tpl_var_set rule tpl_args)
 374
 375     tpl_var_set = mkVarSet tpl_vars'
 376         -- Actually we should probably include the free vars of tpl_args,
 377         -- but I can't be bothered
 378
 379     new_rhs_fvs = (exprFreeVars rhs' `minusVarSet` tpl_var_set) `delVarSet` id
 380         -- Hack alert!
 381         -- Don't include the Id in its own rhs free-var set.
 382         -- Otherwise the occurrence analyser makes bindings recursive
 383         -- that shoudn't be.  E.g.
 384         --      RULE:  f (f x y) z  ==>  f x (f y z)
 385
 386 addIdSpecialisations :: Id -> [([CoreBndr], [CoreExpr], CoreExpr)] -> Id
 387 addIdSpecialisations id spec_stuff
 388   = setIdSpecialisation id new_rules
 389   where
 390     rule_name = _PK_ ("SPEC " ++ showSDoc (ppr id))
 391     new_rules = foldr add (getIdSpecialisation id) spec_stuff
 392     add (vars, args, rhs) rules = addRule id rules (Rule rule_name vars args rhs)
 393 \end{code}
 394
 395
 396 %************************************************************************
 397 %*                                                                      *
 398 \subsection{Preparing the rule base
 399 %*                                                                      *
 400 %************************************************************************
 401
 402 \begin{code}
 403 data ProtoCoreRule
 404   = ProtoCoreRule
 405         Bool            -- True <=> this rule was defined in this module,
 406         Id              -- What Id is it for
 407         CoreRule        -- The rule itself
 408
 409
 410 pprProtoCoreRule (ProtoCoreRule _ fn rule) = pprCoreRule (Just fn) rule
 411
 412 lookupRule :: InScopeSet -> Id -> [CoreExpr] -> Maybe (FAST_STRING, CoreExpr, [CoreExpr])
 413 lookupRule in_scope fn args
 414   = case getIdSpecialisation fn of
 415         Rules rules _ -> matchRules in_scope rules args
 416
 417 orphanRule :: ProtoCoreRule -> Bool
 418 -- An "orphan rule" is one that is defined in this
 419 -- module, but of ran *imported* function.  We need
 420 -- to track these separately when generating the interface file
 421 orphanRule (ProtoCoreRule local fn _)
 422   = local && not (isLocallyDefined fn)
 423 \end{code}
 424
 425
 426 %************************************************************************
 427 %*                                                                      *
 428 \subsection{Getting the rules ready}
 429 %*                                                                      *
 430 %************************************************************************
 431
 432 \begin{code}
 433 type RuleBase = (IdSet,         -- Imported Ids that have rules attached
 434                  IdSet)         -- Ids (whether local or imported) mentioned on
 435                                 -- LHS of some rule; these should be black listed
 436
 437 -- The rule Ids and LHS Ids are black-listed; that is, they aren't inlined
 438 -- so that the opportunity to apply the rule isn't lost too soon
 439
 440 prepareRuleBase :: [CoreBind] -> [ProtoCoreRule] -> ([CoreBind], RuleBase)
 441 prepareRuleBase binds rules
 442   = (map zap_bind binds, (imported_rule_ids, rule_lhs_fvs))
 443   where
 444     (rule_ids, rule_lhs_fvs) = foldr add_rule (emptyVarSet, emptyVarSet) rules
 445     imported_rule_ids = filterVarSet (not . isLocallyDefined) rule_ids
 446
 447         -- rule_fvs is the set of all variables mentioned in rules
 448     rule_fvs = foldVarSet (unionVarSet . idRuleVars) rule_lhs_fvs rule_ids
 449
 450         -- Attach the rules for each locally-defined Id to that Id.
 451         --      - This makes the rules easier to look up
 452         --      - It means that transformation rules and specialisations for
 453         --        locally defined Ids are handled uniformly
 454         --      - It keeps alive things that are referred to only from a rule
 455         --        (the occurrence analyser knows about rules attached to Ids)
 456         --      - It makes sure that, when we apply a rule, the free vars
 457         --        of the RHS are more likely to be in scope
 458         --
 459         -- The LHS and RHS Ids are marked 'no-discard'.
 460         -- This means that the binding won't be discarded EVEN if the binding
 461         -- ends up being trivial (v = w) -- the simplifier would usually just
 462         -- substitute w for v throughout, but we don't apply the substitution to
 463         -- the rules (maybe we should?), so this substitution would make the rule
 464         -- bogus.
 465     zap_bind (NonRec b r) = NonRec (zap_bndr b) r
 466     zap_bind (Rec prs)    = Rec [(zap_bndr b, r) | (b,r) <- prs]
 467
 468     zap_bndr bndr = case lookupVarSet rule_ids bndr of
 469                           Just bndr'                           -> setIdNoDiscard bndr'
 470                           Nothing | bndr `elemVarSet` rule_fvs -> setIdNoDiscard bndr
 471                                   | otherwise                  -> bndr
 472
 473 add_rule (ProtoCoreRule _ id rule)
 474          (rule_id_set, rule_fvs)
 475   = (rule_id_set `extendVarSet` new_id,
 476      rule_fvs `unionVarSet` extendVarSet lhs_fvs id)
 477   where
 478     new_id = case lookupVarSet rule_id_set id of
 479                 Just id' -> addRuleToId id' rule
 480                 Nothing  -> addRuleToId id  rule
 481     lhs_fvs = ruleSomeLhsFreeVars isId rule
 482         -- Find *all* the free Ids of the LHS, not just
 483         -- locally defined ones!!
 484
 485 addRuleToId id rule = setIdSpecialisation id (addRule id (getIdSpecialisation id) rule)
 486 \end{code}
 487