compiler/coreSyn/CoreSubst.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   4 %
   5
   6 Utility functions on @Core@ syntax
   7
   8 \begin{code}
   9 module CoreSubst (
  10         -- * Main data types
  11         Subst, TvSubstEnv, IdSubstEnv, InScopeSet,
  12
  13         -- ** Substituting into expressions and related types
  14         deShadowBinds, substSpec, substRulesForImportedIds,
  15         substTy, substExpr, substExprSC, substBind, substBindSC,
  16         substUnfolding, substUnfoldingSC,
  17         substUnfoldingSource, lookupIdSubst, lookupTvSubst, substIdOcc,
  18
  19         -- ** Operations on substitutions
  20         emptySubst, mkEmptySubst, mkSubst, mkOpenSubst, substInScope, isEmptySubst,
  21         extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
  22         extendSubst, extendSubstList, zapSubstEnv,
  23         addInScopeSet, extendInScope, extendInScopeList, extendInScopeIds,
  24         isInScope, setInScope,
  25         delBndr, delBndrs,
  26
  27         -- ** Substituting and cloning binders
  28         substBndr, substBndrs, substRecBndrs,
  29         cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs,
  30
  31         -- ** Simple expression optimiser
  32         simpleOptPgm, simpleOptExpr, simpleOptExprWith
  33     ) where
  34
  35 #include "HsVersions.h"
  36
  37 import CoreSyn
  38 import CoreFVs
  39 import CoreUtils
  40 import PprCore
  41 import OccurAnal( occurAnalyseExpr, occurAnalysePgm )
  42
  43 import qualified Type
  44 import Type     ( Type, TvSubst(..), TvSubstEnv )
  45 import Coercion    ( isIdentityCoercion )
  46 import OptCoercion ( optCoercion )
  47 import VarSet
  48 import VarEnv
  49 import Id
  50 import Name     ( Name )
  51 import Var      ( Var, TyVar, setVarUnique )
  52 import IdInfo
  53 import Unique
  54 import UniqSupply
  55 import Maybes
  56 import ErrUtils
  57 import DynFlags   ( DynFlags, DynFlag(..) )
  58 import BasicTypes ( isAlwaysActive )
  59 import Outputable
  60 import PprCore          ()              -- Instances
  61 import FastString
  62
  63 import Data.List
  64 \end{code}
  65
  66
  67 %************************************************************************
  68 %*                                                                      *
  69 \subsection{Substitutions}
  70 %*                                                                      *
  71 %************************************************************************
  72
  73 \begin{code}
  74 -- | A substitution environment, containing both 'Id' and 'TyVar' substitutions.
  75 --
  76 -- Some invariants apply to how you use the substitution:
  77 --
  78 -- 1. #in_scope_invariant# The in-scope set contains at least those 'Id's and 'TyVar's that will be in scope /after/
  79 -- applying the substitution to a term. Precisely, the in-scope set must be a superset of the free vars of the
  80 -- substitution range that might possibly clash with locally-bound variables in the thing being substituted in.
  81 --
  82 -- 2. #apply_once# You may apply the substitution only /once/
  83 --
  84 -- There are various ways of setting up the in-scope set such that the first of these invariants hold:
  85 --
  86 -- * Arrange that the in-scope set really is all the things in scope
  87 --
  88 -- * Arrange that it's the free vars of the range of the substitution
  89 --
  90 -- * Make it empty, if you know that all the free vars of the substitution are fresh, and hence can't possibly clash
  91 data Subst
  92   = Subst InScopeSet  -- Variables in in scope (both Ids and TyVars) /after/
  93                       -- applying the substitution
  94           IdSubstEnv  -- Substitution for Ids
  95           TvSubstEnv  -- Substitution for TyVars
  96
  97         -- INVARIANT 1: See #in_scope_invariant#
  98         -- This is what lets us deal with name capture properly
  99         -- It's a hard invariant to check...
 100         --
 101         -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
 102         --              Types.TvSubstEnv
 103         --
 104         -- INVARIANT 3: See Note [Extending the Subst]
 105 \end{code}
 106
 107 Note [Extending the Subst]
 108 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 109 For a core Subst, which binds Ids as well, we make a different choice for Ids
 110 than we do for TyVars.
 111
 112 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
 113
 114 For Ids, we have a different invariant
 115         The IdSubstEnv is extended *only* when the Unique on an Id changes
 116         Otherwise, we just extend the InScopeSet
 117
 118 In consequence:
 119
 120 * If the TvSubstEnv and IdSubstEnv are both empty, substExpr would be a
 121   no-op, so substExprSC ("short cut") does nothing.
 122
 123   However, substExpr still goes ahead and substitutes.  Reason: we may
 124   want to replace existing Ids with new ones from the in-scope set, to
 125   avoid space leaks.
 126
 127 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
 128
 129 * In lookupIdSubst, we *must* look up the Id in the in-scope set, because
 130   it may contain non-trivial changes.  Example:
 131         (/\a. \x:a. ...x...) Int
 132   We extend the TvSubstEnv with [a |-> Int]; but x's unique does not change
 133   so we only extend the in-scope set.  Then we must look up in the in-scope
 134   set when we find the occurrence of x.
 135
 136 * The requirement to look up the Id in the in-scope set means that we
 137   must NOT take no-op short cut when the IdSubst is empty.
 138   We must still look up every Id in the in-scope set.
 139
 140 * (However, we don't need to do so for expressions found in the IdSubst
 141   itself, whose range is assumed to be correct wrt the in-scope set.)
 142
 143 Why do we make a different choice for the IdSubstEnv than the TvSubstEnv?
 144
 145 * For Ids, we change the IdInfo all the time (e.g. deleting the
 146   unfolding), and adding it back later, so using the TyVar convention
 147   would entail extending the substitution almost all the time
 148
 149 * The simplifier wants to look up in the in-scope set anyway, in case it
 150   can see a better unfolding from an enclosing case expression
 151
 152 * For TyVars, only coercion variables can possibly change, and they are
 153   easy to spot
 154
 155 \begin{code}
 156 -- | An environment for substituting for 'Id's
 157 type IdSubstEnv = IdEnv CoreExpr
 158
 159 ----------------------------
 160 isEmptySubst :: Subst -> Bool
 161 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
 162
 163 emptySubst :: Subst
 164 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
 165
 166 mkEmptySubst :: InScopeSet -> Subst
 167 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
 168
 169 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
 170 mkSubst in_scope tvs ids = Subst in_scope ids tvs
 171
 172 -- getTvSubst :: Subst -> TvSubst
 173 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
 174
 175 -- getTvSubstEnv :: Subst -> TvSubstEnv
 176 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
 177 --
 178 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
 179 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
 180
 181 -- | Find the in-scope set: see "CoreSubst#in_scope_invariant"
 182 substInScope :: Subst -> InScopeSet
 183 substInScope (Subst in_scope _ _) = in_scope
 184
 185 -- | Remove all substitutions for 'Id's and 'Var's that might have been built up
 186 -- while preserving the in-scope set
 187 zapSubstEnv :: Subst -> Subst
 188 zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
 189
 190 -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the in-scope set is
 191 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
 192 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
 193 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
 194 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
 195
 196 -- | Adds multiple 'Id' substitutions to the 'Subst': see also 'extendIdSubst'
 197 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
 198 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
 199
 200 -- | Add a substitution for a 'TyVar' to the 'Subst': you must ensure that the in-scope set is
 201 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
 202 extendTvSubst :: Subst -> TyVar -> Type -> Subst
 203 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
 204
 205 -- | Adds multiple 'TyVar' substitutions to the 'Subst': see also 'extendTvSubst'
 206 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
 207 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
 208
 209 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to the 'Var' being added. See also
 210 -- 'extendIdSubst' and 'extendTvSubst'
 211 extendSubst :: Subst -> Var -> CoreArg -> Subst
 212 extendSubst (Subst in_scope ids tvs) tv (Type ty)
 213   = ASSERT( isTyCoVar tv ) Subst in_scope ids (extendVarEnv tvs tv ty)
 214 extendSubst (Subst in_scope ids tvs) id expr
 215   = ASSERT( isId id ) Subst in_scope (extendVarEnv ids id expr) tvs
 216
 217 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to all the 'Var's being added. See also 'extendSubst'
 218 extendSubstList :: Subst -> [(Var,CoreArg)] -> Subst
 219 extendSubstList subst []              = subst
 220 extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var rhs) prs
 221
 222 -- | Find the substitution for an 'Id' in the 'Subst'
 223 lookupIdSubst :: SDoc -> Subst -> Id -> CoreExpr
 224 lookupIdSubst doc (Subst in_scope ids _) v
 225   | not (isLocalId v) = Var v
 226   | Just e  <- lookupVarEnv ids       v = e
 227   | Just v' <- lookupInScope in_scope v = Var v'
 228         -- Vital! See Note [Extending the Subst]
 229   | otherwise = WARN( True, ptext (sLit "CoreSubst.lookupIdSubst") <+> ppr v $$ ppr in_scope $$ doc)
 230                 Var v
 231
 232 -- | Find the substitution for a 'TyVar' in the 'Subst'
 233 lookupTvSubst :: Subst -> TyVar -> Type
 234 lookupTvSubst (Subst _ _ tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
 235
 236 delBndr :: Subst -> Var -> Subst
 237 delBndr (Subst in_scope tvs ids) v
 238   | isId v    = Subst in_scope tvs (delVarEnv ids v)
 239   | otherwise = Subst in_scope (delVarEnv tvs v) ids
 240
 241 delBndrs :: Subst -> [Var] -> Subst
 242 delBndrs (Subst in_scope tvs ids) vs
 243   = Subst in_scope (delVarEnvList tvs vs_tv) (delVarEnvList ids vs_id)
 244   where
 245     (vs_id, vs_tv) = partition isId vs
 246
 247 -- | Simultaneously substitute for a bunch of variables
 248 --   No left-right shadowing
 249 --   ie the substitution for   (\x \y. e) a1 a2
 250 --      so neither x nor y scope over a1 a2
 251 mkOpenSubst :: InScopeSet -> [(Var,CoreArg)] -> Subst
 252 mkOpenSubst in_scope pairs = Subst in_scope
 253                                    (mkVarEnv [(id,e)  | (id, e) <- pairs, isId id])
 254                                    (mkVarEnv [(tv,ty) | (tv, Type ty) <- pairs])
 255
 256 ------------------------------
 257 isInScope :: Var -> Subst -> Bool
 258 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
 259
 260 -- | Add the 'Var' to the in-scope set, but do not remove
 261 -- any existing substitutions for it
 262 addInScopeSet :: Subst -> VarSet -> Subst
 263 addInScopeSet (Subst in_scope ids tvs) vs
 264   = Subst (in_scope `extendInScopeSetSet` vs) ids tvs
 265
 266 -- | Add the 'Var' to the in-scope set: as a side effect,
 267 -- and remove any existing substitutions for it
 268 extendInScope :: Subst -> Var -> Subst
 269 extendInScope (Subst in_scope ids tvs) v
 270   = Subst (in_scope `extendInScopeSet` v)
 271           (ids `delVarEnv` v) (tvs `delVarEnv` v)
 272
 273 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
 274 extendInScopeList :: Subst -> [Var] -> Subst
 275 extendInScopeList (Subst in_scope ids tvs) vs
 276   = Subst (in_scope `extendInScopeSetList` vs)
 277           (ids `delVarEnvList` vs) (tvs `delVarEnvList` vs)
 278
 279 -- | Optimized version of 'extendInScopeList' that can be used if you are certain
 280 -- all the things being added are 'Id's and hence none are 'TyVar's
 281 extendInScopeIds :: Subst -> [Id] -> Subst
 282 extendInScopeIds (Subst in_scope ids tvs) vs
 283   = Subst (in_scope `extendInScopeSetList` vs)
 284           (ids `delVarEnvList` vs) tvs
 285
 286 setInScope :: Subst -> InScopeSet -> Subst
 287 setInScope (Subst _ ids tvs) in_scope = Subst in_scope ids tvs
 288 \end{code}
 289
 290 Pretty printing, for debugging only
 291
 292 \begin{code}
 293 instance Outputable Subst where
 294   ppr (Subst in_scope ids tvs)
 295         =  ptext (sLit "<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
 296         $$ ptext (sLit " IdSubst   =") <+> ppr ids
 297         $$ ptext (sLit " TvSubst   =") <+> ppr tvs
 298          <> char '>'
 299 \end{code}
 300
 301
 302 %************************************************************************
 303 %*                                                                      *
 304         Substituting expressions
 305 %*                                                                      *
 306 %************************************************************************
 307
 308 \begin{code}
 309 -- | Apply a substititon to an entire 'CoreExpr'. Rememeber, you may only
 310 -- apply the substitution /once/: see "CoreSubst#apply_once"
 311 --
 312 -- Do *not* attempt to short-cut in the case of an empty substitution!
 313 -- See Note [Extending the Subst]
 314 substExprSC :: SDoc -> Subst -> CoreExpr -> CoreExpr
 315 substExprSC _doc subst orig_expr
 316   | isEmptySubst subst = orig_expr
 317   | otherwise          = -- pprTrace "enter subst-expr" (doc $$ ppr orig_expr) $
 318                          subst_expr subst orig_expr
 319
 320 substExpr :: SDoc -> Subst -> CoreExpr -> CoreExpr
 321 substExpr _doc subst orig_expr = subst_expr subst orig_expr
 322
 323 subst_expr :: Subst -> CoreExpr -> CoreExpr
 324 subst_expr subst expr
 325   = go expr
 326   where
 327     go (Var v)         = lookupIdSubst (text "subst_expr") subst v
 328     go (Type ty)       = Type (substTy subst ty)
 329     go (Lit lit)       = Lit lit
 330     go (App fun arg)   = App (go fun) (go arg)
 331     go (Note note e)   = Note (go_note note) (go e)
 332     go (Cast e co)     = Cast (go e) (optCoercion (getTvSubst subst) co)
 333        -- Do not optimise even identity coercions
 334        -- Reason: substitution applies to the LHS of RULES, and
 335        --         if you "optimise" an identity coercion, you may
 336        --         lose a binder. We optimise the LHS of rules at
 337        --         construction time
 338
 339     go (Lam bndr body) = Lam bndr' (subst_expr subst' body)
 340                        where
 341                          (subst', bndr') = substBndr subst bndr
 342
 343     go (Let bind body) = Let bind' (subst_expr subst' body)
 344                        where
 345                          (subst', bind') = substBind subst bind
 346
 347     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
 348                                  where
 349                                  (subst', bndr') = substBndr subst bndr
 350
 351     go_alt subst (con, bndrs, rhs) = (con, bndrs', subst_expr subst' rhs)
 352                                  where
 353                                    (subst', bndrs') = substBndrs subst bndrs
 354
 355     go_note note             = note
 356
 357 -- | Apply a substititon to an entire 'CoreBind', additionally returning an updated 'Subst'
 358 -- that should be used by subsequent substitutons.
 359 substBind, substBindSC :: Subst -> CoreBind -> (Subst, CoreBind)
 360
 361 substBindSC subst bind    -- Short-cut if the substitution is empty
 362   | not (isEmptySubst subst)
 363   = substBind subst bind
 364   | otherwise
 365   = case bind of
 366        NonRec bndr rhs -> (subst', NonRec bndr' rhs)
 367           where
 368             (subst', bndr') = substBndr subst bndr
 369        Rec pairs -> (subst', Rec (bndrs' `zip` rhss'))
 370           where
 371             (bndrs, rhss)    = unzip pairs
 372             (subst', bndrs') = substRecBndrs subst bndrs
 373             rhss' | isEmptySubst subst' = rhss
 374                   | otherwise           = map (subst_expr subst') rhss
 375
 376 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (subst_expr subst rhs))
 377                                   where
 378                                     (subst', bndr') = substBndr subst bndr
 379
 380 substBind subst (Rec pairs) = (subst', Rec (bndrs' `zip` rhss'))
 381                             where
 382                                 (bndrs, rhss)    = unzip pairs
 383                                 (subst', bndrs') = substRecBndrs subst bndrs
 384                                 rhss' = map (subst_expr subst') rhss
 385 \end{code}
 386
 387 \begin{code}
 388 -- | De-shadowing the program is sometimes a useful pre-pass. It can be done simply
 389 -- by running over the bindings with an empty substitution, becuase substitution
 390 -- returns a result that has no-shadowing guaranteed.
 391 --
 392 -- (Actually, within a single /type/ there might still be shadowing, because
 393 -- 'substTy' is a no-op for the empty substitution, but that's probably OK.)
 394 --
 395 -- [Aug 09] This function is not used in GHC at the moment, but seems so
 396 --          short and simple that I'm going to leave it here
 397 deShadowBinds :: [CoreBind] -> [CoreBind]
 398 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
 399 \end{code}
 400
 401
 402 %************************************************************************
 403 %*                                                                      *
 404         Substituting binders
 405 %*                                                                      *
 406 %************************************************************************
 407
 408 Remember that substBndr and friends are used when doing expression
 409 substitution only.  Their only business is substitution, so they
 410 preserve all IdInfo (suitably substituted).  For example, we *want* to
 411 preserve occ info in rules.
 412
 413 \begin{code}
 414 -- | Substitutes a 'Var' for another one according to the 'Subst' given, returning
 415 -- the result and an updated 'Subst' that should be used by subsequent substitutons.
 416 -- 'IdInfo' is preserved by this process, although it is substituted into appropriately.
 417 substBndr :: Subst -> Var -> (Subst, Var)
 418 substBndr subst bndr
 419   | isTyCoVar bndr  = substTyVarBndr subst bndr
 420   | otherwise       = substIdBndr (text "var-bndr") subst subst bndr
 421
 422 -- | Applies 'substBndr' to a number of 'Var's, accumulating a new 'Subst' left-to-right
 423 substBndrs :: Subst -> [Var] -> (Subst, [Var])
 424 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 425
 426 -- | Substitute in a mutually recursive group of 'Id's
 427 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
 428 substRecBndrs subst bndrs
 429   = (new_subst, new_bndrs)
 430   where         -- Here's the reason we need to pass rec_subst to subst_id
 431     (new_subst, new_bndrs) = mapAccumL (substIdBndr (text "rec-bndr") new_subst) subst bndrs
 432 \end{code}
 433
 434
 435 \begin{code}
 436 substIdBndr :: SDoc
 437             -> Subst            -- ^ Substitution to use for the IdInfo
 438             -> Subst -> Id      -- ^ Substitition and Id to transform
 439             -> (Subst, Id)      -- ^ Transformed pair
 440                                 -- NB: unfolding may be zapped
 441
 442 substIdBndr _doc rec_subst subst@(Subst in_scope env tvs) old_id
 443   = -- pprTrace "substIdBndr" (doc $$ ppr old_id $$ ppr in_scope) $
 444     (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 445   where
 446     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
 447     id2 | no_type_change = id1
 448         | otherwise      = setIdType id1 (substTy subst old_ty)
 449
 450     old_ty = idType old_id
 451     no_type_change = isEmptyVarEnv tvs ||
 452                      isEmptyVarSet (Type.tyVarsOfType old_ty)
 453
 454         -- new_id has the right IdInfo
 455         -- The lazy-set is because we're in a loop here, with
 456         -- rec_subst, when dealing with a mutually-recursive group
 457     new_id = maybeModifyIdInfo mb_new_info id2
 458     mb_new_info = substIdInfo rec_subst id2 (idInfo id2)
 459         -- NB: unfolding info may be zapped
 460
 461         -- Extend the substitution if the unique has changed
 462         -- See the notes with substTyVarBndr for the delVarEnv
 463     new_env | no_change = delVarEnv env old_id
 464             | otherwise = extendVarEnv env old_id (Var new_id)
 465
 466     no_change = id1 == old_id
 467         -- See Note [Extending the Subst]
 468         -- it's /not/ necessary to check mb_new_info and no_type_change
 469 \end{code}
 470
 471 Now a variant that unconditionally allocates a new unique.
 472 It also unconditionally zaps the OccInfo.
 473
 474 \begin{code}
 475 -- | Very similar to 'substBndr', but it always allocates a new 'Unique' for
 476 -- each variable in its output.  It substitutes the IdInfo though.
 477 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
 478 cloneIdBndr subst us old_id
 479   = clone_id subst subst (old_id, uniqFromSupply us)
 480
 481 -- | Applies 'cloneIdBndr' to a number of 'Id's, accumulating a final
 482 -- substitution from left to right
 483 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 484 cloneIdBndrs subst us ids
 485   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
 486
 487 -- | Clone a mutually recursive group of 'Id's
 488 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 489 cloneRecIdBndrs subst us ids
 490   = (subst', ids')
 491   where
 492     (subst', ids') = mapAccumL (clone_id subst') subst
 493                                (ids `zip` uniqsFromSupply us)
 494
 495 -- Just like substIdBndr, except that it always makes a new unique
 496 -- It is given the unique to use
 497 clone_id    :: Subst                    -- Substitution for the IdInfo
 498             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
 499             -> (Subst, Id)              -- Transformed pair
 500
 501 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
 502   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 503   where
 504     id1     = setVarUnique old_id uniq
 505     id2     = substIdType subst id1
 506     new_id  = maybeModifyIdInfo (substIdInfo rec_subst id2 (idInfo old_id)) id2
 507     new_env = extendVarEnv env old_id (Var new_id)
 508 \end{code}
 509
 510
 511 %************************************************************************
 512 %*                                                                      *
 513                 Types
 514 %*                                                                      *
 515 %************************************************************************
 516
 517 For types we just call the corresponding function in Type, but we have
 518 to repackage the substitution, from a Subst to a TvSubst
 519
 520 \begin{code}
 521 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
 522 substTyVarBndr (Subst in_scope id_env tv_env) tv
 523   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
 524         (TvSubst in_scope' tv_env', tv')
 525            -> (Subst in_scope' id_env tv_env', tv')
 526
 527 -- | See 'Type.substTy'
 528 substTy :: Subst -> Type -> Type
 529 substTy subst ty = Type.substTy (getTvSubst subst) ty
 530
 531 getTvSubst :: Subst -> TvSubst
 532 getTvSubst (Subst in_scope _id_env tv_env) = TvSubst in_scope tv_env
 533 \end{code}
 534
 535
 536 %************************************************************************
 537 %*                                                                      *
 538 \section{IdInfo substitution}
 539 %*                                                                      *
 540 %************************************************************************
 541
 542 \begin{code}
 543 substIdType :: Subst -> Id -> Id
 544 substIdType subst@(Subst _ _ tv_env) id
 545   | isEmptyVarEnv tv_env || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
 546   | otherwise   = setIdType id (substTy subst old_ty)
 547                 -- The tyVarsOfType is cheaper than it looks
 548                 -- because we cache the free tyvars of the type
 549                 -- in a Note in the id's type itself
 550   where
 551     old_ty = idType id
 552
 553 ------------------
 554 -- | Substitute into some 'IdInfo' with regard to the supplied new 'Id'.
 555 substIdInfo :: Subst -> Id -> IdInfo -> Maybe IdInfo
 556 substIdInfo subst new_id info
 557   | nothing_to_do = Nothing
 558   | otherwise     = Just (info `setSpecInfo`      substSpec subst new_id old_rules
 559                                `setUnfoldingInfo` substUnfolding subst old_unf)
 560   where
 561     old_rules     = specInfo info
 562     old_unf       = unfoldingInfo info
 563     nothing_to_do = isEmptySpecInfo old_rules && isClosedUnfolding old_unf
 564
 565
 566 ------------------
 567 -- | Substitutes for the 'Id's within an unfolding
 568 substUnfolding, substUnfoldingSC :: Subst -> Unfolding -> Unfolding
 569         -- Seq'ing on the returned Unfolding is enough to cause
 570         -- all the substitutions to happen completely
 571
 572 substUnfoldingSC subst unf       -- Short-cut version
 573   | isEmptySubst subst = unf
 574   | otherwise          = substUnfolding subst unf
 575
 576 substUnfolding subst (DFunUnfolding ar con args)
 577   = DFunUnfolding ar con (map (substExpr (text "dfun-unf") subst) args)
 578
 579 substUnfolding subst unf@(CoreUnfolding { uf_tmpl = tmpl, uf_src = src })
 580         -- Retain an InlineRule!
 581   | not (isStableSource src)  -- Zap an unstable unfolding, to save substitution work
 582   = NoUnfolding
 583   | otherwise                 -- But keep a stable one!
 584   = seqExpr new_tmpl `seq`
 585     new_src `seq`
 586     unf { uf_tmpl = new_tmpl, uf_src = new_src }
 587   where
 588     new_tmpl = substExpr (text "subst-unf") subst tmpl
 589     new_src  = substUnfoldingSource subst src
 590
 591 substUnfolding _ unf = unf      -- NoUnfolding, OtherCon
 592
 593 -------------------
 594 substUnfoldingSource :: Subst -> UnfoldingSource -> UnfoldingSource
 595 substUnfoldingSource (Subst in_scope ids _) (InlineWrapper wkr)
 596   | Just wkr_expr <- lookupVarEnv ids wkr
 597   = case wkr_expr of
 598       Var w1 -> InlineWrapper w1
 599       _other -> -- WARN( True, text "Interesting! CoreSubst.substWorker1:" <+> ppr wkr
 600                 --             <+> ifPprDebug (equals <+> ppr wkr_expr) )
 601                               -- Note [Worker inlining]
 602                 InlineStable  -- It's not a wrapper any more, but still inline it!
 603
 604   | Just w1  <- lookupInScope in_scope wkr = InlineWrapper w1
 605   | otherwise = -- WARN( True, text "Interesting! CoreSubst.substWorker2:" <+> ppr wkr )
 606                 -- This can legitimately happen.  The worker has been inlined and
 607                 -- dropped as dead code, because we don't treat the UnfoldingSource
 608                 -- as an "occurrence".
 609                 -- Note [Worker inlining]
 610                 InlineStable
 611
 612 substUnfoldingSource _ src = src
 613
 614 ------------------
 615 substIdOcc :: Subst -> Id -> Id
 616 -- These Ids should not be substituted to non-Ids
 617 substIdOcc subst v = case lookupIdSubst (text "substIdOcc") subst v of
 618                         Var v' -> v'
 619                         other  -> pprPanic "substIdOcc" (vcat [ppr v <+> ppr other, ppr subst])
 620
 621 ------------------
 622 -- | Substitutes for the 'Id's within the 'WorkerInfo' given the new function 'Id'
 623 substSpec :: Subst -> Id -> SpecInfo -> SpecInfo
 624 substSpec subst new_id (SpecInfo rules rhs_fvs)
 625   = seqSpecInfo new_spec `seq` new_spec
 626   where
 627     subst_ru_fn = const (idName new_id)
 628     new_spec = SpecInfo (map (substRule subst subst_ru_fn) rules)
 629                          (substVarSet subst rhs_fvs)
 630
 631 ------------------
 632 substRulesForImportedIds :: Subst -> [CoreRule] -> [CoreRule]
 633 substRulesForImportedIds subst rules
 634   = map (substRule subst not_needed) rules
 635   where
 636     not_needed name = pprPanic "substRulesForImportedIds" (ppr name)
 637
 638 ------------------
 639 substRule :: Subst -> (Name -> Name) -> CoreRule -> CoreRule
 640
 641 -- The subst_ru_fn argument is applied to substitute the ru_fn field
 642 -- of the rule:
 643 --    - Rules for *imported* Ids never change ru_fn
 644 --    - Rules for *local* Ids are in the IdInfo for that Id,
 645 --      and the ru_fn field is simply replaced by the new name
 646 --      of the Id
 647
 648 substRule _ _ rule@(BuiltinRule {}) = rule
 649 substRule subst subst_ru_fn rule@(Rule { ru_bndrs = bndrs, ru_args = args
 650                                        , ru_fn = fn_name, ru_rhs = rhs
 651                                        , ru_local = is_local })
 652   = rule { ru_bndrs = bndrs',
 653            ru_fn    = if is_local
 654                         then subst_ru_fn fn_name
 655                         else fn_name,
 656            ru_args  = map (substExpr (text "subst-rule" <+> ppr fn_name) subst') args,
 657            ru_rhs   = simpleOptExprWith subst' rhs }
 658            -- Do simple optimisation on RHS, in case substitution lets
 659            -- you improve it.  The real simplifier never gets to look at it.
 660   where
 661     (subst', bndrs') = substBndrs subst bndrs
 662
 663 ------------------
 664 substVarSet :: Subst -> VarSet -> VarSet
 665 substVarSet subst fvs
 666   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
 667   where
 668     subst_fv subst fv
 669         | isId fv   = exprFreeVars (lookupIdSubst (text "substVarSet") subst fv)
 670         | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)
 671 \end{code}
 672
 673 Note [Worker inlining]
 674 ~~~~~~~~~~~~~~~~~~~~~~
 675 A worker can get sustituted away entirely.
 676         - it might be trivial
 677         - it might simply be very small
 678 We do not treat an InlWrapper as an 'occurrence' in the occurence
 679 analyser, so it's possible that the worker is not even in scope any more.
 680
 681 In all all these cases we simply drop the special case, returning to
 682 InlVanilla.  The WARN is just so I can see if it happens a lot.
 683
 684
 685 %************************************************************************
 686 %*                                                                      *
 687         The Very Simple Optimiser
 688 %*                                                                      *
 689 %************************************************************************
 690
 691 \begin{code}
 692 simpleOptExpr :: CoreExpr -> CoreExpr
 693 -- Do simple optimisation on an expression
 694 -- The optimisation is very straightforward: just
 695 -- inline non-recursive bindings that are used only once,
 696 -- or where the RHS is trivial
 697 --
 698 -- The result is NOT guaranteed occurence-analysed, becuase
 699 -- in  (let x = y in ....) we substitute for x; so y's occ-info
 700 -- may change radically
 701
 702 simpleOptExpr expr
 703   = -- pprTrace "simpleOptExpr" (ppr init_subst $$ ppr expr)
 704     simpleOptExprWith init_subst expr
 705   where
 706     init_subst = mkEmptySubst (mkInScopeSet (exprFreeVars expr))
 707         -- It's potentially important to make a proper in-scope set
 708         -- Consider  let x = ..y.. in \y. ...x...
 709         -- Then we should remember to clone y before substituting
 710         -- for x.  It's very unlikely to occur, because we probably
 711         -- won't *be* substituting for x if it occurs inside a
 712         -- lambda.
 713         --
 714         -- It's a bit painful to call exprFreeVars, because it makes
 715         -- three passes instead of two (occ-anal, and go)
 716
 717 simpleOptExprWith :: Subst -> InExpr -> OutExpr
 718 simpleOptExprWith subst expr = simple_opt_expr subst (occurAnalyseExpr expr)
 719
 720 ----------------------
 721 simpleOptPgm :: DynFlags -> [CoreBind] -> [CoreRule] -> IO ([CoreBind], [CoreRule])
 722 simpleOptPgm dflags binds rules
 723   = do { dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
 724                        (pprCoreBindings occ_anald_binds);
 725
 726        ; return (reverse binds', substRulesForImportedIds subst' rules) }
 727   where
 728     occ_anald_binds  = occurAnalysePgm Nothing {- No rules active -}
 729                                        rules binds
 730     (subst', binds') = foldl do_one (emptySubst, []) occ_anald_binds
 731
 732     do_one (subst, binds') bind
 733       = case simple_opt_bind subst bind of
 734           (subst', Nothing)    -> (subst', binds')
 735           (subst', Just bind') -> (subst', bind':binds')
 736
 737 ----------------------
 738 type InVar   = Var
 739 type OutVar  = Var
 740 type InId    = Id
 741 type OutId   = Id
 742 type InExpr  = CoreExpr
 743 type OutExpr = CoreExpr
 744
 745 -- In these functions the substitution maps InVar -> OutExpr
 746
 747 ----------------------
 748 simple_opt_expr :: Subst -> InExpr -> OutExpr
 749 simple_opt_expr subst expr
 750   = go expr
 751   where
 752     go (Var v)          = lookupIdSubst (text "simpleOptExpr") subst v
 753     go (App e1 e2)      = simple_app subst e1 [go e2]
 754     go (Type ty)        = Type (substTy subst ty)
 755     go (Lit lit)        = Lit lit
 756     go (Note note e)    = Note note (go e)
 757     go (Cast e co)      | isIdentityCoercion co' = go e
 758                         | otherwise              = Cast (go e) co'
 759                         where
 760                           co' = substTy subst co
 761
 762     go (Let bind body) = case simple_opt_bind subst bind of
 763                            (subst', Nothing)   -> simple_opt_expr subst' body
 764                            (subst', Just bind) -> Let bind (simple_opt_expr subst' body)
 765
 766     go lam@(Lam {})     = go_lam [] subst lam
 767     go (Case e b ty as) = Case (go e) b' (substTy subst ty)
 768                                (map (go_alt subst') as)
 769                         where
 770                           (subst', b') = subst_opt_bndr subst b
 771
 772     ----------------------
 773     go_alt subst (con, bndrs, rhs)
 774       = (con, bndrs', simple_opt_expr subst' rhs)
 775       where
 776         (subst', bndrs') = subst_opt_bndrs subst bndrs
 777
 778     ----------------------
 779     -- go_lam tries eta reduction
 780     go_lam bs' subst (Lam b e)
 781        = go_lam (b':bs') subst' e
 782        where
 783          (subst', b') = subst_opt_bndr subst b
 784     go_lam bs' subst e
 785        | Just etad_e <- tryEtaReduce bs e' = etad_e
 786        | otherwise                         = mkLams bs e'
 787        where
 788          bs = reverse bs'
 789          e' = simple_opt_expr subst e
 790
 791 ----------------------
 792 -- simple_app collects arguments for beta reduction
 793 simple_app :: Subst -> InExpr -> [OutExpr] -> CoreExpr
 794 simple_app subst (App e1 e2) as
 795   = simple_app subst e1 (simple_opt_expr subst e2 : as)
 796 simple_app subst (Lam b e) (a:as)
 797   = case maybe_substitute subst b a of
 798       Just ext_subst -> simple_app ext_subst e as
 799       Nothing        -> Let (NonRec b2 a) (simple_app subst' e as)
 800   where
 801     (subst', b') = subst_opt_bndr subst b
 802     b2 = add_info subst' b b'
 803 simple_app subst e as
 804   = foldl App (simple_opt_expr subst e) as
 805
 806 ----------------------
 807 simple_opt_bind :: Subst -> CoreBind -> (Subst, Maybe CoreBind)
 808 simple_opt_bind subst (Rec prs)
 809   = (subst'', Just (Rec (reverse rev_prs')))
 810   where
 811     (subst', bndrs')    = subst_opt_bndrs subst (map fst prs)
 812     (subst'', rev_prs') = foldl do_pr (subst', []) (prs `zip` bndrs')
 813     do_pr (subst, prs) ((b,r), b')
 814        = case maybe_substitute subst b r2 of
 815            Just subst' -> (subst', prs)
 816            Nothing     -> (subst,  (b2,r2):prs)
 817        where
 818          b2 = add_info subst b b'
 819          r2 = simple_opt_expr subst r
 820
 821 simple_opt_bind subst (NonRec b r)
 822   = case maybe_substitute subst b r' of
 823       Just ext_subst -> (ext_subst, Nothing)
 824       Nothing        -> (subst', Just (NonRec b2 r'))
 825   where
 826     r' = simple_opt_expr subst r
 827     (subst', b') = subst_opt_bndr subst b
 828     b2 = add_info subst' b b'
 829
 830 ----------------------
 831 maybe_substitute :: Subst -> InVar -> OutExpr -> Maybe Subst
 832     -- (maybe_substitute subst in_var out_rhs)
 833     --   either extends subst with (in_var -> out_rhs)
 834     --   or     returns Nothing
 835 maybe_substitute subst b r
 836   | Type ty <- r        -- let a::* = TYPE ty in <body>
 837   = ASSERT( isTyCoVar b )
 838     Just (extendTvSubst subst b ty)
 839
 840   | isId b              -- let x = e in <body>
 841   , safe_to_inline (idOccInfo b)
 842   , isAlwaysActive (idInlineActivation b)       -- Note [Inline prag in simplOpt]
 843   , not (isStableUnfolding (idUnfolding b))
 844   , not (isExportedId b)
 845   = Just (extendIdSubst subst b r)
 846
 847   | otherwise
 848   = Nothing
 849   where
 850         -- Unconditionally safe to inline
 851     safe_to_inline :: OccInfo -> Bool
 852     safe_to_inline (IAmALoopBreaker {})     = False
 853     safe_to_inline IAmDead                  = True
 854     safe_to_inline (OneOcc in_lam one_br _) = (not in_lam && one_br) || exprIsTrivial r
 855     safe_to_inline NoOccInfo                = exprIsTrivial r
 856
 857 ----------------------
 858 subst_opt_bndr :: Subst -> InVar -> (Subst, OutVar)
 859 subst_opt_bndr subst bndr
 860   | isTyCoVar bndr  = substTyVarBndr subst bndr
 861   | otherwise       = subst_opt_id_bndr subst bndr
 862
 863 subst_opt_id_bndr :: Subst -> InId -> (Subst, OutId)
 864 -- Nuke all fragile IdInfo, unfolding, and RULES;
 865 --    it gets added back later by add_info
 866 -- Rather like SimplEnv.substIdBndr
 867 --
 868 -- It's important to zap fragile OccInfo (which CoreSubst.SubstIdBndr
 869 -- carefully does not do) because simplOptExpr invalidates it
 870
 871 subst_opt_id_bndr subst@(Subst in_scope id_subst tv_subst) old_id
 872   = (Subst new_in_scope new_id_subst tv_subst, new_id)
 873   where
 874     id1    = uniqAway in_scope old_id
 875     id2    = setIdType id1 (substTy subst (idType old_id))
 876     new_id = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
 877                                         -- and fragile OccInfo
 878     new_in_scope = in_scope `extendInScopeSet` new_id
 879
 880         -- Extend the substitution if the unique has changed,
 881         -- or there's some useful occurrence information
 882         -- See the notes with substTyVarBndr for the delSubstEnv
 883     new_id_subst | new_id /= old_id
 884                  = extendVarEnv id_subst old_id (Var new_id)
 885                  | otherwise
 886                  = delVarEnv id_subst old_id
 887
 888 ----------------------
 889 subst_opt_bndrs :: Subst -> [InVar] -> (Subst, [OutVar])
 890 subst_opt_bndrs subst bndrs
 891   = mapAccumL subst_opt_bndr subst bndrs
 892
 893 ----------------------
 894 add_info :: Subst -> InVar -> OutVar -> OutVar
 895 add_info subst old_bndr new_bndr
 896  | isTyCoVar old_bndr = new_bndr
 897  | otherwise          = maybeModifyIdInfo mb_new_info new_bndr
 898  where
 899    mb_new_info = substIdInfo subst new_bndr (idInfo old_bndr)
 900 \end{code}
 901
 902 Note [Inline prag in simplOpt]
 903 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 904 If there's an INLINE/NOINLINE pragma that restricts the phase in
 905 which the binder can be inlined, we don't inline here; after all,
 906 we don't know what phase we're in.  Here's an example
 907
 908   foo :: Int -> Int -> Int
 909   {-# INLINE foo #-}
 910   foo m n = inner m
 911      where
 912        {-# INLINE [1] inner #-}
 913        inner m = m+n
 914
 915   bar :: Int -> Int
 916   bar n = foo n 1
 917
 918 When inlining 'foo' in 'bar' we want the let-binding for 'inner'
 919 to remain visible until Phase 1
 920