ghc/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, substRules, 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                           CoreRules(..), CoreRule(..),
  28                           isEmptyCoreRules, seqRules, 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, setSpecInfo,
  40                           unfoldingInfo, setUnfoldingInfo,
  41                           WorkerInfo(..), workerExists, workerInfo, setWorkerInfo, WorkerInfo
  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 (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
 190                                     where
 191                                       (subst', bndr') = substBndr subst bndr
 192
 193     go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
 194                               where
 195                                 (subst', bndrs') = substRecBndrs subst (map fst pairs)
 196                                 pairs'  = bndrs' `zip` rhss'
 197                                 rhss'   = map (substExpr subst' . snd) pairs
 198
 199     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
 200                                  where
 201                                  (subst', bndr') = substBndr subst bndr
 202
 203     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
 204                                  where
 205                                    (subst', bndrs') = substBndrs subst bndrs
 206
 207     go_note (Coerce ty1 ty2) = Coerce (substTy subst ty1) (substTy subst ty2)
 208     go_note note             = note
 209 \end{code}
 210
 211
 212 %************************************************************************
 213 %*                                                                      *
 214         Substituting binders
 215 %*                                                                      *
 216 %************************************************************************
 217
 218 Remember that substBndr and friends are used when doing expression
 219 substitution only.  Their only business is substitution, so they
 220 preserve all IdInfo (suitably substituted).  For example, we *want* to
 221 preserve occ info in rules.
 222
 223 \begin{code}
 224 substBndr :: Subst -> Var -> (Subst, Var)
 225 substBndr subst bndr
 226   | isTyVar bndr  = substTyVarBndr subst bndr
 227   | otherwise     = substIdBndr subst subst bndr
 228
 229 substBndrs :: Subst -> [Var] -> (Subst, [Var])
 230 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 231
 232 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
 233 -- Substitute a mutually recursive group
 234 substRecBndrs subst bndrs
 235   = (new_subst, new_bndrs)
 236   where         -- Here's the reason we need to pass rec_subst to subst_id
 237     (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
 238 \end{code}
 239
 240
 241 \begin{code}
 242 substIdBndr :: Subst            -- Substitution to use for the IdInfo
 243             -> Subst -> Id      -- Substitition and Id to transform
 244             -> (Subst, Id)      -- Transformed pair
 245
 246 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
 247   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 248   where
 249     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
 250     id2 = substIdType subst id1         -- id2 has its type zapped
 251
 252         -- new_id has the right IdInfo
 253         -- The lazy-set is because we're in a loop here, with
 254         -- rec_subst, when dealing with a mutually-recursive group
 255     new_id = maybeModifyIdInfo (substIdInfo rec_subst) id2
 256
 257         -- Extend the substitution if the unique has changed
 258         -- See the notes with substTyVarBndr for the delSubstEnv
 259     new_env | new_id /= old_id  = extendVarEnv env old_id (Var new_id)
 260             | otherwise         = delVarEnv env old_id
 261 \end{code}
 262
 263 Now a variant that unconditionally allocates a new unique.
 264 It also unconditionally zaps the OccInfo.
 265
 266 \begin{code}
 267 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
 268 cloneIdBndr subst us old_id
 269   = clone_id subst subst (old_id, uniqFromSupply us)
 270
 271 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 272 cloneIdBndrs subst us ids
 273   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
 274
 275 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 276 cloneRecIdBndrs subst us ids
 277   = (subst', ids')
 278   where
 279     (subst', ids') = mapAccumL (clone_id subst') subst
 280                                (ids `zip` uniqsFromSupply us)
 281
 282 -- Just like substIdBndr, except that it always makes a new unique
 283 -- It is given the unique to use
 284 clone_id    :: Subst                    -- Substitution for the IdInfo
 285             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
 286             -> (Subst, Id)              -- Transformed pair
 287
 288 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
 289   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 290   where
 291     id1     = setVarUnique old_id uniq
 292     id2     = substIdType subst id1
 293     new_id  = maybeModifyIdInfo (substIdInfo rec_subst) id2
 294     new_env = extendVarEnv env old_id (Var new_id)
 295 \end{code}
 296
 297
 298 %************************************************************************
 299 %*                                                                      *
 300                 Types
 301 %*                                                                      *
 302 %************************************************************************
 303
 304 For types we just call the corresponding function in Type, but we have
 305 to repackage the substitution, from a Subst to a TvSubst
 306
 307 \begin{code}
 308 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
 309 substTyVarBndr (Subst in_scope id_env tv_env) tv
 310   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
 311         (TvSubst in_scope' tv_env', tv')
 312            -> (Subst in_scope' id_env tv_env', tv')
 313
 314 substTy :: Subst -> Type -> Type
 315 substTy (Subst in_scope id_env tv_env) ty
 316   = Type.substTy (TvSubst in_scope tv_env) ty
 317 \end{code}
 318
 319
 320 %************************************************************************
 321 %*                                                                      *
 322 \section{IdInfo substitution}
 323 %*                                                                      *
 324 %************************************************************************
 325
 326 \begin{code}
 327 substIdType :: Subst -> Id -> Id
 328 substIdType subst@(Subst in_scope id_env tv_env) id
 329   | isEmptyVarEnv tv_env || isEmptyVarSet (tyVarsOfType old_ty) = id
 330   | otherwise   = setIdType id (substTy subst old_ty)
 331                 -- The tyVarsOfType is cheaper than it looks
 332                 -- because we cache the free tyvars of the type
 333                 -- in a Note in the id's type itself
 334   where
 335     old_ty = idType id
 336
 337 ------------------
 338 substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
 339 -- Always zaps the unfolding, to save substitution work
 340 substIdInfo  subst info
 341   | nothing_to_do = Nothing
 342   | otherwise     = Just (info `setSpecInfo`      substRules  subst old_rules
 343                                `setWorkerInfo`    substWorker subst old_wrkr
 344                                `setUnfoldingInfo` noUnfolding)
 345   where
 346     old_rules     = specInfo info
 347     old_wrkr      = workerInfo info
 348     nothing_to_do = isEmptyCoreRules old_rules &&
 349                     not (workerExists old_wrkr) &&
 350                     not (hasUnfolding (unfoldingInfo info))
 351
 352
 353 ------------------
 354 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 355         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 356         -- substitutions to happen completely
 357
 358 substWorker subst NoWorker
 359   = NoWorker
 360 substWorker subst (HasWorker w a)
 361   = case lookupIdSubst subst w of
 362         Var w1 -> HasWorker w1 a
 363         other  -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
 364                   NoWorker      -- Worker has got substituted away altogether
 365                                 -- (This can happen if it's trivial,
 366                                 --  via postInlineUnconditionally, hence warning)
 367
 368 ------------------
 369 substRules :: Subst -> CoreRules -> CoreRules
 370
 371 substRules subst rules
 372  | isEmptySubst subst = rules
 373 substRules subst (Rules rules rhs_fvs)
 374   = seqRules new_rules `seq` new_rules
 375   where
 376     new_rules = Rules (map do_subst rules) (substVarSet subst rhs_fvs)
 377
 378     do_subst rule@(BuiltinRule _ _) = rule
 379     do_subst (Rule name act tpl_vars lhs_args rhs)
 380         = Rule name act tpl_vars'
 381                (map (substExpr subst') lhs_args)
 382                (substExpr subst' rhs)
 383         where
 384           (subst', tpl_vars') = substBndrs subst tpl_vars
 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}