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