2 % (c) The University of Glasgow 2006
3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
6 Utility functions on @Core@ syntax
11 Subst, TvSubstEnv, IdSubstEnv, InScopeSet,
14 substTy, substExpr, substSpec, substWorker,
15 lookupIdSubst, lookupTvSubst,
17 emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst,
18 extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
19 extendInScope, extendInScopeIds,
23 substBndr, substBndrs, substRecBndrs,
24 cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
27 #include "HsVersions.h"
34 import Type ( Type, TvSubst(..), TvSubstEnv )
38 import Var ( Var, TyVar, setVarUnique )
44 import PprCore () -- Instances
50 %************************************************************************
52 \subsection{Substitutions}
54 %************************************************************************
58 = Subst InScopeSet -- Variables in in scope (both Ids and TyVars)
59 IdSubstEnv -- Substitution for Ids
60 TvSubstEnv -- Substitution for TyVars
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
74 -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
77 -- INVARIANT 3: See Note [Extending the Subst]
80 Note [Extending the Subst]
81 ~~~~~~~~~~~~~~~~~~~~~~~~~~
82 For a core Subst, which binds Ids as well, we make a different choice for Ids
83 than we do for TyVars.
85 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
87 For Ids, we have a different invariant
88 The IdSubstEnv is extended *only* when the Unique on an Id changes
89 Otherwise, we just extend the InScopeSet
93 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
95 * If the TvSubstEnv and IdSubstEnv are both empty, substExpr does nothing
96 (Note that the above rule for substIdBndr maintains this property. If
97 the incoming envts are both empty, then substituting the type and
98 IdInfo can't change anything.)
100 * In lookupIdSubst, we *must* look up the Id in the in-scope set, because
101 it may contain non-trivial changes. Example:
102 (/\a. \x:a. ...x...) Int
103 We extend the TvSubstEnv with [a |-> Int]; but x's unique does not change
104 so we only extend the in-scope set. Then we must look up in the in-scope
105 set when we find the occurrence of x.
107 Why do we make a different choice for the IdSubstEnv than the TvSubstEnv?
109 * For Ids, we change the IdInfo all the time (e.g. deleting the
110 unfolding), and adding it back later, so using the TyVar convention
111 would entail extending the substitution almost all the time
113 * The simplifier wants to look up in the in-scope set anyway, in case it
114 can see a better unfolding from an enclosing case expression
116 * For TyVars, only coercion variables can possibly change, and they are
120 type IdSubstEnv = IdEnv CoreExpr
122 ----------------------------
123 isEmptySubst :: Subst -> Bool
124 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
127 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
129 mkEmptySubst :: InScopeSet -> Subst
130 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
132 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
133 mkSubst in_scope tvs ids = Subst in_scope ids tvs
135 -- getTvSubst :: Subst -> TvSubst
136 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
138 -- getTvSubstEnv :: Subst -> TvSubstEnv
139 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
141 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
142 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
144 substInScope :: Subst -> InScopeSet
145 substInScope (Subst in_scope _ _) = in_scope
147 -- zapSubstEnv :: Subst -> Subst
148 -- zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
150 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
151 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
152 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
154 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
155 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
157 extendTvSubst :: Subst -> TyVar -> Type -> Subst
158 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
160 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
161 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
163 lookupIdSubst :: Subst -> Id -> CoreExpr
164 lookupIdSubst (Subst in_scope ids tvs) v
165 | not (isLocalId v) = Var v
166 | Just e <- lookupVarEnv ids v = e
167 | Just v' <- lookupInScope in_scope v = Var v'
168 -- Vital! See Note [Extending the Subst]
169 | otherwise = WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v )
172 lookupTvSubst :: Subst -> TyVar -> Type
173 lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
175 ------------------------------
176 isInScope :: Var -> Subst -> Bool
177 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
179 extendInScope :: Subst -> Var -> Subst
180 extendInScope (Subst in_scope ids tvs) v
181 = Subst (in_scope `extendInScopeSet` v)
182 (ids `delVarEnv` v) (tvs `delVarEnv` v)
184 extendInScopeIds :: Subst -> [Id] -> Subst
185 extendInScopeIds (Subst in_scope ids tvs) vs
186 = Subst (in_scope `extendInScopeSetList` vs)
187 (ids `delVarEnvList` vs) tvs
190 Pretty printing, for debugging only
193 instance Outputable Subst where
194 ppr (Subst in_scope ids tvs)
195 = ptext SLIT("<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
196 $$ ptext SLIT(" IdSubst =") <+> ppr ids
197 $$ ptext SLIT(" TvSubst =") <+> ppr tvs
202 %************************************************************************
204 Substituting expressions
206 %************************************************************************
209 substExpr :: Subst -> CoreExpr -> CoreExpr
213 go (Var v) = lookupIdSubst subst v
214 go (Type ty) = Type (substTy subst ty)
215 go (Lit lit) = Lit lit
216 go (App fun arg) = App (go fun) (go arg)
217 go (Note note e) = Note (go_note note) (go e)
218 go (Cast e co) = Cast (go e) (substTy subst co)
219 go (Lam bndr body) = Lam bndr' (substExpr subst' body)
221 (subst', bndr') = substBndr subst bndr
223 go (Let bind body) = Let bind' (substExpr subst' body)
225 (subst', bind') = substBind subst bind
227 go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
229 (subst', bndr') = substBndr subst bndr
231 go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
233 (subst', bndrs') = substBndrs subst bndrs
237 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
238 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
240 (subst', bndr') = substBndr subst bndr
242 substBind subst (Rec pairs) = (subst', Rec pairs')
244 (subst', bndrs') = substRecBndrs subst (map fst pairs)
245 pairs' = bndrs' `zip` rhss'
246 rhss' = map (substExpr subst' . snd) pairs
249 De-shadowing the program is sometimes a useful pre-pass. It can be done simply
250 by running over the bindings with an empty substitution, becuase substitution
251 returns a result that has no-shadowing guaranteed.
253 (Actually, within a single *type* there might still be shadowing, because
254 substType is a no-op for the empty substitution, but that's OK.)
257 deShadowBinds :: [CoreBind] -> [CoreBind]
258 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
262 %************************************************************************
266 %************************************************************************
268 Remember that substBndr and friends are used when doing expression
269 substitution only. Their only business is substitution, so they
270 preserve all IdInfo (suitably substituted). For example, we *want* to
271 preserve occ info in rules.
274 substBndr :: Subst -> Var -> (Subst, Var)
276 | isTyVar bndr = substTyVarBndr subst bndr
277 | otherwise = substIdBndr subst subst bndr
279 substBndrs :: Subst -> [Var] -> (Subst, [Var])
280 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
282 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
283 -- Substitute a mutually recursive group
284 substRecBndrs subst bndrs
285 = (new_subst, new_bndrs)
286 where -- Here's the reason we need to pass rec_subst to subst_id
287 (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
292 substIdBndr :: Subst -- Substitution to use for the IdInfo
293 -> Subst -> Id -- Substitition and Id to transform
294 -> (Subst, Id) -- Transformed pair
296 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
297 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
299 id1 = uniqAway in_scope old_id -- id1 is cloned if necessary
300 id2 | no_type_change = id1
301 | otherwise = setIdType id1 (substTy subst old_ty)
303 old_ty = idType old_id
304 no_type_change = isEmptyVarEnv tvs ||
305 isEmptyVarSet (Type.tyVarsOfType old_ty)
307 -- new_id has the right IdInfo
308 -- The lazy-set is because we're in a loop here, with
309 -- rec_subst, when dealing with a mutually-recursive group
310 new_id = maybeModifyIdInfo mb_new_info id2
311 mb_new_info = substIdInfo rec_subst (idInfo id2)
313 -- Extend the substitution if the unique has changed
314 -- See the notes with substTyVarBndr for the delVarEnv
315 new_env | no_change = delVarEnv env old_id
316 | otherwise = extendVarEnv env old_id (Var new_id)
318 no_change = id1 == old_id
319 -- See Note [Extending the Subst]
320 -- *not* necessary to check mb_new_info and no_type_change
323 Now a variant that unconditionally allocates a new unique.
324 It also unconditionally zaps the OccInfo.
327 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
328 cloneIdBndr subst us old_id
329 = clone_id subst subst (old_id, uniqFromSupply us)
331 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
332 cloneIdBndrs subst us ids
333 = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
335 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
336 cloneRecIdBndrs subst us ids
339 (subst', ids') = mapAccumL (clone_id subst') subst
340 (ids `zip` uniqsFromSupply us)
342 -- Just like substIdBndr, except that it always makes a new unique
343 -- It is given the unique to use
344 clone_id :: Subst -- Substitution for the IdInfo
345 -> Subst -> (Id, Unique) -- Substitition and Id to transform
346 -> (Subst, Id) -- Transformed pair
348 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
349 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
351 id1 = setVarUnique old_id uniq
352 id2 = substIdType subst id1
353 new_id = maybeModifyIdInfo (substIdInfo rec_subst (idInfo old_id)) id2
354 new_env = extendVarEnv env old_id (Var new_id)
358 %************************************************************************
362 %************************************************************************
364 For types we just call the corresponding function in Type, but we have
365 to repackage the substitution, from a Subst to a TvSubst
368 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
369 substTyVarBndr (Subst in_scope id_env tv_env) tv
370 = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
371 (TvSubst in_scope' tv_env', tv')
372 -> (Subst in_scope' id_env tv_env', tv')
374 substTy :: Subst -> Type -> Type
375 substTy (Subst in_scope id_env tv_env) ty
376 = Type.substTy (TvSubst in_scope tv_env) ty
380 %************************************************************************
382 \section{IdInfo substitution}
384 %************************************************************************
387 substIdType :: Subst -> Id -> Id
388 substIdType subst@(Subst in_scope id_env tv_env) id
389 | isEmptyVarEnv tv_env || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
390 | otherwise = setIdType id (substTy subst old_ty)
391 -- The tyVarsOfType is cheaper than it looks
392 -- because we cache the free tyvars of the type
393 -- in a Note in the id's type itself
398 substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
399 -- Always zaps the unfolding, to save substitution work
400 substIdInfo subst info
401 | nothing_to_do = Nothing
402 | otherwise = Just (info `setSpecInfo` substSpec subst old_rules
403 `setWorkerInfo` substWorker subst old_wrkr
404 `setUnfoldingInfo` noUnfolding)
406 old_rules = specInfo info
407 old_wrkr = workerInfo info
408 nothing_to_do = isEmptySpecInfo old_rules &&
409 not (workerExists old_wrkr) &&
410 not (hasUnfolding (unfoldingInfo info))
414 substWorker :: Subst -> WorkerInfo -> WorkerInfo
415 -- Seq'ing on the returned WorkerInfo is enough to cause all the
416 -- substitutions to happen completely
418 substWorker subst NoWorker
420 substWorker subst (HasWorker w a)
421 = case lookupIdSubst subst w of
422 Var w1 -> HasWorker w1 a
423 other -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
424 NoWorker -- Worker has got substituted away altogether
425 -- (This can happen if it's trivial,
426 -- via postInlineUnconditionally, hence warning)
429 substSpec :: Subst -> SpecInfo -> SpecInfo
431 substSpec subst spec@(SpecInfo rules rhs_fvs)
435 = seqSpecInfo new_rules `seq` new_rules
437 new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
439 do_subst rule@(BuiltinRule {}) = rule
440 do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
441 = rule { ru_bndrs = bndrs',
442 ru_args = map (substExpr subst') args,
443 ru_rhs = substExpr subst' rhs }
445 (subst', bndrs') = substBndrs subst bndrs
448 substVarSet subst fvs
449 = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
452 | isId fv = exprFreeVars (lookupIdSubst subst fv)
453 | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)