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 notes with
77 type IdSubstEnv = IdEnv CoreExpr
79 ----------------------------
80 isEmptySubst :: Subst -> Bool
81 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
84 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
86 mkEmptySubst :: InScopeSet -> Subst
87 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
89 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
90 mkSubst in_scope tvs ids = Subst in_scope ids tvs
92 -- getTvSubst :: Subst -> TvSubst
93 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
95 -- getTvSubstEnv :: Subst -> TvSubstEnv
96 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
98 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
99 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
101 substInScope :: Subst -> InScopeSet
102 substInScope (Subst in_scope _ _) = in_scope
104 -- zapSubstEnv :: Subst -> Subst
105 -- zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
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
111 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
112 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
114 extendTvSubst :: Subst -> TyVar -> Type -> Subst
115 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
117 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
118 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
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
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
134 If you put this back in, you should worry about the
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
142 Nothing -> WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v )
146 lookupTvSubst :: Subst -> TyVar -> Type
147 lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
149 ------------------------------
150 isInScope :: Var -> Subst -> Bool
151 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
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)
158 extendInScopeIds :: Subst -> [Id] -> Subst
159 extendInScopeIds (Subst in_scope ids tvs) vs
160 = Subst (in_scope `extendInScopeSetList` vs)
161 (ids `delVarEnvList` vs) tvs
164 Pretty printing, for debugging only
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
176 %************************************************************************
178 Substituting expressions
180 %************************************************************************
183 substExpr :: Subst -> CoreExpr -> CoreExpr
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)
195 (subst', bndr') = substBndr subst bndr
197 go (Let bind body) = Let bind' (substExpr subst' body)
199 (subst', bind') = substBind subst bind
201 go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
203 (subst', bndr') = substBndr subst bndr
205 go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
207 (subst', bndrs') = substBndrs subst bndrs
211 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
212 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
214 (subst', bndr') = substBndr subst bndr
216 substBind subst (Rec pairs) = (subst', Rec pairs')
218 (subst', bndrs') = substRecBndrs subst (map fst pairs)
219 pairs' = bndrs' `zip` rhss'
220 rhss' = map (substExpr subst' . snd) pairs
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.
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.)
231 deShadowBinds :: [CoreBind] -> [CoreBind]
232 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
236 %************************************************************************
240 %************************************************************************
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.
248 substBndr :: Subst -> Var -> (Subst, Var)
250 | isTyVar bndr = substTyVarBndr subst bndr
251 | otherwise = substIdBndr subst subst bndr
253 substBndrs :: Subst -> [Var] -> (Subst, [Var])
254 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
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
266 substIdBndr :: Subst -- Substitution to use for the IdInfo
267 -> Subst -> Id -- Substitition and Id to transform
268 -> (Subst, Id) -- Transformed pair
270 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
271 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
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)
277 old_ty = idType old_id
278 no_type_change = isEmptyVarEnv tvs ||
279 isEmptyVarSet (Type.tyVarsOfType old_ty)
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)
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)
292 no_change = False -- id1 == old_id && isNothing mb_new_info && no_type_change
295 Now a variant that unconditionally allocates a new unique.
296 It also unconditionally zaps the OccInfo.
299 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
300 cloneIdBndr subst us old_id
301 = clone_id subst subst (old_id, uniqFromSupply us)
303 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
304 cloneIdBndrs subst us ids
305 = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
307 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
308 cloneRecIdBndrs subst us ids
311 (subst', ids') = mapAccumL (clone_id subst') subst
312 (ids `zip` uniqsFromSupply us)
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
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)
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)
330 %************************************************************************
334 %************************************************************************
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
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')
346 substTy :: Subst -> Type -> Type
347 substTy (Subst in_scope id_env tv_env) ty
348 = Type.substTy (TvSubst in_scope tv_env) ty
352 %************************************************************************
354 \section{IdInfo substitution}
356 %************************************************************************
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
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)
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))
386 substWorker :: Subst -> WorkerInfo -> WorkerInfo
387 -- Seq'ing on the returned WorkerInfo is enough to cause all the
388 -- substitutions to happen completely
390 substWorker subst 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)
401 substSpec :: Subst -> SpecInfo -> SpecInfo
403 substSpec subst spec@(SpecInfo rules rhs_fvs)
407 = seqSpecInfo new_rules `seq` new_rules
409 new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
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 }
417 (subst', bndrs') = substBndrs subst bndrs
420 substVarSet subst fvs
421 = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
424 | isId fv = exprFreeVars (lookupIdSubst subst fv)
425 | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)