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