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(..), 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, idInfo, 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         = case lookupVarEnv ids v of
 128                           Just e  -> e
 129                           Nothing -> Var v
 130
 131 {-      We used to have to look up in the in-scope set,
 132         because GADTs were implicit in the intermediate language
 133         But with FC, the type of an Id does not change in its scope
 134         The worst that can happen if we don't look up in the in-scope set
 135         is that we don't propagate IdInfo as vigorously as we might.
 136         But that'll happen (when it's useful) in SimplEnv.substId
 137
 138         If you put this back in, you should worry about the
 139                 Just e -> e
 140         case above too!
 141
 142     case lookupInScope in_scope v of {
 143         -- Watch out!  Must get the Id from the in-scope set,
 144         -- because its type there may differ
 145         Just v  -> Var v ;
 146         Nothing -> WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v )
 147                    Var v
 148 -}
 149
 150 lookupTvSubst :: Subst -> TyVar -> Type
 151 lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` mkTyVarTy v
 152
 153 ------------------------------
 154 isInScope :: Var -> Subst -> Bool
 155 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
 156
 157 extendInScope :: Subst -> Var -> Subst
 158 extendInScope (Subst in_scope ids tvs) v
 159   = Subst (in_scope `extendInScopeSet` v)
 160           (ids `delVarEnv` v) (tvs `delVarEnv` v)
 161
 162 extendInScopeIds :: Subst -> [Id] -> Subst
 163 extendInScopeIds (Subst in_scope ids tvs) vs
 164   = Subst (in_scope `extendInScopeSetList` vs)
 165           (ids `delVarEnvList` vs) tvs
 166 \end{code}
 167
 168 Pretty printing, for debugging only
 169
 170 \begin{code}
 171 instance Outputable Subst where
 172   ppr (Subst in_scope ids tvs)
 173         =  ptext SLIT("<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
 174         $$ ptext SLIT(" IdSubst   =") <+> ppr ids
 175         $$ ptext SLIT(" TvSubst   =") <+> ppr tvs
 176          <> char '>'
 177 \end{code}
 178
 179
 180 %************************************************************************
 181 %*                                                                      *
 182         Substituting expressions
 183 %*                                                                      *
 184 %************************************************************************
 185
 186 \begin{code}
 187 substExpr :: Subst -> CoreExpr -> CoreExpr
 188 substExpr subst expr
 189   = go expr
 190   where
 191     go (Var v)         = lookupIdSubst subst v
 192     go (Type ty)       = Type (substTy subst ty)
 193     go (Lit lit)       = Lit lit
 194     go (App fun arg)   = App (go fun) (go arg)
 195     go (Note note e)   = Note (go_note note) (go e)
 196     go (Cast e co)     = Cast (go e) (substTy subst co)
 197     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
 198                        where
 199                          (subst', bndr') = substBndr subst bndr
 200
 201     go (Let bind body) = Let bind' (substExpr subst' body)
 202                        where
 203                          (subst', bind') = substBind subst bind
 204
 205     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
 206                                  where
 207                                  (subst', bndr') = substBndr subst bndr
 208
 209     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
 210                                  where
 211                                    (subst', bndrs') = substBndrs subst bndrs
 212
 213     go_note note             = note
 214
 215 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
 216 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
 217                                   where
 218                                     (subst', bndr') = substBndr subst bndr
 219
 220 substBind subst (Rec pairs) = (subst', Rec pairs')
 221                             where
 222                                 (subst', bndrs') = substRecBndrs subst (map fst pairs)
 223                                 pairs'  = bndrs' `zip` rhss'
 224                                 rhss'   = map (substExpr subst' . snd) pairs
 225 \end{code}
 226
 227 De-shadowing the program is sometimes a useful pre-pass.  It can be done simply
 228 by running over the bindings with an empty substitution, becuase substitution
 229 returns a result that has no-shadowing guaranteed.
 230
 231 (Actually, within a single *type* there might still be shadowing, because
 232 substType is a no-op for the empty substitution, but that's OK.)
 233
 234 \begin{code}
 235 deShadowBinds :: [CoreBind] -> [CoreBind]
 236 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
 237 \end{code}
 238
 239
 240 %************************************************************************
 241 %*                                                                      *
 242         Substituting binders
 243 %*                                                                      *
 244 %************************************************************************
 245
 246 Remember that substBndr and friends are used when doing expression
 247 substitution only.  Their only business is substitution, so they
 248 preserve all IdInfo (suitably substituted).  For example, we *want* to
 249 preserve occ info in rules.
 250
 251 \begin{code}
 252 substBndr :: Subst -> Var -> (Subst, Var)
 253 substBndr subst bndr
 254   | isTyVar bndr  = substTyVarBndr subst bndr
 255   | otherwise     = substIdBndr subst subst bndr
 256
 257 substBndrs :: Subst -> [Var] -> (Subst, [Var])
 258 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 259
 260 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
 261 -- Substitute a mutually recursive group
 262 substRecBndrs subst bndrs
 263   = (new_subst, new_bndrs)
 264   where         -- Here's the reason we need to pass rec_subst to subst_id
 265     (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
 266 \end{code}
 267
 268
 269 \begin{code}
 270 substIdBndr :: Subst            -- Substitution to use for the IdInfo
 271             -> Subst -> Id      -- Substitition and Id to transform
 272             -> (Subst, Id)      -- Transformed pair
 273
 274 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
 275   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 276   where
 277     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
 278     id2 | no_type_change = id1
 279         | otherwise      = setIdType id1 (substTy subst old_ty)
 280
 281     old_ty = idType old_id
 282     no_type_change = isEmptyVarEnv tvs || isEmptyVarSet (tyVarsOfType old_ty)
 283
 284         -- new_id has the right IdInfo
 285         -- The lazy-set is because we're in a loop here, with
 286         -- rec_subst, when dealing with a mutually-recursive group
 287     new_id = maybeModifyIdInfo mb_new_info id2
 288     mb_new_info = substIdInfo rec_subst (idInfo id2)
 289
 290         -- Extend the substitution if the unique has changed
 291         -- See the notes with substTyVarBndr for the delVarEnv
 292     new_env | no_change = delVarEnv env old_id
 293             | otherwise = extendVarEnv env old_id (Var new_id)
 294
 295     no_change = False -- id1 == old_id && isNothing mb_new_info && no_type_change
 296 \end{code}
 297
 298 Now a variant that unconditionally allocates a new unique.
 299 It also unconditionally zaps the OccInfo.
 300
 301 \begin{code}
 302 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
 303 cloneIdBndr subst us old_id
 304   = clone_id subst subst (old_id, uniqFromSupply us)
 305
 306 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 307 cloneIdBndrs subst us ids
 308   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
 309
 310 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
 311 cloneRecIdBndrs subst us ids
 312   = (subst', ids')
 313   where
 314     (subst', ids') = mapAccumL (clone_id subst') subst
 315                                (ids `zip` uniqsFromSupply us)
 316
 317 -- Just like substIdBndr, except that it always makes a new unique
 318 -- It is given the unique to use
 319 clone_id    :: Subst                    -- Substitution for the IdInfo
 320             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
 321             -> (Subst, Id)              -- Transformed pair
 322
 323 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
 324   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
 325   where
 326     id1     = setVarUnique old_id uniq
 327     id2     = substIdType subst id1
 328     new_id  = maybeModifyIdInfo (substIdInfo rec_subst (idInfo old_id)) id2
 329     new_env = extendVarEnv env old_id (Var new_id)
 330 \end{code}
 331
 332
 333 %************************************************************************
 334 %*                                                                      *
 335                 Types
 336 %*                                                                      *
 337 %************************************************************************
 338
 339 For types we just call the corresponding function in Type, but we have
 340 to repackage the substitution, from a Subst to a TvSubst
 341
 342 \begin{code}
 343 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
 344 substTyVarBndr (Subst in_scope id_env tv_env) tv
 345   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
 346         (TvSubst in_scope' tv_env', tv')
 347            -> (Subst in_scope' id_env tv_env', tv')
 348
 349 substTy :: Subst -> Type -> Type
 350 substTy (Subst in_scope id_env tv_env) ty
 351   = Type.substTy (TvSubst in_scope tv_env) ty
 352 \end{code}
 353
 354
 355 %************************************************************************
 356 %*                                                                      *
 357 \section{IdInfo substitution}
 358 %*                                                                      *
 359 %************************************************************************
 360
 361 \begin{code}
 362 substIdType :: Subst -> Id -> Id
 363 substIdType subst@(Subst in_scope id_env tv_env) id
 364   | isEmptyVarEnv tv_env || isEmptyVarSet (tyVarsOfType old_ty) = id
 365   | otherwise   = setIdType id (substTy subst old_ty)
 366                 -- The tyVarsOfType is cheaper than it looks
 367                 -- because we cache the free tyvars of the type
 368                 -- in a Note in the id's type itself
 369   where
 370     old_ty = idType id
 371
 372 ------------------
 373 substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
 374 -- Always zaps the unfolding, to save substitution work
 375 substIdInfo  subst info
 376   | nothing_to_do = Nothing
 377   | otherwise     = Just (info `setSpecInfo`      substSpec  subst old_rules
 378                                `setWorkerInfo`    substWorker subst old_wrkr
 379                                `setUnfoldingInfo` noUnfolding)
 380   where
 381     old_rules     = specInfo info
 382     old_wrkr      = workerInfo info
 383     nothing_to_do = isEmptySpecInfo old_rules &&
 384                     not (workerExists old_wrkr) &&
 385                     not (hasUnfolding (unfoldingInfo info))
 386
 387
 388 ------------------
 389 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 390         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 391         -- substitutions to happen completely
 392
 393 substWorker subst NoWorker
 394   = NoWorker
 395 substWorker subst (HasWorker w a)
 396   = case lookupIdSubst subst w of
 397         Var w1 -> HasWorker w1 a
 398         other  -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
 399                   NoWorker      -- Worker has got substituted away altogether
 400                                 -- (This can happen if it's trivial,
 401                                 --  via postInlineUnconditionally, hence warning)
 402
 403 ------------------
 404 substSpec :: Subst -> SpecInfo -> SpecInfo
 405
 406 substSpec subst spec@(SpecInfo rules rhs_fvs)
 407   | isEmptySubst subst
 408   = spec
 409   | otherwise
 410   = seqSpecInfo new_rules `seq` new_rules
 411   where
 412     new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
 413
 414     do_subst rule@(BuiltinRule {}) = rule
 415     do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
 416         = rule { ru_bndrs = bndrs',
 417                  ru_args  = map (substExpr subst') args,
 418                  ru_rhs   = substExpr subst' rhs }
 419         where
 420           (subst', bndrs') = substBndrs subst bndrs
 421
 422 ------------------
 423 substVarSet subst fvs
 424   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
 425   where
 426     subst_fv subst fv
 427         | isId fv   = exprFreeVars (lookupIdSubst subst fv)
 428         | otherwise = tyVarsOfType (lookupTvSubst subst fv)
 429 \end{code}