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 subst_arg args)
 578   where
 579     subst_arg = fmap (substExpr (text "dfun-unf") subst)
 580
 581 substUnfolding subst unf@(CoreUnfolding { uf_tmpl = tmpl, uf_src = src })
 582         -- Retain an InlineRule!
 583   | not (isStableSource src)  -- Zap an unstable unfolding, to save substitution work
 584   = NoUnfolding
 585   | otherwise                 -- But keep a stable one!
 586   = seqExpr new_tmpl `seq`
 587     new_src `seq`
 588     unf { uf_tmpl = new_tmpl, uf_src = new_src }
 589   where
 590     new_tmpl = substExpr (text "subst-unf") subst tmpl
 591     new_src  = substUnfoldingSource subst src
 592
 593 substUnfolding _ unf = unf      -- NoUnfolding, OtherCon
 594
 595 -------------------
 596 substUnfoldingSource :: Subst -> UnfoldingSource -> UnfoldingSource
 597 substUnfoldingSource (Subst in_scope ids _) (InlineWrapper wkr)
 598   | Just wkr_expr <- lookupVarEnv ids wkr
 599   = case wkr_expr of
 600       Var w1 -> InlineWrapper w1
 601       _other -> -- WARN( True, text "Interesting! CoreSubst.substWorker1:" <+> ppr wkr
 602                 --             <+> ifPprDebug (equals <+> ppr wkr_expr) )
 603                               -- Note [Worker inlining]
 604                 InlineStable  -- It's not a wrapper any more, but still inline it!
 605
 606   | Just w1  <- lookupInScope in_scope wkr = InlineWrapper w1
 607   | otherwise = -- WARN( True, text "Interesting! CoreSubst.substWorker2:" <+> ppr wkr )
 608                 -- This can legitimately happen.  The worker has been inlined and
 609                 -- dropped as dead code, because we don't treat the UnfoldingSource
 610                 -- as an "occurrence".
 611                 -- Note [Worker inlining]
 612                 InlineStable
 613
 614 substUnfoldingSource _ src = src
 615
 616 ------------------
 617 substIdOcc :: Subst -> Id -> Id
 618 -- These Ids should not be substituted to non-Ids
 619 substIdOcc subst v = case lookupIdSubst (text "substIdOcc") subst v of
 620                         Var v' -> v'
 621                         other  -> pprPanic "substIdOcc" (vcat [ppr v <+> ppr other, ppr subst])
 622
 623 ------------------
 624 -- | Substitutes for the 'Id's within the 'WorkerInfo' given the new function 'Id'
 625 substSpec :: Subst -> Id -> SpecInfo -> SpecInfo
 626 substSpec subst new_id (SpecInfo rules rhs_fvs)
 627   = seqSpecInfo new_spec `seq` new_spec
 628   where
 629     subst_ru_fn = const (idName new_id)
 630     new_spec = SpecInfo (map (substRule subst subst_ru_fn) rules)
 631                          (substVarSet subst rhs_fvs)
 632
 633 ------------------
 634 substRulesForImportedIds :: Subst -> [CoreRule] -> [CoreRule]
 635 substRulesForImportedIds subst rules
 636   = map (substRule subst not_needed) rules
 637   where
 638     not_needed name = pprPanic "substRulesForImportedIds" (ppr name)
 639
 640 ------------------
 641 substRule :: Subst -> (Name -> Name) -> CoreRule -> CoreRule
 642
 643 -- The subst_ru_fn argument is applied to substitute the ru_fn field
 644 -- of the rule:
 645 --    - Rules for *imported* Ids never change ru_fn
 646 --    - Rules for *local* Ids are in the IdInfo for that Id,
 647 --      and the ru_fn field is simply replaced by the new name
 648 --      of the Id
 649
 650 substRule _ _ rule@(BuiltinRule {}) = rule
 651 substRule subst subst_ru_fn rule@(Rule { ru_bndrs = bndrs, ru_args = args
 652                                        , ru_fn = fn_name, ru_rhs = rhs
 653                                        , ru_local = is_local })
 654   = rule { ru_bndrs = bndrs',
 655            ru_fn    = if is_local
 656                         then subst_ru_fn fn_name
 657                         else fn_name,
 658            ru_args  = map (substExpr (text "subst-rule" <+> ppr fn_name) subst') args,
 659            ru_rhs   = simpleOptExprWith subst' rhs }
 660            -- Do simple optimisation on RHS, in case substitution lets
 661            -- you improve it.  The real simplifier never gets to look at it.
 662   where
 663     (subst', bndrs') = substBndrs subst bndrs
 664
 665 ------------------
 666 substVarSet :: Subst -> VarSet -> VarSet
 667 substVarSet subst fvs
 668   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
 669   where
 670     subst_fv subst fv
 671         | isId fv   = exprFreeVars (lookupIdSubst (text "substVarSet") subst fv)
 672         | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)
 673 \end{code}
 674
 675 Note [Worker inlining]
 676 ~~~~~~~~~~~~~~~~~~~~~~
 677 A worker can get sustituted away entirely.
 678         - it might be trivial
 679         - it might simply be very small
 680 We do not treat an InlWrapper as an 'occurrence' in the occurence
 681 analyser, so it's possible that the worker is not even in scope any more.
 682
 683 In all all these cases we simply drop the special case, returning to
 684 InlVanilla.  The WARN is just so I can see if it happens a lot.
 685
 686
 687 %************************************************************************
 688 %*                                                                      *
 689         The Very Simple Optimiser
 690 %*                                                                      *
 691 %************************************************************************
 692
 693 \begin{code}
 694 simpleOptExpr :: CoreExpr -> CoreExpr
 695 -- Do simple optimisation on an expression
 696 -- The optimisation is very straightforward: just
 697 -- inline non-recursive bindings that are used only once,
 698 -- or where the RHS is trivial
 699 --
 700 -- The result is NOT guaranteed occurence-analysed, becuase
 701 -- in  (let x = y in ....) we substitute for x; so y's occ-info
 702 -- may change radically
 703
 704 simpleOptExpr expr
 705   = -- pprTrace "simpleOptExpr" (ppr init_subst $$ ppr expr)
 706     simpleOptExprWith init_subst expr
 707   where
 708     init_subst = mkEmptySubst (mkInScopeSet (exprFreeVars expr))
 709         -- It's potentially important to make a proper in-scope set
 710         -- Consider  let x = ..y.. in \y. ...x...
 711         -- Then we should remember to clone y before substituting
 712         -- for x.  It's very unlikely to occur, because we probably
 713         -- won't *be* substituting for x if it occurs inside a
 714         -- lambda.
 715         --
 716         -- It's a bit painful to call exprFreeVars, because it makes
 717         -- three passes instead of two (occ-anal, and go)
 718
 719 simpleOptExprWith :: Subst -> InExpr -> OutExpr
 720 simpleOptExprWith subst expr = simple_opt_expr subst (occurAnalyseExpr expr)
 721
 722 ----------------------
 723 simpleOptPgm :: DynFlags -> [CoreBind] -> [CoreRule] -> IO ([CoreBind], [CoreRule])
 724 simpleOptPgm dflags binds rules
 725   = do { dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
 726                        (pprCoreBindings occ_anald_binds);
 727
 728        ; return (reverse binds', substRulesForImportedIds subst' rules) }
 729   where
 730     occ_anald_binds  = occurAnalysePgm Nothing {- No rules active -}
 731                                        rules binds
 732     (subst', binds') = foldl do_one (emptySubst, []) occ_anald_binds
 733
 734     do_one (subst, binds') bind
 735       = case simple_opt_bind subst bind of
 736           (subst', Nothing)    -> (subst', binds')
 737           (subst', Just bind') -> (subst', bind':binds')
 738
 739 ----------------------
 740 type InVar   = Var
 741 type OutVar  = Var
 742 type InId    = Id
 743 type OutId   = Id
 744 type InExpr  = CoreExpr
 745 type OutExpr = CoreExpr
 746
 747 -- In these functions the substitution maps InVar -> OutExpr
 748
 749 ----------------------
 750 simple_opt_expr :: Subst -> InExpr -> OutExpr
 751 simple_opt_expr subst expr
 752   = go expr
 753   where
 754     go (Var v)          = lookupIdSubst (text "simpleOptExpr") subst v
 755     go (App e1 e2)      = simple_app subst e1 [go e2]
 756     go (Type ty)        = Type (substTy subst ty)
 757     go (Lit lit)        = Lit lit
 758     go (Note note e)    = Note note (go e)
 759     go (Cast e co)      | isIdentityCoercion co' = go e
 760                         | otherwise              = Cast (go e) co'
 761                         where
 762                           co' = substTy subst co
 763
 764     go (Let bind body) = case simple_opt_bind subst bind of
 765                            (subst', Nothing)   -> simple_opt_expr subst' body
 766                            (subst', Just bind) -> Let bind (simple_opt_expr subst' body)
 767
 768     go lam@(Lam {})     = go_lam [] subst lam
 769     go (Case e b ty as) = Case (go e) b' (substTy subst ty)
 770                                (map (go_alt subst') as)
 771                         where
 772                           (subst', b') = subst_opt_bndr subst b
 773
 774     ----------------------
 775     go_alt subst (con, bndrs, rhs)
 776       = (con, bndrs', simple_opt_expr subst' rhs)
 777       where
 778         (subst', bndrs') = subst_opt_bndrs subst bndrs
 779
 780     ----------------------
 781     -- go_lam tries eta reduction
 782     go_lam bs' subst (Lam b e)
 783        = go_lam (b':bs') subst' e
 784        where
 785          (subst', b') = subst_opt_bndr subst b
 786     go_lam bs' subst e
 787        | Just etad_e <- tryEtaReduce bs e' = etad_e
 788        | otherwise                         = mkLams bs e'
 789        where
 790          bs = reverse bs'
 791          e' = simple_opt_expr subst e
 792
 793 ----------------------
 794 -- simple_app collects arguments for beta reduction
 795 simple_app :: Subst -> InExpr -> [OutExpr] -> CoreExpr
 796 simple_app subst (App e1 e2) as
 797   = simple_app subst e1 (simple_opt_expr subst e2 : as)
 798 simple_app subst (Lam b e) (a:as)
 799   = case maybe_substitute subst b a of
 800       Just ext_subst -> simple_app ext_subst e as
 801       Nothing        -> Let (NonRec b2 a) (simple_app subst' e as)
 802   where
 803     (subst', b') = subst_opt_bndr subst b
 804     b2 = add_info subst' b b'
 805 simple_app subst e as
 806   = foldl App (simple_opt_expr subst e) as
 807
 808 ----------------------
 809 simple_opt_bind :: Subst -> CoreBind -> (Subst, Maybe CoreBind)
 810 simple_opt_bind subst (Rec prs)
 811   = (subst'', Just (Rec (reverse rev_prs')))
 812   where
 813     (subst', bndrs')    = subst_opt_bndrs subst (map fst prs)
 814     (subst'', rev_prs') = foldl do_pr (subst', []) (prs `zip` bndrs')
 815     do_pr (subst, prs) ((b,r), b')
 816        = case maybe_substitute subst b r2 of
 817            Just subst' -> (subst', prs)
 818            Nothing     -> (subst,  (b2,r2):prs)
 819        where
 820          b2 = add_info subst b b'
 821          r2 = simple_opt_expr subst r
 822
 823 simple_opt_bind subst (NonRec b r)
 824   = case maybe_substitute subst b r' of
 825       Just ext_subst -> (ext_subst, Nothing)
 826       Nothing        -> (subst', Just (NonRec b2 r'))
 827   where
 828     r' = simple_opt_expr subst r
 829     (subst', b') = subst_opt_bndr subst b
 830     b2 = add_info subst' b b'
 831
 832 ----------------------
 833 maybe_substitute :: Subst -> InVar -> OutExpr -> Maybe Subst
 834     -- (maybe_substitute subst in_var out_rhs)
 835     --   either extends subst with (in_var -> out_rhs)
 836     --   or     returns Nothing
 837 maybe_substitute subst b r
 838   | Type ty <- r        -- let a::* = TYPE ty in <body>
 839   = ASSERT( isTyCoVar b )
 840     Just (extendTvSubst subst b ty)
 841
 842   | isId b              -- let x = e in <body>
 843   , safe_to_inline (idOccInfo b)
 844   , isAlwaysActive (idInlineActivation b)       -- Note [Inline prag in simplOpt]
 845   , not (isStableUnfolding (idUnfolding b))
 846   , not (isExportedId b)
 847   = Just (extendIdSubst subst b r)
 848
 849   | otherwise
 850   = Nothing
 851   where
 852         -- Unconditionally safe to inline
 853     safe_to_inline :: OccInfo -> Bool
 854     safe_to_inline (IAmALoopBreaker {})     = False
 855     safe_to_inline IAmDead                  = True
 856     safe_to_inline (OneOcc in_lam one_br _) = (not in_lam && one_br) || exprIsTrivial r
 857     safe_to_inline NoOccInfo                = exprIsTrivial r
 858
 859 ----------------------
 860 subst_opt_bndr :: Subst -> InVar -> (Subst, OutVar)
 861 subst_opt_bndr subst bndr
 862   | isTyCoVar bndr  = substTyVarBndr subst bndr
 863   | otherwise       = subst_opt_id_bndr subst bndr
 864
 865 subst_opt_id_bndr :: Subst -> InId -> (Subst, OutId)
 866 -- Nuke all fragile IdInfo, unfolding, and RULES;
 867 --    it gets added back later by add_info
 868 -- Rather like SimplEnv.substIdBndr
 869 --
 870 -- It's important to zap fragile OccInfo (which CoreSubst.SubstIdBndr
 871 -- carefully does not do) because simplOptExpr invalidates it
 872
 873 subst_opt_id_bndr subst@(Subst in_scope id_subst tv_subst) old_id
 874   = (Subst new_in_scope new_id_subst tv_subst, new_id)
 875   where
 876     id1    = uniqAway in_scope old_id
 877     id2    = setIdType id1 (substTy subst (idType old_id))
 878     new_id = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
 879                                         -- and fragile OccInfo
 880     new_in_scope = in_scope `extendInScopeSet` new_id
 881
 882         -- Extend the substitution if the unique has changed,
 883         -- or there's some useful occurrence information
 884         -- See the notes with substTyVarBndr for the delSubstEnv
 885     new_id_subst | new_id /= old_id
 886                  = extendVarEnv id_subst old_id (Var new_id)
 887                  | otherwise
 888                  = delVarEnv id_subst old_id
 889
 890 ----------------------
 891 subst_opt_bndrs :: Subst -> [InVar] -> (Subst, [OutVar])
 892 subst_opt_bndrs subst bndrs
 893   = mapAccumL subst_opt_bndr subst bndrs
 894
 895 ----------------------
 896 add_info :: Subst -> InVar -> OutVar -> OutVar
 897 add_info subst old_bndr new_bndr
 898  | isTyCoVar old_bndr = new_bndr
 899  | otherwise          = maybeModifyIdInfo mb_new_info new_bndr
 900  where
 901    mb_new_info = substIdInfo subst new_bndr (idInfo old_bndr)
 902 \end{code}
 903
 904 Note [Inline prag in simplOpt]
 905 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 906 If there's an INLINE/NOINLINE pragma that restricts the phase in
 907 which the binder can be inlined, we don't inline here; after all,
 908 we don't know what phase we're in.  Here's an example
 909
 910   foo :: Int -> Int -> Int
 911   {-# INLINE foo #-}
 912   foo m n = inner m
 913      where
 914        {-# INLINE [1] inner #-}
 915        inner m = m+n
 916
 917   bar :: Int -> Int
 918   bar n = foo n 1
 919
 920 When inlining 'foo' in 'bar' we want the let-binding for 'inner'
 921 to remain visible until Phase 1
 922