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