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