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,
13 -- ** Substituting into expressions and related types
15 substTy, substExpr, substSpec, substWorker,
16 lookupIdSubst, lookupTvSubst,
18 -- ** Operations on substitutions
19 emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst,
20 extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
21 extendSubst, extendSubstList, zapSubstEnv,
22 extendInScope, extendInScopeList, extendInScopeIds,
25 -- ** Substituting and cloning binders
26 substBndr, substBndrs, substRecBndrs,
27 cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
30 #include "HsVersions.h"
37 import Type ( Type, TvSubst(..), TvSubstEnv )
41 import Var ( Var, TyVar, setVarUnique )
47 import PprCore () -- Instances
54 %************************************************************************
56 \subsection{Substitutions}
58 %************************************************************************
61 -- | A substitution environment, containing both 'Id' and 'TyVar' substitutions.
63 -- Some invariants apply to how you use the substitution:
65 -- 1. #in_scope_invariant# The in-scope set contains at least those 'Id's and 'TyVar's that will be in scope /after/
66 -- applying the substitution to a term. Precisely, the in-scope set must be a superset of the free vars of the
67 -- substitution range that might possibly clash with locally-bound variables in the thing being substituted in.
69 -- 2. #apply_once# You may apply the substitution only /once/
71 -- There are various ways of setting up the in-scope set such that the first of these invariants hold:
73 -- * Arrange that the in-scope set really is all the things in scope
75 -- * Arrange that it's the free vars of the range of the substitution
77 -- * Make it empty, if you know that all the free vars of the substitution are fresh, and hence can't possibly clash
79 = Subst InScopeSet -- Variables in in scope (both Ids and TyVars) /after/
80 -- applying the substitution
81 IdSubstEnv -- Substitution for Ids
82 TvSubstEnv -- Substitution for TyVars
84 -- INVARIANT 1: See #in_scope_invariant#
85 -- This is what lets us deal with name capture properly
86 -- It's a hard invariant to check...
88 -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
91 -- INVARIANT 3: See Note [Extending the Subst]
94 Note [Extending the Subst]
95 ~~~~~~~~~~~~~~~~~~~~~~~~~~
96 For a core Subst, which binds Ids as well, we make a different choice for Ids
97 than we do for TyVars.
99 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
101 For Ids, we have a different invariant
102 The IdSubstEnv is extended *only* when the Unique on an Id changes
103 Otherwise, we just extend the InScopeSet
107 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
109 * If the TvSubstEnv and IdSubstEnv are both empty, substExpr does nothing
110 (Note that the above rule for substIdBndr maintains this property. If
111 the incoming envts are both empty, then substituting the type and
112 IdInfo can't change anything.)
114 * In lookupIdSubst, we *must* look up the Id in the in-scope set, because
115 it may contain non-trivial changes. Example:
116 (/\a. \x:a. ...x...) Int
117 We extend the TvSubstEnv with [a |-> Int]; but x's unique does not change
118 so we only extend the in-scope set. Then we must look up in the in-scope
119 set when we find the occurrence of x.
121 * The requirement to look up the Id in the in-scope set means that we
122 must NOT take no-op short cut in the case the substitution is empty.
123 We must still look up every Id in the in-scope set.
125 * (However, we don't need to do so for expressions found in the IdSubst
126 itself, whose range is assumed to be correct wrt the in-scope set.)
128 Why do we make a different choice for the IdSubstEnv than the TvSubstEnv?
130 * For Ids, we change the IdInfo all the time (e.g. deleting the
131 unfolding), and adding it back later, so using the TyVar convention
132 would entail extending the substitution almost all the time
134 * The simplifier wants to look up in the in-scope set anyway, in case it
135 can see a better unfolding from an enclosing case expression
137 * For TyVars, only coercion variables can possibly change, and they are
141 -- | An environment for substituting for 'Id's
142 type IdSubstEnv = IdEnv CoreExpr
144 ----------------------------
145 isEmptySubst :: Subst -> Bool
146 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
149 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
151 mkEmptySubst :: InScopeSet -> Subst
152 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
154 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
155 mkSubst in_scope tvs ids = Subst in_scope ids tvs
157 -- getTvSubst :: Subst -> TvSubst
158 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
160 -- getTvSubstEnv :: Subst -> TvSubstEnv
161 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
163 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
164 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
166 -- | Find the in-scope set: see "CoreSubst#in_scope_invariant"
167 substInScope :: Subst -> InScopeSet
168 substInScope (Subst in_scope _ _) = in_scope
170 -- | Remove all substitutions for 'Id's and 'Var's that might have been built up
171 -- while preserving the in-scope set
172 zapSubstEnv :: Subst -> Subst
173 zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
175 -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the in-scope set is
176 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
177 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
178 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
179 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
181 -- | Adds multiple 'Id' substitutions to the 'Subst': see also 'extendIdSubst'
182 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
183 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
185 -- | Add a substitution for a 'TyVar' to the 'Subst': you must ensure that the in-scope set is
186 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
187 extendTvSubst :: Subst -> TyVar -> Type -> Subst
188 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
190 -- | Adds multiple 'TyVar' substitutions to the 'Subst': see also 'extendTvSubst'
191 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
192 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
194 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to the 'Var' being added. See also
195 -- 'extendIdSubst' and 'extendTvSubst'
196 extendSubst :: Subst -> Var -> CoreArg -> Subst
197 extendSubst (Subst in_scope ids tvs) tv (Type ty)
198 = ASSERT( isTyVar tv ) Subst in_scope ids (extendVarEnv tvs tv ty)
199 extendSubst (Subst in_scope ids tvs) id expr
200 = ASSERT( isId id ) Subst in_scope (extendVarEnv ids id expr) tvs
202 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to all the 'Var's being added. See also 'extendSubst'
203 extendSubstList :: Subst -> [(Var,CoreArg)] -> Subst
204 extendSubstList subst [] = subst
205 extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var rhs) prs
207 -- | Find the substitution for an 'Id' in the 'Subst'
208 lookupIdSubst :: Subst -> Id -> CoreExpr
209 lookupIdSubst (Subst in_scope ids _) v
210 | not (isLocalId v) = Var v
211 | Just e <- lookupVarEnv ids v = e
212 | Just v' <- lookupInScope in_scope v = Var v'
213 -- Vital! See Note [Extending the Subst]
214 | otherwise = WARN( True, ptext (sLit "CoreSubst.lookupIdSubst") <+> ppr v )
217 -- | Find the substitution for a 'TyVar' in the 'Subst'
218 lookupTvSubst :: Subst -> TyVar -> Type
219 lookupTvSubst (Subst _ _ tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
221 ------------------------------
222 isInScope :: Var -> Subst -> Bool
223 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
225 -- | Add the 'Var' to the in-scope set: as a side effect, removes any existing substitutions for it
226 extendInScope :: Subst -> Var -> Subst
227 extendInScope (Subst in_scope ids tvs) v
228 = Subst (in_scope `extendInScopeSet` v)
229 (ids `delVarEnv` v) (tvs `delVarEnv` v)
231 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
232 extendInScopeList :: Subst -> [Var] -> Subst
233 extendInScopeList (Subst in_scope ids tvs) vs
234 = Subst (in_scope `extendInScopeSetList` vs)
235 (ids `delVarEnvList` vs) (tvs `delVarEnvList` vs)
237 -- | Optimized version of 'extendInScopeList' that can be used if you are certain
238 -- all the things being added are 'Id's and hence none are 'TyVar's
239 extendInScopeIds :: Subst -> [Id] -> Subst
240 extendInScopeIds (Subst in_scope ids tvs) vs
241 = Subst (in_scope `extendInScopeSetList` vs)
242 (ids `delVarEnvList` vs) tvs
245 Pretty printing, for debugging only
248 instance Outputable Subst where
249 ppr (Subst in_scope ids tvs)
250 = ptext (sLit "<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
251 $$ ptext (sLit " IdSubst =") <+> ppr ids
252 $$ ptext (sLit " TvSubst =") <+> ppr tvs
257 %************************************************************************
259 Substituting expressions
261 %************************************************************************
264 -- | Apply a substititon to an entire 'CoreExpr'. Rememeber, you may only
265 -- apply the substitution /once/: see "CoreSubst#apply_once"
267 -- Do *not* attempt to short-cut in the case of an empty substitution!
268 -- See Note [Extending the Subst]
269 substExpr :: Subst -> CoreExpr -> CoreExpr
273 go (Var v) = lookupIdSubst subst v
274 go (Type ty) = Type (substTy subst ty)
275 go (Lit lit) = Lit lit
276 go (App fun arg) = App (go fun) (go arg)
277 go (Note note e) = Note (go_note note) (go e)
278 go (Cast e co) = Cast (go e) (substTy subst co)
279 go (Lam bndr body) = Lam bndr' (substExpr subst' body)
281 (subst', bndr') = substBndr subst bndr
283 go (Let bind body) = Let bind' (substExpr subst' body)
285 (subst', bind') = substBind subst bind
287 go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
289 (subst', bndr') = substBndr subst bndr
291 go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
293 (subst', bndrs') = substBndrs subst bndrs
297 -- | Apply a substititon to an entire 'CoreBind', additionally returning an updated 'Subst'
298 -- that should be used by subsequent substitutons.
299 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
300 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
302 (subst', bndr') = substBndr subst bndr
304 substBind subst (Rec pairs) = (subst', Rec pairs')
306 (subst', bndrs') = substRecBndrs subst (map fst pairs)
307 pairs' = bndrs' `zip` rhss'
308 rhss' = map (substExpr subst' . snd) pairs
312 -- | De-shadowing the program is sometimes a useful pre-pass. It can be done simply
313 -- by running over the bindings with an empty substitution, becuase substitution
314 -- returns a result that has no-shadowing guaranteed.
316 -- (Actually, within a single /type/ there might still be shadowing, because
317 -- 'substTy' is a no-op for the empty substitution, but that's probably OK.)
318 deShadowBinds :: [CoreBind] -> [CoreBind]
319 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
323 %************************************************************************
327 %************************************************************************
329 Remember that substBndr and friends are used when doing expression
330 substitution only. Their only business is substitution, so they
331 preserve all IdInfo (suitably substituted). For example, we *want* to
332 preserve occ info in rules.
335 -- | Substitutes a 'Var' for another one according to the 'Subst' given, returning
336 -- the result and an updated 'Subst' that should be used by subsequent substitutons.
337 -- 'IdInfo' is preserved by this process, although it is substituted into appropriately.
338 substBndr :: Subst -> Var -> (Subst, Var)
340 | isTyVar bndr = substTyVarBndr subst bndr
341 | otherwise = substIdBndr subst subst bndr
343 -- | Applies 'substBndr' to a number of 'Var's, accumulating a new 'Subst' left-to-right
344 substBndrs :: Subst -> [Var] -> (Subst, [Var])
345 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
347 -- | Substitute in a mutually recursive group of 'Id's
348 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
349 substRecBndrs subst bndrs
350 = (new_subst, new_bndrs)
351 where -- Here's the reason we need to pass rec_subst to subst_id
352 (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
357 substIdBndr :: Subst -- ^ Substitution to use for the IdInfo
358 -> Subst -> Id -- ^ Substitition and Id to transform
359 -> (Subst, Id) -- ^ Transformed pair
360 -- NB: unfolding may be zapped
362 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
363 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
365 id1 = uniqAway in_scope old_id -- id1 is cloned if necessary
366 id2 | no_type_change = id1
367 | otherwise = setIdType id1 (substTy subst old_ty)
369 old_ty = idType old_id
370 no_type_change = isEmptyVarEnv tvs ||
371 isEmptyVarSet (Type.tyVarsOfType old_ty)
373 -- new_id has the right IdInfo
374 -- The lazy-set is because we're in a loop here, with
375 -- rec_subst, when dealing with a mutually-recursive group
376 new_id = maybeModifyIdInfo mb_new_info id2
377 mb_new_info = substIdInfo rec_subst id2 (idInfo id2)
378 -- NB: unfolding info may be zapped
380 -- Extend the substitution if the unique has changed
381 -- See the notes with substTyVarBndr for the delVarEnv
382 new_env | no_change = delVarEnv env old_id
383 | otherwise = extendVarEnv env old_id (Var new_id)
385 no_change = id1 == old_id
386 -- See Note [Extending the Subst]
387 -- it's /not/ necessary to check mb_new_info and no_type_change
390 Now a variant that unconditionally allocates a new unique.
391 It also unconditionally zaps the OccInfo.
394 -- | Very similar to 'substBndr', but it always allocates a new 'Unique' for
395 -- each variable in its output and removes all 'IdInfo'
396 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
397 cloneIdBndr subst us old_id
398 = clone_id subst subst (old_id, uniqFromSupply us)
400 -- | Applies 'cloneIdBndr' to a number of 'Id's, accumulating a final
401 -- substitution from left to right
402 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
403 cloneIdBndrs subst us ids
404 = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
406 -- | Clone a mutually recursive group of 'Id's
407 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
408 cloneRecIdBndrs subst us ids
411 (subst', ids') = mapAccumL (clone_id subst') subst
412 (ids `zip` uniqsFromSupply us)
414 -- Just like substIdBndr, except that it always makes a new unique
415 -- It is given the unique to use
416 clone_id :: Subst -- Substitution for the IdInfo
417 -> Subst -> (Id, Unique) -- Substitition and Id to transform
418 -> (Subst, Id) -- Transformed pair
420 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
421 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
423 id1 = setVarUnique old_id uniq
424 id2 = substIdType subst id1
425 new_id = maybeModifyIdInfo (substIdInfo rec_subst id2 (idInfo old_id)) id2
426 new_env = extendVarEnv env old_id (Var new_id)
430 %************************************************************************
434 %************************************************************************
436 For types we just call the corresponding function in Type, but we have
437 to repackage the substitution, from a Subst to a TvSubst
440 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
441 substTyVarBndr (Subst in_scope id_env tv_env) tv
442 = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
443 (TvSubst in_scope' tv_env', tv')
444 -> (Subst in_scope' id_env tv_env', tv')
446 -- | See 'Type.substTy'
447 substTy :: Subst -> Type -> Type
448 substTy (Subst in_scope _id_env tv_env) ty
449 = Type.substTy (TvSubst in_scope tv_env) ty
453 %************************************************************************
455 \section{IdInfo substitution}
457 %************************************************************************
460 substIdType :: Subst -> Id -> Id
461 substIdType subst@(Subst _ _ tv_env) id
462 | isEmptyVarEnv tv_env || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
463 | otherwise = setIdType id (substTy subst old_ty)
464 -- The tyVarsOfType is cheaper than it looks
465 -- because we cache the free tyvars of the type
466 -- in a Note in the id's type itself
471 -- | Substitute into some 'IdInfo' with regard to the supplied new 'Id'.
472 -- Always zaps the unfolding, to save substitution work
473 substIdInfo :: Subst -> Id -> IdInfo -> Maybe IdInfo
474 substIdInfo subst new_id info
475 | nothing_to_do = Nothing
476 | otherwise = Just (info `setSpecInfo` substSpec subst new_id old_rules
477 `setWorkerInfo` substWorker subst old_wrkr
478 `setUnfoldingInfo` noUnfolding)
480 old_rules = specInfo info
481 old_wrkr = workerInfo info
482 nothing_to_do = isEmptySpecInfo old_rules &&
483 not (workerExists old_wrkr) &&
484 not (hasUnfolding (unfoldingInfo info))
488 -- | Substitutes for the 'Id's within the 'WorkerInfo'
489 substWorker :: Subst -> WorkerInfo -> WorkerInfo
490 -- Seq'ing on the returned WorkerInfo is enough to cause all the
491 -- substitutions to happen completely
493 substWorker _ NoWorker
495 substWorker subst (HasWorker w a)
496 = case lookupIdSubst subst w of
497 Var w1 -> HasWorker w1 a
498 other -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
499 NoWorker -- Worker has got substituted away altogether
500 -- (This can happen if it's trivial,
501 -- via postInlineUnconditionally, hence warning)
504 -- | Substitutes for the 'Id's within the 'WorkerInfo' given the new function 'Id'
505 substSpec :: Subst -> Id -> SpecInfo -> SpecInfo
506 substSpec subst new_fn (SpecInfo rules rhs_fvs)
507 = seqSpecInfo new_rules `seq` new_rules
509 new_name = idName new_fn
510 new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
512 do_subst rule@(BuiltinRule {}) = rule
513 do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
514 = rule { ru_bndrs = bndrs',
515 ru_fn = new_name, -- Important: the function may have changed its name!
516 ru_args = map (substExpr subst') args,
517 ru_rhs = substExpr subst' rhs }
519 (subst', bndrs') = substBndrs subst bndrs
522 substVarSet :: Subst -> VarSet -> VarSet
523 substVarSet subst fvs
524 = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
527 | isId fv = exprFreeVars (lookupIdSubst subst fv)
528 | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)