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