ghc/compiler/coreSyn/Subst.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[CoreUtils]{Utility functions on @Core@ syntax}
   5
   6 \begin{code}
   7 module Subst (
   8         -- In-scope set
   9         InScopeSet, emptyInScopeSet, mkInScopeSet,
  10         extendInScopeSet, extendInScopeSetList,
  11         lookupInScope, elemInScopeSet, uniqAway,
  12
  13
  14         -- Substitution stuff
  15         Subst, TyVarSubst, IdSubst,
  16         emptySubst, mkSubst, substEnv, substInScope,
  17         lookupSubst, lookupIdSubst, isEmptySubst, extendSubst, extendSubstList,
  18         zapSubstEnv, setSubstEnv,
  19         setInScope,
  20         extendInScope, extendInScopeList, extendNewInScope, extendNewInScopeList,
  21         isInScope, modifyInScope,
  22
  23         bindSubst, unBindSubst, bindSubstList, unBindSubstList,
  24
  25         -- Binders
  26         simplBndr, simplBndrs, simplLetId, simplLamBndr, simplIdInfo,
  27         substAndCloneId, substAndCloneIds, substAndCloneRecIds,
  28
  29         -- Type stuff
  30         mkTyVarSubst, mkTopTyVarSubst,
  31         substTyWith, substTy, substTheta, deShadowTy,
  32
  33         -- Expression stuff
  34         substExpr, substRules
  35     ) where
  36
  37 #include "HsVersions.h"
  38
  39 import CmdLineOpts      ( opt_PprStyle_Debug )
  40 import CoreSyn          ( Expr(..), Bind(..), Note(..), CoreExpr,
  41                           CoreRules(..), CoreRule(..),
  42                           isEmptyCoreRules, seqRules, hasUnfolding, noUnfolding, hasSomeUnfolding,
  43                           Unfolding(..)
  44                         )
  45 import CoreFVs          ( exprFreeVars )
  46 import TypeRep          ( Type(..), TyNote(..) )  -- friend
  47 import Type             ( ThetaType, SourceType(..), PredType,
  48                           tyVarsOfType, tyVarsOfTypes, mkAppTy,
  49                         )
  50 import VarSet
  51 import VarEnv
  52 import Var              ( setVarUnique, isId, mustHaveLocalBinding )
  53 import Id               ( idType, idInfo, setIdInfo, setIdType,
  54                           idUnfolding, setIdUnfolding,
  55                           idOccInfo, maybeModifyIdInfo )
  56 import IdInfo           ( IdInfo, vanillaIdInfo,
  57                           occInfo, isFragileOcc, setOccInfo,
  58                           specInfo, setSpecInfo,
  59                           setArityInfo, unknownArity,
  60                           unfoldingInfo, setUnfoldingInfo,
  61                           WorkerInfo(..), workerExists, workerInfo, setWorkerInfo, WorkerInfo,
  62                           lbvarInfo, LBVarInfo(..), setLBVarInfo, hasNoLBVarInfo
  63                         )
  64 import BasicTypes       ( OccInfo(..) )
  65 import Unique           ( Unique, Uniquable(..), deriveUnique )
  66 import UniqSet          ( elemUniqSet_Directly )
  67 import UniqSupply       ( UniqSupply, uniqFromSupply, uniqsFromSupply )
  68 import Var              ( Var, Id, TyVar, isTyVar )
  69 import Outputable
  70 import PprCore          ()              -- Instances
  71 import UniqFM           ( ufmToList )   -- Yuk (add a new op to VarEnv)
  72 import Util             ( mapAccumL, foldl2, seqList )
  73 import FastTypes
  74 \end{code}
  75
  76
  77 %************************************************************************
  78 %*                                                                      *
  79 \subsection{The in-scope set}
  80 %*                                                                      *
  81 %************************************************************************
  82
  83 \begin{code}
  84 data InScopeSet = InScope (VarEnv Var) FastInt
  85         -- The Int# is a kind of hash-value used by uniqAway
  86         -- For example, it might be the size of the set
  87         -- INVARIANT: it's not zero; we use it as a multiplier in uniqAway
  88
  89 emptyInScopeSet :: InScopeSet
  90 emptyInScopeSet = InScope emptyVarSet 1#
  91
  92 mkInScopeSet :: VarEnv Var -> InScopeSet
  93 mkInScopeSet in_scope = InScope in_scope 1#
  94
  95 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
  96 extendInScopeSet (InScope in_scope n) v = InScope (extendVarEnv in_scope v v) (n +# 1#)
  97
  98 extendInScopeSetList :: InScopeSet -> [Var] -> InScopeSet
  99 extendInScopeSetList (InScope in_scope n) vs
 100    = InScope (foldl (\s v -> extendVarEnv s v v) in_scope vs)
 101                     (n +# iUnbox (length vs))
 102
 103 modifyInScopeSet :: InScopeSet -> Var -> Var -> InScopeSet
 104 -- Exploit the fact that the in-scope "set" is really a map
 105 --      Make old_v map to new_v
 106 modifyInScopeSet (InScope in_scope n) old_v new_v = InScope (extendVarEnv in_scope old_v new_v) (n +# 1#)
 107
 108 delInScopeSet :: InScopeSet -> Var -> InScopeSet
 109 delInScopeSet (InScope in_scope n) v = InScope (in_scope `delVarEnv` v) n
 110
 111 elemInScopeSet :: Var -> InScopeSet -> Bool
 112 elemInScopeSet v (InScope in_scope n) = v `elemVarEnv` in_scope
 113
 114 lookupInScope :: InScopeSet -> Var -> Var
 115 -- It's important to look for a fixed point
 116 -- When we see (case x of y { I# v -> ... })
 117 -- we add  [x -> y] to the in-scope set (Simplify.simplCaseBinder).
 118 -- When we lookup up an occurrence of x, we map to y, but then
 119 -- we want to look up y in case it has acquired more evaluation information by now.
 120 lookupInScope (InScope in_scope n) v
 121   = go v
 122   where
 123     go v = case lookupVarEnv in_scope v of
 124                 Just v' | v == v'   -> v'       -- Reached a fixed point
 125                         | otherwise -> go v'
 126                 Nothing             -> WARN( mustHaveLocalBinding v, ppr v )
 127                                        v
 128 \end{code}
 129
 130 \begin{code}
 131 uniqAway :: InScopeSet -> Var -> Var
 132 -- (uniqAway in_scope v) finds a unique that is not used in the
 133 -- in-scope set, and gives that to v.  It starts with v's current unique, of course,
 134 -- in the hope that it won't have to change it, nad thereafter uses a combination
 135 -- of that and the hash-code found in the in-scope set
 136 uniqAway (InScope set n) var
 137   | not (var `elemVarSet` set) = var                            -- Nothing to do
 138   | otherwise                  = try 1#
 139   where
 140     orig_unique = getUnique var
 141     try k
 142 #ifdef DEBUG
 143           | k ># 1000#
 144           = pprPanic "uniqAway loop:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 145 #endif
 146           | uniq `elemUniqSet_Directly` set = try (k +# 1#)
 147 #ifdef DEBUG
 148           | opt_PprStyle_Debug && k ># 3#
 149           = pprTrace "uniqAway:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n))
 150             setVarUnique var uniq
 151 #endif
 152           | otherwise = setVarUnique var uniq
 153           where
 154             uniq = deriveUnique orig_unique (iBox (n *# k))
 155 \end{code}
 156
 157
 158 %************************************************************************
 159 %*                                                                      *
 160 \subsection{Substitutions}
 161 %*                                                                      *
 162 %************************************************************************
 163
 164 \begin{code}
 165 data Subst = Subst InScopeSet           -- In scope
 166                    SubstEnv             -- Substitution itself
 167         -- INVARIANT 1: The (domain of the) in-scope set is a superset
 168         --              of the free vars of the range of the substitution
 169         --              that might possibly clash with locally-bound variables
 170         --              in the thing being substituted in.
 171         -- This is what lets us deal with name capture properly
 172         -- It's a hard invariant to check...
 173         -- There are various ways of causing it to happen:
 174         --      - arrange that the in-scope set really is all the things in scope
 175         --      - arrange that it's the free vars of the range of the substitution
 176         --      - make it empty because all the free vars of the subst are fresh,
 177         --              and hence can't possibly clash.a
 178         --
 179         -- INVARIANT 2: No variable is both in scope and in the domain of the substitution
 180         --              Equivalently, the substitution is idempotent
 181         --      [Sep 2000: Lies, all lies.  The substitution now does contain
 182         --                 mappings x77 -> DoneId x77 occ
 183         --                 to record x's occurrence information.]
 184         --      [Also watch out: the substitution can contain x77 -> DoneEx (Var x77)
 185         --       Consider let x = case k of I# x77 -> ... in
 186         --                let y = case k of I# x77 -> ... in ...
 187         --       and suppose the body is strict in both x and y.  Then the simplifier
 188         --       will pull the first (case k) to the top; so the second (case k) will
 189         --       cancel out, mapping x77 to, well, x77!  But one is an in-Id and the
 190         --       other is an out-Id. So the substitution is idempotent in the sense
 191         --       that we *must not* repeatedly apply it.]
 192
 193 type IdSubst    = Subst
 194 \end{code}
 195
 196 The general plan about the substitution and in-scope set for Ids is as follows
 197
 198 * substId always adds new_id to the in-scope set.
 199   new_id has a correctly-substituted type, occ info
 200
 201 * substId adds a binding (DoneId new_id occ) to the substitution if
 202         EITHER the Id's unique has changed
 203         OR     the Id has interesting occurrence information
 204   So in effect you can only get to interesting occurrence information
 205   by looking up the *old* Id; it's not really attached to the new id
 206   at all.
 207
 208   Note, though that the substitution isn't necessarily extended
 209   if the type changes.  Why not?  Because of the next point:
 210
 211 * We *always, always* finish by looking up in the in-scope set
 212   any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
 213   Reason: so that we never finish up with a "old" Id in the result.
 214   An old Id might point to an old unfolding and so on... which gives a space leak.
 215
 216   [The DoneEx and DoneVar hits map to "new" stuff.]
 217
 218 * It follows that substExpr must not do a no-op if the substitution is empty.
 219   substType is free to do so, however.
 220
 221 * When we come to a let-binding (say) we generate new IdInfo, including an
 222   unfolding, attach it to the binder, and add this newly adorned binder to
 223   the in-scope set.  So all subsequent occurrences of the binder will get mapped
 224   to the full-adorned binder, which is also the one put in the binding site.
 225
 226 * The in-scope "set" usually maps x->x; we use it simply for its domain.
 227   But sometimes we have two in-scope Ids that are synomyms, and should
 228   map to the same target:  x->x, y->x.  Notably:
 229         case y of x { ... }
 230   That's why the "set" is actually a VarEnv Var
 231
 232
 233 \begin{code}
 234 isEmptySubst :: Subst -> Bool
 235 isEmptySubst (Subst _ env) = isEmptySubstEnv env
 236
 237 emptySubst :: Subst
 238 emptySubst = Subst emptyInScopeSet emptySubstEnv
 239
 240 mkSubst :: InScopeSet -> SubstEnv -> Subst
 241 mkSubst in_scope env = Subst in_scope env
 242
 243 substEnv :: Subst -> SubstEnv
 244 substEnv (Subst _ env) = env
 245
 246 substInScope :: Subst -> InScopeSet
 247 substInScope (Subst in_scope _) = in_scope
 248
 249 zapSubstEnv :: Subst -> Subst
 250 zapSubstEnv (Subst in_scope env) = Subst in_scope emptySubstEnv
 251
 252 extendSubst :: Subst -> Var -> SubstResult -> Subst
 253 extendSubst (Subst in_scope env) v r = Subst in_scope (extendSubstEnv env v r)
 254
 255 extendSubstList :: Subst -> [Var] -> [SubstResult] -> Subst
 256 extendSubstList (Subst in_scope env) v r = Subst in_scope (extendSubstEnvList env v r)
 257
 258 lookupSubst :: Subst -> Var -> Maybe SubstResult
 259 lookupSubst (Subst _ env) v = lookupSubstEnv env v
 260
 261 lookupIdSubst :: Subst -> Id -> SubstResult
 262 -- Does the lookup in the in-scope set too
 263 lookupIdSubst (Subst in_scope env) v
 264   = case lookupSubstEnv env v of
 265         Just (DoneId v' occ) -> DoneId (lookupInScope in_scope v') occ
 266         Just res             -> res
 267         Nothing              -> DoneId v' (idOccInfo v')
 268                                 -- We don't use DoneId for LoopBreakers, so the idOccInfo is
 269                                 -- very important!  If isFragileOcc returned True for
 270                                 -- loop breakers we could avoid this call, but at the expense
 271                                 -- of adding more to the substitution, and building new Ids
 272                                 -- in substId a bit more often than really necessary
 273                              where
 274                                     v' = lookupInScope in_scope v
 275
 276 isInScope :: Var -> Subst -> Bool
 277 isInScope v (Subst in_scope _) = v `elemInScopeSet` in_scope
 278
 279 modifyInScope :: Subst -> Var -> Var -> Subst
 280 modifyInScope (Subst in_scope env) old_v new_v = Subst (modifyInScopeSet in_scope old_v new_v) env
 281         -- make old_v map to new_v
 282
 283 extendInScope :: Subst -> Var -> Subst
 284         -- Add a new variable as in-scope
 285         -- Remember to delete any existing binding in the substitution!
 286 extendInScope (Subst in_scope env) v = Subst (in_scope `extendInScopeSet` v)
 287                                              (env `delSubstEnv` v)
 288
 289 extendInScopeList :: Subst -> [Var] -> Subst
 290 extendInScopeList (Subst in_scope env) vs = Subst (extendInScopeSetList in_scope vs)
 291                                                   (delSubstEnvList env vs)
 292
 293 -- The "New" variants are guaranteed to be adding freshly-allocated variables
 294 -- It's not clear that the gain (not needing to delete it from the substitution)
 295 -- is worth the extra proof obligation
 296 extendNewInScope :: Subst -> Var -> Subst
 297 extendNewInScope (Subst in_scope env) v = Subst (in_scope `extendInScopeSet` v) env
 298
 299 extendNewInScopeList :: Subst -> [Var] -> Subst
 300 extendNewInScopeList (Subst in_scope env) vs = Subst (in_scope `extendInScopeSetList` vs) env
 301
 302 -------------------------------
 303 bindSubst :: Subst -> Var -> Var -> Subst
 304 -- Extend with a substitution, v1 -> Var v2
 305 -- and extend the in-scopes with v2
 306 bindSubst (Subst in_scope env) old_bndr new_bndr
 307   = Subst (in_scope `extendInScopeSet` new_bndr)
 308           (extendSubstEnv env old_bndr subst_result)
 309   where
 310     subst_result | isId old_bndr = DoneEx (Var new_bndr)
 311                  | otherwise     = DoneTy (TyVarTy new_bndr)
 312
 313 unBindSubst :: Subst -> Var -> Var -> Subst
 314 -- Reverse the effect of bindSubst
 315 -- If old_bndr was already in the substitution, this doesn't quite work
 316 unBindSubst (Subst in_scope env) old_bndr new_bndr
 317   = Subst (in_scope `delInScopeSet` new_bndr) (delSubstEnv env old_bndr)
 318
 319 -- And the "List" forms
 320 bindSubstList :: Subst -> [Var] -> [Var] -> Subst
 321 bindSubstList subst old_bndrs new_bndrs
 322   = foldl2 bindSubst subst old_bndrs new_bndrs
 323
 324 unBindSubstList :: Subst -> [Var] -> [Var] -> Subst
 325 unBindSubstList subst old_bndrs new_bndrs
 326   = foldl2 unBindSubst subst old_bndrs new_bndrs
 327
 328
 329 -------------------------------
 330 setInScope :: Subst     -- Take env part from here
 331            -> InScopeSet
 332            -> Subst
 333 setInScope (Subst in_scope1 env1) in_scope2
 334   = Subst in_scope2 env1
 335
 336 setSubstEnv :: Subst            -- Take in-scope part from here
 337             -> SubstEnv         -- ... and env part from here
 338             -> Subst
 339 setSubstEnv (Subst in_scope1 _) env2 = Subst in_scope1 env2
 340 \end{code}
 341
 342 Pretty printing, for debugging only
 343
 344 \begin{code}
 345 instance Outputable SubstResult where
 346   ppr (DoneEx e)   = ptext SLIT("DoneEx") <+> ppr e
 347   ppr (DoneId v _) = ptext SLIT("DoneId") <+> ppr v
 348   ppr (ContEx _ e) = ptext SLIT("ContEx") <+> ppr e
 349   ppr (DoneTy t)   = ptext SLIT("DoneTy") <+> ppr t
 350
 351 instance Outputable SubstEnv where
 352   ppr se = brackets (fsep (punctuate comma (map ppr_elt (ufmToList (substEnvEnv se)))))
 353         where
 354            ppr_elt (uniq,sr) = ppr uniq <+> ptext SLIT("->") <+> ppr sr
 355
 356 instance Outputable Subst where
 357   ppr (Subst (InScope in_scope _) se)
 358         =  ptext SLIT("<InScope =") <+> braces   (fsep (map ppr (rngVarEnv in_scope)))
 359         $$ ptext SLIT(" Subst   =") <+> ppr se <> char '>'
 360 \end{code}
 361
 362 %************************************************************************
 363 %*                                                                      *
 364 \subsection{Type substitution}
 365 %*                                                                      *
 366 %************************************************************************
 367
 368 \begin{code}
 369 type TyVarSubst = Subst -- TyVarSubst are expected to have range elements
 370         -- (We could have a variant of Subst, but it doesn't seem worth it.)
 371
 372 -- mkTyVarSubst generates the in-scope set from
 373 -- the types given; but it's just a thunk so with a bit of luck
 374 -- it'll never be evaluated
 375 mkTyVarSubst :: [TyVar] -> [Type] -> Subst
 376 mkTyVarSubst tyvars tys = Subst (mkInScopeSet (tyVarsOfTypes tys))
 377                                 (zipTyEnv tyvars tys)
 378
 379 -- mkTopTyVarSubst is called when doing top-level substitutions.
 380 -- Here we expect that the free vars of the range of the
 381 -- substitution will be empty.
 382 mkTopTyVarSubst :: [TyVar] -> [Type] -> Subst
 383 mkTopTyVarSubst tyvars tys = Subst emptyInScopeSet (zipTyEnv tyvars tys)
 384
 385 zipTyEnv tyvars tys
 386 #ifdef DEBUG
 387   | length tyvars /= length tys
 388   = pprTrace "mkTopTyVarSubst" (ppr tyvars $$ ppr tys) emptySubstEnv
 389   | otherwise
 390 #endif
 391   = zip_ty_env tyvars tys emptySubstEnv
 392
 393 zip_ty_env []       []       env = env
 394 zip_ty_env (tv:tvs) (ty:tys) env = zip_ty_env tvs tys (extendSubstEnv env tv (DoneTy ty))
 395         -- There used to be a special case for when
 396         --      ty == TyVarTy tv
 397         -- (a not-uncommon case) in which case the substitution was dropped.
 398         -- But the type-tidier changes the print-name of a type variable without
 399         -- changing the unique, and that led to a bug.   Why?  Pre-tidying, we had
 400         -- a type {Foo t}, where Foo is a one-method class.  So Foo is really a newtype.
 401         -- And it happened that t was the type variable of the class.  Post-tiding,
 402         -- it got turned into {Foo t2}.  The ext-core printer expanded this using
 403         -- sourceTypeRep, but that said "Oh, t == t2" because they have the same unique,
 404         -- and so generated a rep type mentioning t not t2.
 405         --
 406         -- Simplest fix is to nuke the "optimisation"
 407 \end{code}
 408
 409 substTy works with general Substs, so that it can be called from substExpr too.
 410
 411 \begin{code}
 412 substTyWith :: [TyVar] -> [Type] -> Type -> Type
 413 substTyWith tvs tys = substTy (mkTyVarSubst tvs tys)
 414
 415 substTy :: Subst -> Type  -> Type
 416 substTy subst ty | isEmptySubst subst = ty
 417                  | otherwise          = subst_ty subst ty
 418
 419 deShadowTy :: Type -> Type              -- Remove any shadowing from the type
 420 deShadowTy ty = subst_ty emptySubst ty
 421
 422 substTheta :: TyVarSubst -> ThetaType -> ThetaType
 423 substTheta subst theta
 424   | isEmptySubst subst = theta
 425   | otherwise          = map (substPred subst) theta
 426
 427 substPred :: TyVarSubst -> PredType -> PredType
 428 substPred = substSourceType
 429
 430 substSourceType subst (IParam n ty)     = IParam n (subst_ty subst ty)
 431 substSourceType subst (ClassP clas tys) = ClassP clas (map (subst_ty subst) tys)
 432 substSourceType subst (NType  tc   tys) = NType  tc   (map (subst_ty subst) tys)
 433
 434 subst_ty subst ty
 435    = go ty
 436   where
 437     go (TyConApp tc tys)           = let args = map go tys
 438                                      in  args `seqList` TyConApp tc args
 439
 440     go (SourceTy p)                = SourceTy $! (substSourceType subst p)
 441
 442     go (NoteTy (SynNote ty1) ty2)  = NoteTy (SynNote $! (go ty1)) $! (go ty2)
 443     go (NoteTy (FTVNote _) ty2)    = go ty2             -- Discard the free tyvar note
 444
 445     go (FunTy arg res)             = (FunTy $! (go arg)) $! (go res)
 446     go (AppTy fun arg)             = mkAppTy (go fun) $! (go arg)
 447     go ty@(TyVarTy tv)             = case (lookupSubst subst tv) of
 448                                         Nothing            -> ty
 449                                         Just (DoneTy ty')  -> ty'
 450
 451     go (ForAllTy tv ty)            = case substTyVar subst tv of
 452                                         (subst', tv') -> ForAllTy tv' $! (subst_ty subst' ty)
 453 \end{code}
 454
 455 Here is where we invent a new binder if necessary.
 456
 457 \begin{code}
 458 substTyVar :: Subst -> TyVar -> (Subst, TyVar)
 459 substTyVar subst@(Subst in_scope env) old_var
 460   | old_var == new_var  -- No need to clone
 461                         -- But we *must* zap any current substitution for the variable.
 462                         --  For example:
 463                         --      (\x.e) with id_subst = [x |-> e']
 464                         -- Here we must simply zap the substitution for x
 465                         --
 466                         -- The new_id isn't cloned, but it may have a different type
 467                         -- etc, so we must return it, not the old id
 468   = (Subst (in_scope `extendInScopeSet` new_var)
 469            (delSubstEnv env old_var),
 470      new_var)
 471
 472   | otherwise   -- The new binder is in scope so
 473                 -- we'd better rename it away from the in-scope variables
 474                 -- Extending the substitution to do this renaming also
 475                 -- has the (correct) effect of discarding any existing
 476                 -- substitution for that variable
 477   = (Subst (in_scope `extendInScopeSet` new_var)
 478            (extendSubstEnv env old_var (DoneTy (TyVarTy new_var))),
 479      new_var)
 480   where
 481     new_var = uniqAway in_scope old_var
 482         -- The uniqAway part makes sure the new variable is not already in scope
 483 \end{code}
 484
 485
 486 %************************************************************************
 487 %*                                                                      *
 488 \section{Expression substitution}
 489 %*                                                                      *
 490 %************************************************************************
 491
 492 This expression substituter deals correctly with name capture.
 493
 494 BUT NOTE that substExpr silently discards the
 495         unfolding, and
 496         spec env
 497 IdInfo attached to any binders in the expression.  It's quite
 498 tricky to do them 'right' in the case of mutually recursive bindings,
 499 and so far has proved unnecessary.
 500
 501 \begin{code}
 502 substExpr :: Subst -> CoreExpr -> CoreExpr
 503 substExpr subst expr
 504         -- NB: we do not do a no-op when the substitution is empty,
 505         -- because we always want to substitute the variables in the
 506         -- in-scope set for their occurrences.  Why?
 507         --      (a) because they may contain more information
 508         --      (b) because leaving an un-substituted Id might cause
 509         --          a space leak (its unfolding might point to an old version
 510         --          of its right hand side).
 511
 512   = go expr
 513   where
 514     go (Var v) = -- See the notes at the top, with the Subst data type declaration
 515                  case lookupIdSubst subst v of
 516
 517                     ContEx env' e' -> substExpr (setSubstEnv subst env') e'
 518                     DoneId v _     -> Var v
 519                     DoneEx e'      -> e'
 520
 521     go (Type ty)      = Type (go_ty ty)
 522     go (Lit lit)      = Lit lit
 523     go (App fun arg)  = App (go fun) (go arg)
 524     go (Note note e)  = Note (go_note note) (go e)
 525
 526     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
 527                        where
 528                          (subst', bndr') = substBndr subst bndr
 529
 530     go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
 531                                     where
 532                                       (subst', bndr') = substBndr subst bndr
 533
 534     go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
 535                               where
 536                                 (subst', bndrs') = substRecIds subst (map fst pairs)
 537                                 pairs'  = bndrs' `zip` rhss'
 538                                 rhss'   = map (substExpr subst' . snd) pairs
 539
 540     go (Case scrut bndr alts) = Case (go scrut) bndr' (map (go_alt subst') alts)
 541                               where
 542                                 (subst', bndr') = substBndr subst bndr
 543
 544     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
 545                                  where
 546                                    (subst', bndrs') = substBndrs subst bndrs
 547
 548     go_note (Coerce ty1 ty2) = Coerce (go_ty ty1) (go_ty ty2)
 549     go_note note             = note
 550
 551     go_ty ty = substTy subst ty
 552
 553 \end{code}
 554
 555
 556 %************************************************************************
 557 %*                                                                      *
 558 \section{Substituting an Id binder}
 559 %*                                                                      *
 560 %************************************************************************
 561
 562 \begin{code}
 563 -- simplBndr and simplLetId are used by the simplifier
 564
 565 simplBndr :: Subst -> Var -> (Subst, Var)
 566 -- Used for lambda and case-bound variables
 567 -- Clone Id if necessary, substitute type
 568 -- Return with IdInfo already substituted, but (fragile) occurrence info zapped
 569 -- The substitution is extended only if the variable is cloned, because
 570 -- we *don't* need to use it to track occurrence info.
 571 simplBndr subst bndr
 572   | isTyVar bndr  = substTyVar subst bndr
 573   | otherwise     = subst_id isFragileOcc subst subst bndr
 574
 575 simplBndrs :: Subst -> [Var] -> (Subst, [Var])
 576 simplBndrs subst bndrs = mapAccumL simplBndr subst bndrs
 577
 578 simplLamBndr :: Subst -> Var -> (Subst, Var)
 579 -- Used for lambda binders.  These sometimes have unfoldings added by
 580 -- the worker/wrapper pass that must be preserved, becuase they can't
 581 -- be reconstructed from context.  For example:
 582 --      f x = case x of (a,b) -> fw a b x
 583 --      fw a b x{=(a,b)} = ...
 584 -- The "{=(a,b)}" is an unfolding we can't reconstruct otherwise.
 585 simplLamBndr subst bndr
 586   | not (isId bndr && hasSomeUnfolding old_unf)
 587   = simplBndr subst bndr        -- Normal case
 588   | otherwise
 589   = (subst', bndr' `setIdUnfolding` substUnfolding subst old_unf)
 590   where
 591     old_unf = idUnfolding bndr
 592     (subst', bndr') = subst_id isFragileOcc subst subst bndr
 593
 594
 595 simplLetId :: Subst -> Id -> (Subst, Id)
 596 -- Clone Id if necessary
 597 -- Substitute its type
 598 -- Return an Id with completely zapped IdInfo
 599 --      [A subsequent substIdInfo will restore its IdInfo]
 600 -- Augment the subtitution
 601 --      if the unique changed, *or*
 602 --      if there's interesting occurrence info
 603
 604 simplLetId subst@(Subst in_scope env) old_id
 605   = (Subst (in_scope `extendInScopeSet` new_id) new_env, new_id)
 606   where
 607     old_info = idInfo old_id
 608     id1     = uniqAway in_scope old_id
 609     id2     = substIdType subst id1
 610     new_id  = setIdInfo id2 vanillaIdInfo
 611
 612         -- Extend the substitution if the unique has changed,
 613         -- or there's some useful occurrence information
 614         -- See the notes with substTyVar for the delSubstEnv
 615     occ_info = occInfo old_info
 616     new_env | new_id /= old_id || isFragileOcc occ_info
 617             = extendSubstEnv env old_id (DoneId new_id occ_info)
 618             | otherwise
 619             = delSubstEnv env old_id
 620
 621 simplIdInfo :: Subst -> IdInfo -> IdInfo
 622   -- Used by the simplifier to compute new IdInfo for a let(rec) binder,
 623   -- subsequent to simplLetId having zapped its IdInfo
 624 simplIdInfo subst old_info
 625   = case substIdInfo subst isFragileOcc zapped_old_info of
 626         Just new_info -> new_info
 627         Nothing       -> old_info
 628   where
 629     zapped_old_info = old_info `setArityInfo` unknownArity
 630         -- Like unfolding, arity gets set later
 631         -- Maybe we should do this in substIdInfo?
 632 \end{code}
 633
 634 \begin{code}
 635 -- substBndr and friends are used when doing expression substitution only
 636 -- In this case we can *preserve* occurrence information, and indeed we *want*
 637 -- to do so else lose useful occ info in rules.  Hence the calls to
 638 -- simpl_id with keepOccInfo
 639
 640 substBndr :: Subst -> Var -> (Subst, Var)
 641 substBndr subst bndr
 642   | isTyVar bndr  = substTyVar subst bndr
 643   | otherwise     = subst_id keepOccInfo subst subst bndr
 644
 645 substBndrs :: Subst -> [Var] -> (Subst, [Var])
 646 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 647
 648 substRecIds :: Subst -> [Id] -> (Subst, [Id])
 649 -- Substitute a mutually recursive group
 650 substRecIds subst bndrs
 651   = (new_subst, new_bndrs)
 652   where
 653         -- Here's the reason we need to pass rec_subst to subst_id
 654     (new_subst, new_bndrs) = mapAccumL (subst_id keepOccInfo new_subst) subst bndrs
 655
 656 keepOccInfo occ = False -- Never fragile
 657 \end{code}
 658
 659
 660 \begin{code}
 661 subst_id :: (OccInfo -> Bool)   -- True <=> the OccInfo is fragile
 662          -> Subst               -- Substitution to use for the IdInfo
 663          -> Subst -> Id         -- Substitition and Id to transform
 664          -> (Subst, Id)         -- Transformed pair
 665
 666 -- Returns with:
 667 --      * Unique changed if necessary
 668 --      * Type substituted
 669 --      * Unfolding zapped
 670 --      * Rules, worker, lbvar info all substituted
 671 --      * Occurrence info zapped if is_fragile_occ returns True
 672 --      * The in-scope set extended with the returned Id
 673 --      * The substitution extended with a DoneId if unique changed
 674 --        In this case, the var in the DoneId is the same as the
 675 --        var returned
 676
 677 subst_id is_fragile_occ rec_subst subst@(Subst in_scope env) old_id
 678   = (Subst (in_scope `extendInScopeSet` new_id) new_env, new_id)
 679   where
 680         -- id1 is cloned if necessary
 681     id1 = uniqAway in_scope old_id
 682
 683         -- id2 has its type zapped
 684     id2 = substIdType subst id1
 685
 686         -- new_id has the right IdInfo
 687         -- The lazy-set is because we're in a loop here, with
 688         -- rec_subst, when dealing with a mutually-recursive group
 689     new_id = maybeModifyIdInfo (substIdInfo rec_subst is_fragile_occ) id2
 690
 691         -- Extend the substitution if the unique has changed
 692         -- See the notes with substTyVar for the delSubstEnv
 693     new_env | new_id /= old_id
 694             = extendSubstEnv env old_id (DoneId new_id (idOccInfo old_id))
 695             | otherwise
 696             = delSubstEnv env old_id
 697 \end{code}
 698
 699 Now a variant that unconditionally allocates a new unique.
 700 It also unconditionally zaps the OccInfo.
 701
 702 \begin{code}
 703 subst_clone_id :: Subst                 -- Substitution to use (lazily) for the rules and worker
 704                -> Subst -> (Id, Unique) -- Substitition and Id to transform
 705                -> (Subst, Id)           -- Transformed pair
 706
 707 subst_clone_id rec_subst subst@(Subst in_scope env) (old_id, uniq)
 708   = (Subst (in_scope `extendInScopeSet` new_id) new_env, new_id)
 709   where
 710     id1  = setVarUnique old_id uniq
 711     id2  = substIdType subst id1
 712
 713     new_id  = maybeModifyIdInfo (substIdInfo rec_subst isFragileOcc) id2
 714     new_env = extendSubstEnv env old_id (DoneId new_id NoOccInfo)
 715
 716 substAndCloneIds :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 717 substAndCloneIds subst us ids
 718   = mapAccumL (subst_clone_id subst) subst (ids `zip` uniqsFromSupply us)
 719
 720 substAndCloneRecIds :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 721 substAndCloneRecIds subst us ids
 722   = (subst', ids')
 723   where
 724     (subst', ids') = mapAccumL (subst_clone_id subst') subst
 725                                (ids `zip` uniqsFromSupply us)
 726
 727 substAndCloneId :: Subst -> UniqSupply -> Id -> (Subst, Id)
 728 substAndCloneId subst@(Subst in_scope env) us old_id
 729   = subst_clone_id subst subst (old_id, uniqFromSupply us)
 730 \end{code}
 731
 732
 733 %************************************************************************
 734 %*                                                                      *
 735 \section{IdInfo substitution}
 736 %*                                                                      *
 737 %************************************************************************
 738
 739 \begin{code}
 740 substIdInfo :: Subst
 741             -> (OccInfo -> Bool)        -- True <=> zap the occurrence info
 742             -> IdInfo
 743             -> Maybe IdInfo
 744 -- Substitute the
 745 --      rules
 746 --      worker info
 747 --      LBVar info
 748 -- Zap the unfolding
 749 -- Zap the occ info if instructed to do so
 750 --
 751 -- Seq'ing on the returned IdInfo is enough to cause all the
 752 -- substitutions to happen completely
 753
 754 substIdInfo subst is_fragile_occ info
 755   | nothing_to_do = Nothing
 756   | otherwise     = Just (info `setOccInfo`       (if zap_occ then NoOccInfo else old_occ)
 757                                `setSpecInfo`      substRules  subst old_rules
 758                                `setWorkerInfo`    substWorker subst old_wrkr
 759                                `setUnfoldingInfo` noUnfolding)
 760                         -- setSpecInfo does a seq
 761                         -- setWorkerInfo does a seq
 762   where
 763     nothing_to_do = not zap_occ &&
 764                     isEmptyCoreRules old_rules &&
 765                     not (workerExists old_wrkr) &&
 766                     not (hasUnfolding (unfoldingInfo info))
 767
 768     zap_occ   = is_fragile_occ old_occ
 769     old_occ   = occInfo info
 770     old_rules = specInfo info
 771     old_wrkr  = workerInfo info
 772
 773 ------------------
 774 substIdType :: Subst -> Id -> Id
 775 substIdType subst@(Subst in_scope env) id
 776   |  noTypeSubst env || isEmptyVarSet (tyVarsOfType old_ty) = id
 777   | otherwise                                               = setIdType id (substTy subst old_ty)
 778                 -- The tyVarsOfType is cheaper than it looks
 779                 -- because we cache the free tyvars of the type
 780                 -- in a Note in the id's type itself
 781   where
 782     old_ty = idType id
 783
 784 ------------------
 785 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 786         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 787         -- substitutions to happen completely
 788
 789 substWorker subst NoWorker
 790   = NoWorker
 791 substWorker subst (HasWorker w a)
 792   = case lookupIdSubst subst w of
 793         (DoneId w1 _)     -> HasWorker w1 a
 794         (DoneEx (Var w1)) -> HasWorker w1 a
 795         (DoneEx other)    -> WARN( True, text "substWorker: DoneEx" <+> ppr w )
 796                                   NoWorker      -- Worker has got substituted away altogether
 797         (ContEx se1 e)    -> WARN( True, text "substWorker: ContEx" <+> ppr w <+> ppr e)
 798                                   NoWorker      -- Ditto
 799
 800 ------------------
 801 substUnfolding subst NoUnfolding                 = NoUnfolding
 802 substUnfolding subst (OtherCon cons)             = OtherCon cons
 803 substUnfolding subst (CompulsoryUnfolding rhs)   = CompulsoryUnfolding (substExpr subst rhs)
 804 substUnfolding subst (CoreUnfolding rhs t v w g) = CoreUnfolding (substExpr subst rhs) t v w g
 805
 806 ------------------
 807 substRules :: Subst -> CoreRules -> CoreRules
 808         -- Seq'ing on the returned CoreRules is enough to cause all the
 809         -- substitutions to happen completely
 810
 811 substRules subst rules
 812  | isEmptySubst subst = rules
 813
 814 substRules subst (Rules rules rhs_fvs)
 815   = seqRules new_rules `seq` new_rules
 816   where
 817     new_rules = Rules (map do_subst rules) (substVarSet subst rhs_fvs)
 818
 819     do_subst rule@(BuiltinRule _ _) = rule
 820     do_subst (Rule name act tpl_vars lhs_args rhs)
 821         = Rule name act tpl_vars'
 822                (map (substExpr subst') lhs_args)
 823                (substExpr subst' rhs)
 824         where
 825           (subst', tpl_vars') = substBndrs subst tpl_vars
 826
 827 ------------------
 828 substVarSet subst fvs
 829   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
 830   where
 831     subst_fv subst fv = case lookupIdSubst subst fv of
 832                             DoneId fv' _    -> unitVarSet fv'
 833                             DoneEx expr     -> exprFreeVars expr
 834                             DoneTy ty       -> tyVarsOfType ty
 835                             ContEx se' expr -> substVarSet (setSubstEnv subst se') (exprFreeVars expr)
 836 \end{code}