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         -- Substitution stuff
  11         Subst, TvSubstEnv, IdSubstEnv, InScopeSet,
  12
  13         deShadowBinds,
  14         substTy, substExpr, substSpec, substWorker,
  15         lookupIdSubst, lookupTvSubst,
  16
  17         emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst,
  18         extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
  19         extendInScope, extendInScopeIds,
  20         isInScope,
  21
  22         -- Binders
  23         substBndr, substBndrs, substRecBndrs,
  24         cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
  25     ) where
  26
  27 #include "HsVersions.h"
  28
  29 import CoreSyn
  30 import CoreFVs
  31 import CoreUtils
  32
  33 import qualified Type
  34 import Type     ( Type, TvSubst(..), TvSubstEnv )
  35 import VarSet
  36 import VarEnv
  37 import Id
  38 import Var      ( Var, TyVar, setVarUnique )
  39 import IdInfo
  40 import Unique
  41 import UniqSupply
  42 import Maybes
  43 import Outputable
  44 import PprCore          ()              -- Instances
  45 import Util
  46 import FastTypes
  47 \end{code}
  48
  49
  50 %************************************************************************
  51 %*                                                                      *
  52 \subsection{Substitutions}
  53 %*                                                                      *
  54 %************************************************************************
  55
  56 \begin{code}
  57 data Subst
  58   = Subst InScopeSet    -- Variables in in scope (both Ids and TyVars)
  59           IdSubstEnv    -- Substitution for Ids
  60           TvSubstEnv    -- Substitution for TyVars
  61
  62         -- INVARIANT 1: The (domain of the) in-scope set is a superset
  63         --              of the free vars of the range of the substitution
  64         --              that might possibly clash with locally-bound variables
  65         --              in the thing being substituted in.
  66         -- This is what lets us deal with name capture properly
  67         -- It's a hard invariant to check...
  68         -- There are various ways of causing it to happen:
  69         --      - arrange that the in-scope set really is all the things in scope
  70         --      - arrange that it's the free vars of the range of the substitution
  71         --      - make it empty because all the free vars of the subst are fresh,
  72         --              and hence can't possibly clash.a
  73         --
  74         -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
  75         --              Types.TvSubstEnv
  76         --
  77         -- INVARIANT 3: See Note [Extending the Subst]
  78
  79 {-
  80 Note [Extending the Subst]
  81 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  82 For a core Subst, which binds Ids as well, we make a different choice for Ids
  83 than we do for TyVars.
  84
  85 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
  86
  87 For Ids, we have a different invariant
  88         The IdSubstEnv is extended *only* when the Unique on an Id changes
  89         Otherwise, we just extend the InScopeSet
  90
  91 In consequence:
  92
  93 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
  94
  95 * If the TvSubstEnv and IdSubstEnv are both empty, substExpr does nothing
  96   (Note that the above rule for substIdBndr maintains this property.  If
  97    the incoming envts are both empty, then substituting the type and
  98    IdInfo can't change anything.)
  99
 100 * In lookupIdSubst, we *must* look up the Id in the in-scope set, because
 101   it may contain non-trivial changes.  Example:
 102         (/\a. \x:a. ...x...) Int
 103   We extend the TvSubstEnv with [a |-> Int]; but x's unique does not change
 104   so we only extend the in-scope set.  Then we must look up in the in-scope
 105   set when we find the occurrence of x.
 106
 107 Why do we make a different choice for the IdSubstEnv than the TvSubstEnv?
 108
 109 * For Ids, we change the IdInfo all the time (e.g. deleting the
 110   unfolding), and adding it back later, so using the TyVar convention
 111   would entail extending the substitution almost all the time
 112
 113 * The simplifier wants to look up in the in-scope set anyway, in case it
 114   can see a better unfolding from an enclosing case expression
 115
 116 * For TyVars, only coercion variables can possibly change, and they are
 117   easy to spot
 118 -}
 119
 120 type IdSubstEnv = IdEnv CoreExpr
 121
 122 ----------------------------
 123 isEmptySubst :: Subst -> Bool
 124 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
 125
 126 emptySubst :: Subst
 127 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
 128
 129 mkEmptySubst :: InScopeSet -> Subst
 130 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
 131
 132 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
 133 mkSubst in_scope tvs ids = Subst in_scope ids tvs
 134
 135 -- getTvSubst :: Subst -> TvSubst
 136 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
 137
 138 -- getTvSubstEnv :: Subst -> TvSubstEnv
 139 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
 140 --
 141 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
 142 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
 143
 144 substInScope :: Subst -> InScopeSet
 145 substInScope (Subst in_scope _ _) = in_scope
 146
 147 -- zapSubstEnv :: Subst -> Subst
 148 -- zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
 149
 150 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
 151 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
 152 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
 153
 154 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
 155 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
 156
 157 extendTvSubst :: Subst -> TyVar -> Type -> Subst
 158 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
 159
 160 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
 161 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
 162
 163 lookupIdSubst :: Subst -> Id -> CoreExpr
 164 lookupIdSubst (Subst in_scope ids tvs) v
 165   | not (isLocalId v) = Var v
 166   | Just e  <- lookupVarEnv ids       v = e
 167   | Just v' <- lookupInScope in_scope v = Var v'
 168         -- Vital! See Note [Extending the Subst]
 169   | otherwise = WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v )
 170                 Var v
 171
 172 lookupTvSubst :: Subst -> TyVar -> Type
 173 lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
 174
 175 ------------------------------
 176 isInScope :: Var -> Subst -> Bool
 177 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
 178
 179 extendInScope :: Subst -> Var -> Subst
 180 extendInScope (Subst in_scope ids tvs) v
 181   = Subst (in_scope `extendInScopeSet` v)
 182           (ids `delVarEnv` v) (tvs `delVarEnv` v)
 183
 184 extendInScopeIds :: Subst -> [Id] -> Subst
 185 extendInScopeIds (Subst in_scope ids tvs) vs
 186   = Subst (in_scope `extendInScopeSetList` vs)
 187           (ids `delVarEnvList` vs) tvs
 188 \end{code}
 189
 190 Pretty printing, for debugging only
 191
 192 \begin{code}
 193 instance Outputable Subst where
 194   ppr (Subst in_scope ids tvs)
 195         =  ptext SLIT("<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
 196         $$ ptext SLIT(" IdSubst   =") <+> ppr ids
 197         $$ ptext SLIT(" TvSubst   =") <+> ppr tvs
 198          <> char '>'
 199 \end{code}
 200
 201
 202 %************************************************************************
 203 %*                                                                      *
 204         Substituting expressions
 205 %*                                                                      *
 206 %************************************************************************
 207
 208 \begin{code}
 209 substExpr :: Subst -> CoreExpr -> CoreExpr
 210 substExpr subst expr
 211   = go expr
 212   where
 213     go (Var v)         = lookupIdSubst subst v
 214     go (Type ty)       = Type (substTy subst ty)
 215     go (Lit lit)       = Lit lit
 216     go (App fun arg)   = App (go fun) (go arg)
 217     go (Note note e)   = Note (go_note note) (go e)
 218     go (Cast e co)     = Cast (go e) (substTy subst co)
 219     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
 220                        where
 221                          (subst', bndr') = substBndr subst bndr
 222
 223     go (Let bind body) = Let bind' (substExpr subst' body)
 224                        where
 225                          (subst', bind') = substBind subst bind
 226
 227     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
 228                                  where
 229                                  (subst', bndr') = substBndr subst bndr
 230
 231     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
 232                                  where
 233                                    (subst', bndrs') = substBndrs subst bndrs
 234
 235     go_note note             = note
 236
 237 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
 238 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
 239                                   where
 240                                     (subst', bndr') = substBndr subst bndr
 241
 242 substBind subst (Rec pairs) = (subst', Rec pairs')
 243                             where
 244                                 (subst', bndrs') = substRecBndrs subst (map fst pairs)
 245                                 pairs'  = bndrs' `zip` rhss'
 246                                 rhss'   = map (substExpr subst' . snd) pairs
 247 \end{code}
 248
 249 De-shadowing the program is sometimes a useful pre-pass.  It can be done simply
 250 by running over the bindings with an empty substitution, becuase substitution
 251 returns a result that has no-shadowing guaranteed.
 252
 253 (Actually, within a single *type* there might still be shadowing, because
 254 substType is a no-op for the empty substitution, but that's OK.)
 255
 256 \begin{code}
 257 deShadowBinds :: [CoreBind] -> [CoreBind]
 258 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
 259 \end{code}
 260
 261
 262 %************************************************************************
 263 %*                                                                      *
 264         Substituting binders
 265 %*                                                                      *
 266 %************************************************************************
 267
 268 Remember that substBndr and friends are used when doing expression
 269 substitution only.  Their only business is substitution, so they
 270 preserve all IdInfo (suitably substituted).  For example, we *want* to
 271 preserve occ info in rules.
 272
 273 \begin{code}
 274 substBndr :: Subst -> Var -> (Subst, Var)
 275 substBndr subst bndr
 276   | isTyVar bndr  = substTyVarBndr subst bndr
 277   | otherwise     = substIdBndr subst subst bndr
 278
 279 substBndrs :: Subst -> [Var] -> (Subst, [Var])
 280 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 281
 282 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
 283 -- Substitute a mutually recursive group
 284 substRecBndrs subst bndrs
 285   = (new_subst, new_bndrs)
 286   where         -- Here's the reason we need to pass rec_subst to subst_id
 287     (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
 288 \end{code}
 289
 290
 291 \begin{code}
 292 substIdBndr :: Subst            -- Substitution to use for the IdInfo
 293             -> Subst -> Id      -- Substitition and Id to transform
 294             -> (Subst, Id)      -- Transformed pair
 295
 296 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
 297   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 298   where
 299     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
 300     id2 | no_type_change = id1
 301         | otherwise      = setIdType id1 (substTy subst old_ty)
 302
 303     old_ty = idType old_id
 304     no_type_change = isEmptyVarEnv tvs ||
 305                      isEmptyVarSet (Type.tyVarsOfType old_ty)
 306
 307         -- new_id has the right IdInfo
 308         -- The lazy-set is because we're in a loop here, with
 309         -- rec_subst, when dealing with a mutually-recursive group
 310     new_id = maybeModifyIdInfo mb_new_info id2
 311     mb_new_info = substIdInfo rec_subst (idInfo id2)
 312
 313         -- Extend the substitution if the unique has changed
 314         -- See the notes with substTyVarBndr for the delVarEnv
 315     new_env | no_change = delVarEnv env old_id
 316             | otherwise = extendVarEnv env old_id (Var new_id)
 317
 318     no_change = id1 == old_id
 319         -- See Note [Extending the Subst]
 320         -- *not* necessary to check mb_new_info and no_type_change
 321 \end{code}
 322
 323 Now a variant that unconditionally allocates a new unique.
 324 It also unconditionally zaps the OccInfo.
 325
 326 \begin{code}
 327 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
 328 cloneIdBndr subst us old_id
 329   = clone_id subst subst (old_id, uniqFromSupply us)
 330
 331 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 332 cloneIdBndrs subst us ids
 333   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
 334
 335 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 336 cloneRecIdBndrs subst us ids
 337   = (subst', ids')
 338   where
 339     (subst', ids') = mapAccumL (clone_id subst') subst
 340                                (ids `zip` uniqsFromSupply us)
 341
 342 -- Just like substIdBndr, except that it always makes a new unique
 343 -- It is given the unique to use
 344 clone_id    :: Subst                    -- Substitution for the IdInfo
 345             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
 346             -> (Subst, Id)              -- Transformed pair
 347
 348 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
 349   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 350   where
 351     id1     = setVarUnique old_id uniq
 352     id2     = substIdType subst id1
 353     new_id  = maybeModifyIdInfo (substIdInfo rec_subst (idInfo old_id)) id2
 354     new_env = extendVarEnv env old_id (Var new_id)
 355 \end{code}
 356
 357
 358 %************************************************************************
 359 %*                                                                      *
 360                 Types
 361 %*                                                                      *
 362 %************************************************************************
 363
 364 For types we just call the corresponding function in Type, but we have
 365 to repackage the substitution, from a Subst to a TvSubst
 366
 367 \begin{code}
 368 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
 369 substTyVarBndr (Subst in_scope id_env tv_env) tv
 370   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
 371         (TvSubst in_scope' tv_env', tv')
 372            -> (Subst in_scope' id_env tv_env', tv')
 373
 374 substTy :: Subst -> Type -> Type
 375 substTy (Subst in_scope id_env tv_env) ty
 376   = Type.substTy (TvSubst in_scope tv_env) ty
 377 \end{code}
 378
 379
 380 %************************************************************************
 381 %*                                                                      *
 382 \section{IdInfo substitution}
 383 %*                                                                      *
 384 %************************************************************************
 385
 386 \begin{code}
 387 substIdType :: Subst -> Id -> Id
 388 substIdType subst@(Subst in_scope id_env tv_env) id
 389   | isEmptyVarEnv tv_env || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
 390   | otherwise   = setIdType id (substTy subst old_ty)
 391                 -- The tyVarsOfType is cheaper than it looks
 392                 -- because we cache the free tyvars of the type
 393                 -- in a Note in the id's type itself
 394   where
 395     old_ty = idType id
 396
 397 ------------------
 398 substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
 399 -- Always zaps the unfolding, to save substitution work
 400 substIdInfo  subst info
 401   | nothing_to_do = Nothing
 402   | otherwise     = Just (info `setSpecInfo`      substSpec  subst old_rules
 403                                `setWorkerInfo`    substWorker subst old_wrkr
 404                                `setUnfoldingInfo` noUnfolding)
 405   where
 406     old_rules     = specInfo info
 407     old_wrkr      = workerInfo info
 408     nothing_to_do = isEmptySpecInfo old_rules &&
 409                     not (workerExists old_wrkr) &&
 410                     not (hasUnfolding (unfoldingInfo info))
 411
 412
 413 ------------------
 414 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 415         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 416         -- substitutions to happen completely
 417
 418 substWorker subst NoWorker
 419   = NoWorker
 420 substWorker subst (HasWorker w a)
 421   = case lookupIdSubst subst w of
 422         Var w1 -> HasWorker w1 a
 423         other  -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
 424                   NoWorker      -- Worker has got substituted away altogether
 425                                 -- (This can happen if it's trivial,
 426                                 --  via postInlineUnconditionally, hence warning)
 427
 428 ------------------
 429 substSpec :: Subst -> SpecInfo -> SpecInfo
 430
 431 substSpec subst spec@(SpecInfo rules rhs_fvs)
 432   | isEmptySubst subst
 433   = spec
 434   | otherwise
 435   = seqSpecInfo new_rules `seq` new_rules
 436   where
 437     new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
 438
 439     do_subst rule@(BuiltinRule {}) = rule
 440     do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
 441         = rule { ru_bndrs = bndrs',
 442                  ru_args  = map (substExpr subst') args,
 443                  ru_rhs   = substExpr subst' rhs }
 444         where
 445           (subst', bndrs') = substBndrs subst bndrs
 446
 447 ------------------
 448 substVarSet subst fvs
 449   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
 450   where
 451     subst_fv subst fv
 452         | isId fv   = exprFreeVars (lookupIdSubst subst fv)
 453         | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)
 454 \end{code}