2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[CoreUtils]{Utility functions on @Core@ syntax}
9 Subst, TvSubstEnv, IdSubstEnv, InScopeSet,
11 substTy, substExpr, substSpec, substWorker,
12 lookupIdSubst, lookupTvSubst,
14 emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst,
15 extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
16 extendInScope, extendInScopeIds,
20 substBndr, substBndrs, substRecBndrs,
21 cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
24 #include "HsVersions.h"
26 import CoreSyn ( Expr(..), Bind(..), Note(..), CoreExpr,
27 CoreRule(..), hasUnfolding, noUnfolding
29 import CoreFVs ( exprFreeVars )
30 import CoreUtils ( exprIsTrivial )
32 import qualified Type ( substTy, substTyVarBndr )
33 import Type ( Type, tyVarsOfType, TvSubstEnv, TvSubst(..), mkTyVarTy )
36 import Var ( setVarUnique, isId )
37 import Id ( idType, setIdType, maybeModifyIdInfo, isLocalId )
38 import IdInfo ( IdInfo, SpecInfo(..), specInfo, setSpecInfo, isEmptySpecInfo,
39 unfoldingInfo, setUnfoldingInfo, seqSpecInfo,
40 WorkerInfo(..), workerExists, workerInfo, setWorkerInfo, WorkerInfo
42 import Unique ( Unique )
43 import UniqSupply ( UniqSupply, uniqFromSupply, uniqsFromSupply )
44 import Var ( Var, Id, TyVar, isTyVar )
45 import Maybes ( orElse )
47 import PprCore () -- Instances
48 import Util ( mapAccumL )
53 %************************************************************************
55 \subsection{Substitutions}
57 %************************************************************************
61 = Subst InScopeSet -- Variables in in scope (both Ids and TyVars)
62 IdSubstEnv -- Substitution for Ids
63 TvSubstEnv -- Substitution for TyVars
65 -- INVARIANT 1: The (domain of the) in-scope set is a superset
66 -- of the free vars of the range of the substitution
67 -- that might possibly clash with locally-bound variables
68 -- in the thing being substituted in.
69 -- This is what lets us deal with name capture properly
70 -- It's a hard invariant to check...
71 -- There are various ways of causing it to happen:
72 -- - arrange that the in-scope set really is all the things in scope
73 -- - arrange that it's the free vars of the range of the substitution
74 -- - make it empty because all the free vars of the subst are fresh,
75 -- and hence can't possibly clash.a
77 -- INVARIANT 2: The substitution is apply-once; see notes with
80 type IdSubstEnv = IdEnv CoreExpr
82 ----------------------------
83 isEmptySubst :: Subst -> Bool
84 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
87 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
89 mkEmptySubst :: InScopeSet -> Subst
90 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
92 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
93 mkSubst in_scope tvs ids = Subst in_scope ids tvs
95 -- getTvSubst :: Subst -> TvSubst
96 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
98 -- getTvSubstEnv :: Subst -> TvSubstEnv
99 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
101 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
102 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
104 substInScope :: Subst -> InScopeSet
105 substInScope (Subst in_scope _ _) = in_scope
107 -- zapSubstEnv :: Subst -> Subst
108 -- zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
110 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
111 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
112 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
114 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
115 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
117 extendTvSubst :: Subst -> TyVar -> Type -> Subst
118 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r)
120 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
121 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
123 lookupIdSubst :: Subst -> Id -> CoreExpr
124 lookupIdSubst (Subst in_scope ids tvs) v
125 | not (isLocalId v) = Var v
127 = case lookupVarEnv ids v of {
130 case lookupInScope in_scope v of {
131 -- Watch out! Must get the Id from the in-scope set,
132 -- because its type there may differ
134 Nothing -> WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v )
138 lookupTvSubst :: Subst -> TyVar -> Type
139 lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` mkTyVarTy v
141 ------------------------------
142 isInScope :: Var -> Subst -> Bool
143 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
145 extendInScope :: Subst -> Var -> Subst
146 extendInScope (Subst in_scope ids tvs) v
147 = Subst (in_scope `extendInScopeSet` v)
148 (ids `delVarEnv` v) (tvs `delVarEnv` v)
150 extendInScopeIds :: Subst -> [Id] -> Subst
151 extendInScopeIds (Subst in_scope ids tvs) vs
152 = Subst (in_scope `extendInScopeSetList` vs)
153 (ids `delVarEnvList` vs) tvs
156 Pretty printing, for debugging only
159 instance Outputable Subst where
160 ppr (Subst in_scope ids tvs)
161 = ptext SLIT("<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
162 $$ ptext SLIT(" IdSubst =") <+> ppr ids
163 $$ ptext SLIT(" TvSubst =") <+> ppr tvs
168 %************************************************************************
170 Substituting expressions
172 %************************************************************************
175 substExpr :: Subst -> CoreExpr -> CoreExpr
179 go (Var v) = lookupIdSubst subst v
180 go (Type ty) = Type (substTy subst ty)
181 go (Lit lit) = Lit lit
182 go (App fun arg) = App (go fun) (go arg)
183 go (Note note e) = Note (go_note note) (go e)
184 go (Lam bndr body) = Lam bndr' (substExpr subst' body)
186 (subst', bndr') = substBndr subst bndr
188 go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
190 (subst', bndr') = substBndr subst bndr
192 go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
194 (subst', bndrs') = substRecBndrs subst (map fst pairs)
195 pairs' = bndrs' `zip` rhss'
196 rhss' = map (substExpr subst' . snd) pairs
198 go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
200 (subst', bndr') = substBndr subst bndr
202 go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
204 (subst', bndrs') = substBndrs subst bndrs
206 go_note (Coerce ty1 ty2) = Coerce (substTy subst ty1) (substTy subst ty2)
211 %************************************************************************
215 %************************************************************************
217 Remember that substBndr and friends are used when doing expression
218 substitution only. Their only business is substitution, so they
219 preserve all IdInfo (suitably substituted). For example, we *want* to
220 preserve occ info in rules.
223 substBndr :: Subst -> Var -> (Subst, Var)
225 | isTyVar bndr = substTyVarBndr subst bndr
226 | otherwise = substIdBndr subst subst bndr
228 substBndrs :: Subst -> [Var] -> (Subst, [Var])
229 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
231 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
232 -- Substitute a mutually recursive group
233 substRecBndrs subst bndrs
234 = (new_subst, new_bndrs)
235 where -- Here's the reason we need to pass rec_subst to subst_id
236 (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
241 substIdBndr :: Subst -- Substitution to use for the IdInfo
242 -> Subst -> Id -- Substitition and Id to transform
243 -> (Subst, Id) -- Transformed pair
245 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
246 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
248 id1 = uniqAway in_scope old_id -- id1 is cloned if necessary
249 id2 = substIdType subst id1 -- id2 has its type zapped
251 -- new_id has the right IdInfo
252 -- The lazy-set is because we're in a loop here, with
253 -- rec_subst, when dealing with a mutually-recursive group
254 new_id = maybeModifyIdInfo (substIdInfo rec_subst) id2
256 -- Extend the substitution if the unique has changed
257 -- See the notes with substTyVarBndr for the delSubstEnv
258 new_env | new_id /= old_id = extendVarEnv env old_id (Var new_id)
259 | otherwise = delVarEnv env old_id
262 Now a variant that unconditionally allocates a new unique.
263 It also unconditionally zaps the OccInfo.
266 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
267 cloneIdBndr subst us old_id
268 = clone_id subst subst (old_id, uniqFromSupply us)
270 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
271 cloneIdBndrs subst us ids
272 = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
274 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
275 cloneRecIdBndrs subst us ids
278 (subst', ids') = mapAccumL (clone_id subst') subst
279 (ids `zip` uniqsFromSupply us)
281 -- Just like substIdBndr, except that it always makes a new unique
282 -- It is given the unique to use
283 clone_id :: Subst -- Substitution for the IdInfo
284 -> Subst -> (Id, Unique) -- Substitition and Id to transform
285 -> (Subst, Id) -- Transformed pair
287 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
288 = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
290 id1 = setVarUnique old_id uniq
291 id2 = substIdType subst id1
292 new_id = maybeModifyIdInfo (substIdInfo rec_subst) id2
293 new_env = extendVarEnv env old_id (Var new_id)
297 %************************************************************************
301 %************************************************************************
303 For types we just call the corresponding function in Type, but we have
304 to repackage the substitution, from a Subst to a TvSubst
307 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
308 substTyVarBndr (Subst in_scope id_env tv_env) tv
309 = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
310 (TvSubst in_scope' tv_env', tv')
311 -> (Subst in_scope' id_env tv_env', tv')
313 substTy :: Subst -> Type -> Type
314 substTy (Subst in_scope id_env tv_env) ty
315 = Type.substTy (TvSubst in_scope tv_env) ty
319 %************************************************************************
321 \section{IdInfo substitution}
323 %************************************************************************
326 substIdType :: Subst -> Id -> Id
327 substIdType subst@(Subst in_scope id_env tv_env) id
328 | isEmptyVarEnv tv_env || isEmptyVarSet (tyVarsOfType old_ty) = id
329 | otherwise = setIdType id (substTy subst old_ty)
330 -- The tyVarsOfType is cheaper than it looks
331 -- because we cache the free tyvars of the type
332 -- in a Note in the id's type itself
337 substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
338 -- Always zaps the unfolding, to save substitution work
339 substIdInfo subst info
340 | nothing_to_do = Nothing
341 | otherwise = Just (info `setSpecInfo` substSpec subst old_rules
342 `setWorkerInfo` substWorker subst old_wrkr
343 `setUnfoldingInfo` noUnfolding)
345 old_rules = specInfo info
346 old_wrkr = workerInfo info
347 nothing_to_do = isEmptySpecInfo old_rules &&
348 not (workerExists old_wrkr) &&
349 not (hasUnfolding (unfoldingInfo info))
353 substWorker :: Subst -> WorkerInfo -> WorkerInfo
354 -- Seq'ing on the returned WorkerInfo is enough to cause all the
355 -- substitutions to happen completely
357 substWorker subst NoWorker
359 substWorker subst (HasWorker w a)
360 = case lookupIdSubst subst w of
361 Var w1 -> HasWorker w1 a
362 other -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
363 NoWorker -- Worker has got substituted away altogether
364 -- (This can happen if it's trivial,
365 -- via postInlineUnconditionally, hence warning)
368 substSpec :: Subst -> SpecInfo -> SpecInfo
370 substSpec subst spec@(SpecInfo rules rhs_fvs)
374 = seqSpecInfo new_rules `seq` new_rules
376 new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
378 do_subst rule@(BuiltinRule {}) = rule
379 do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
380 = rule { ru_bndrs = bndrs',
381 ru_args = map (substExpr subst') args,
382 ru_rhs = substExpr subst' rhs }
384 (subst', bndrs') = substBndrs subst bndrs
387 substVarSet subst fvs
388 = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
391 | isId fv = exprFreeVars (lookupIdSubst subst fv)
392 | otherwise = tyVarsOfType (lookupTvSubst subst fv)