2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 \section[CoreUtils]{Utility functions on @Core@ syntax}
9 InScopeSet, emptyInScopeSet,
10 lookupInScope, setInScope, extendInScope, extendInScopes, isInScope,
13 Subst, TyVarSubst, IdSubst,
14 emptySubst, mkSubst, substEnv, substInScope,
15 lookupSubst, isEmptySubst, extendSubst, extendSubstList,
16 zapSubstEnv, setSubstEnv,
18 bindSubst, unBindSubst, bindSubstList, unBindSubstList,
21 substBndr, substBndrs, substTyVar, substId, substIds,
22 substAndCloneId, substAndCloneIds,
25 mkTyVarSubst, mkTopTyVarSubst,
29 substExpr, substIdInfo
32 #include "HsVersions.h"
34 import CoreSyn ( Expr(..), Bind(..), Note(..), CoreExpr, CoreBndr,
35 CoreRules(..), CoreRule(..),
36 emptyCoreRules, isEmptyCoreRules, seqRules
38 import CoreFVs ( exprFreeVars )
39 import Type ( Type(..), ThetaType, TyNote(..),
40 tyVarsOfType, tyVarsOfTypes, mkAppTy
44 import Var ( setVarUnique, isId )
45 import Id ( idType, setIdType )
46 import IdInfo ( IdInfo, zapFragileIdInfo,
47 specInfo, setSpecInfo,
48 workerExists, workerInfo, setWorkerInfo, WorkerInfo
50 import UniqSupply ( UniqSupply, uniqFromSupply, splitUniqSupply )
51 import Var ( Var, IdOrTyVar, Id, TyVar, isTyVar, maybeModifyIdInfo )
53 import Util ( mapAccumL, foldl2, seqList, ($!) )
56 %************************************************************************
58 \subsection{Substitutions}
60 %************************************************************************
63 type InScopeSet = VarSet
65 data Subst = Subst InScopeSet -- In scope
66 SubstEnv -- Substitution itself
67 -- INVARIANT 1: The in-scope set is a superset
68 -- of the free vars of the range of the substitution
69 -- that might possibly clash with locally-bound variables
70 -- in the thing being substituted in.
71 -- This is what lets us deal with name capture properly
72 -- It's a hard invariant to check...
73 -- There are various ways of causing it to happen:
74 -- - 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
76 -- - make it empty because all the free vars of the subst are fresh,
77 -- and hence can't possibly clash.a
79 -- INVARIANT 2: No variable is both in scope and in the domain of the substitution
80 -- Equivalently, the substitution is idempotent
87 emptyInScopeSet :: InScopeSet
88 emptyInScopeSet = emptyVarSet
94 isEmptySubst :: Subst -> Bool
95 isEmptySubst (Subst _ env) = isEmptySubstEnv env
98 emptySubst = Subst emptyVarSet emptySubstEnv
100 mkSubst :: InScopeSet -> SubstEnv -> Subst
101 mkSubst in_scope env = Subst in_scope env
103 substEnv :: Subst -> SubstEnv
104 substEnv (Subst _ env) = env
106 substInScope :: Subst -> InScopeSet
107 substInScope (Subst in_scope _) = in_scope
109 zapSubstEnv :: Subst -> Subst
110 zapSubstEnv (Subst in_scope env) = Subst in_scope emptySubstEnv
112 extendSubst :: Subst -> Var -> SubstResult -> Subst
113 extendSubst (Subst in_scope env) v r = Subst in_scope (extendSubstEnv env v r)
115 extendSubstList :: Subst -> [Var] -> [SubstResult] -> Subst
116 extendSubstList (Subst in_scope env) v r = Subst in_scope (extendSubstEnvList env v r)
118 lookupSubst :: Subst -> Var -> Maybe SubstResult
119 lookupSubst (Subst _ env) v = lookupSubstEnv env v
121 lookupInScope :: Subst -> Var -> Maybe Var
122 lookupInScope (Subst in_scope _) v = lookupVarSet in_scope v
124 isInScope :: Var -> Subst -> Bool
125 isInScope v (Subst in_scope _) = v `elemVarSet` in_scope
127 extendInScope :: Subst -> Var -> Subst
128 extendInScope (Subst in_scope env) v = Subst (extendVarSet in_scope v) env
130 extendInScopes :: Subst -> [Var] -> Subst
131 extendInScopes (Subst in_scope env) vs = Subst (foldl extendVarSet in_scope vs) env
133 -------------------------------
134 bindSubst :: Subst -> Var -> Var -> Subst
135 -- Extend with a substitution, v1 -> Var v2
136 -- and extend the in-scopes with v2
137 bindSubst (Subst in_scope env) old_bndr new_bndr
138 = Subst (in_scope `extendVarSet` new_bndr)
139 (extendSubstEnv env old_bndr subst_result)
141 subst_result | isId old_bndr = DoneEx (Var new_bndr)
142 | otherwise = DoneTy (TyVarTy new_bndr)
144 unBindSubst :: Subst -> Var -> Var -> Subst
145 -- Reverse the effect of bindSubst
146 -- If old_bndr was already in the substitution, this doesn't quite work
147 unBindSubst (Subst in_scope env) old_bndr new_bndr
148 = Subst (in_scope `delVarSet` new_bndr) (delSubstEnv env old_bndr)
150 -- And the "List" forms
151 bindSubstList :: Subst -> [Var] -> [Var] -> Subst
152 bindSubstList subst old_bndrs new_bndrs
153 = foldl2 bindSubst subst old_bndrs new_bndrs
155 unBindSubstList :: Subst -> [Var] -> [Var] -> Subst
156 unBindSubstList subst old_bndrs new_bndrs
157 = foldl2 unBindSubst subst old_bndrs new_bndrs
160 -------------------------------
161 setInScope :: Subst -- Take env part from here
164 setInScope (Subst in_scope1 env1) in_scope2
165 = ASSERT( in_scope1 `subVarSet` in_scope1 )
168 setSubstEnv :: Subst -- Take in-scope part from here
169 -> SubstEnv -- ... and env part from here
171 setSubstEnv (Subst in_scope1 _) env2 = Subst in_scope1 env2
175 %************************************************************************
177 \subsection{Type substitution}
179 %************************************************************************
182 type TyVarSubst = Subst -- TyVarSubst are expected to have range elements
183 -- (We could have a variant of Subst, but it doesn't seem worth it.)
185 -- mkTyVarSubst generates the in-scope set from
186 -- the types given; but it's just a thunk so with a bit of luck
187 -- it'll never be evaluated
188 mkTyVarSubst :: [TyVar] -> [Type] -> Subst
189 mkTyVarSubst tyvars tys = Subst (tyVarsOfTypes tys) (zip_ty_env tyvars tys emptySubstEnv)
191 -- mkTopTyVarSubst is called when doing top-level substitutions.
192 -- Here we expect that the free vars of the range of the
193 -- substitution will be empty.
194 mkTopTyVarSubst :: [TyVar] -> [Type] -> Subst
195 mkTopTyVarSubst tyvars tys = Subst emptyVarSet (zip_ty_env tyvars tys emptySubstEnv)
197 zip_ty_env [] [] env = env
198 zip_ty_env (tv:tvs) (ty:tys) env = zip_ty_env tvs tys (extendSubstEnv env tv (DoneTy ty))
201 substTy works with general Substs, so that it can be called from substExpr too.
204 substTy :: Subst -> Type -> Type
205 substTy subst ty | isEmptySubst subst = ty
206 | otherwise = subst_ty subst ty
208 substTheta :: TyVarSubst -> ThetaType -> ThetaType
209 substTheta subst theta
210 | isEmptySubst subst = theta
211 | otherwise = [(clas, map (subst_ty subst) tys) | (clas, tys) <- theta]
216 go (TyConApp tc tys) = let args = map go tys
217 in args `seqList` TyConApp tc args
218 go (NoteTy (SynNote ty1) ty2) = NoteTy (SynNote $! (go ty1)) $! (go ty2)
219 go (NoteTy (FTVNote _) ty2) = go ty2 -- Discard the free tyvar note
220 go (FunTy arg res) = (FunTy $! (go arg)) $! (go res)
221 go (NoteTy (UsgNote usg) ty2) = (NoteTy $! UsgNote usg) $! go ty2 -- Keep usage annot
222 go (AppTy fun arg) = mkAppTy (go fun) $! (go arg)
223 go ty@(TyVarTy tv) = case (lookupSubst subst tv) of
225 Just (DoneTy ty') -> ty'
227 go (ForAllTy tv ty) = case substTyVar subst tv of
228 (subst', tv') -> ForAllTy tv' $! (subst_ty subst' ty)
231 Here is where we invent a new binder if necessary.
234 substTyVar :: Subst -> TyVar -> (Subst, TyVar)
235 substTyVar subst@(Subst in_scope env) old_var
236 | old_var == new_var -- No need to clone
237 -- But we *must* zap any current substitution for the variable.
239 -- (\x.e) with id_subst = [x |-> e']
240 -- Here we must simply zap the substitution for x
242 -- The new_id isn't cloned, but it may have a different type
243 -- etc, so we must return it, not the old id
244 = (Subst (in_scope `extendVarSet` new_var)
245 (delSubstEnv env old_var),
248 | otherwise -- The new binder is in scope so
249 -- we'd better rename it away from the in-scope variables
250 -- Extending the substitution to do this renaming also
251 -- has the (correct) effect of discarding any existing
252 -- substitution for that variable
253 = (Subst (in_scope `extendVarSet` new_var)
254 (extendSubstEnv env old_var (DoneTy (TyVarTy new_var))),
257 new_var = uniqAway in_scope old_var
258 -- The uniqAway part makes sure the new variable is not already in scope
262 %************************************************************************
264 \section{Expression substitution}
266 %************************************************************************
268 This expression substituter deals correctly with name capture.
270 BUT NOTE that substExpr silently discards the
273 IdInfo attached to any binders in the expression. It's quite
274 tricky to do them 'right' in the case of mutually recursive bindings,
275 and so far has proved unnecessary.
278 substExpr :: Subst -> CoreExpr -> CoreExpr
279 substExpr subst expr | isEmptySubst subst = expr
280 | otherwise = subst_expr subst expr
282 subst_expr subst expr
285 go (Var v) = case lookupSubst subst v of
286 Just (DoneEx e') -> e'
287 Just (ContEx env' e') -> subst_expr (setSubstEnv subst env') e'
288 -- NO! NO! SLPJ 14 July 99
289 Nothing -> case lookupInScope subst v of
292 -- NB: we look up in the in_scope set because the variable
293 -- there may have more info. In particular, when substExpr
294 -- is called from the simplifier, the type inside the *occurrences*
295 -- of a variable may not be right; we should replace it with the
296 -- binder, from the in_scope set.
300 go (Type ty) = Type (go_ty ty)
301 go (Con con args) = Con con (map go args)
302 go (App fun arg) = App (go fun) (go arg)
303 go (Note note e) = Note (go_note note) (go e)
305 go (Lam bndr body) = Lam bndr' (subst_expr subst' body)
307 (subst', bndr') = substBndr subst bndr
309 go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (subst_expr subst' body)
311 (subst', bndr') = substBndr subst bndr
313 go (Let (Rec pairs) body) = Let (Rec pairs') (subst_expr subst' body)
315 (subst', bndrs') = substBndrs subst (map fst pairs)
316 pairs' = bndrs' `zip` rhss'
317 rhss' = map (subst_expr subst' . snd) pairs
319 go (Case scrut bndr alts) = Case (go scrut) bndr' (map (go_alt subst') alts)
321 (subst', bndr') = substBndr subst bndr
323 go_alt subst (con, bndrs, rhs) = (con, bndrs', subst_expr subst' rhs)
325 (subst', bndrs') = substBndrs subst bndrs
327 go_note (Coerce ty1 ty2) = Coerce (go_ty ty1) (go_ty ty2)
330 go_ty ty = substTy subst ty
334 Substituting in binders is a rather tricky part of the whole compiler.
336 When we hit a binder we may need to
337 (a) apply the the type envt (if non-empty) to its type
338 (b) apply the type envt and id envt to its SpecEnv (if it has one)
339 (c) give it a new unique to avoid name clashes
342 substBndr :: Subst -> IdOrTyVar -> (Subst, IdOrTyVar)
344 | isTyVar bndr = substTyVar subst bndr
345 | otherwise = substId subst bndr
347 substBndrs :: Subst -> [IdOrTyVar] -> (Subst, [IdOrTyVar])
348 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
351 substIds :: Subst -> [Id] -> (Subst, [Id])
352 substIds subst bndrs = mapAccumL substId subst bndrs
354 substId :: Subst -> Id -> (Subst, Id)
356 -- Returns an Id with empty unfolding and spec-env.
357 -- It's up to the caller to sort these out.
359 substId subst@(Subst in_scope env) old_id
360 = (Subst (in_scope `extendVarSet` new_id)
361 (extendSubstEnv env old_id (DoneEx (Var new_id))),
364 id_ty = idType old_id
366 -- id1 has its type zapped
367 id1 | noTypeSubst env
368 || isEmptyVarSet (tyVarsOfType id_ty) = old_id
369 -- The tyVarsOfType is cheaper than it looks
370 -- because we cache the free tyvars of the type
371 -- in a Note in the id's type itself
372 | otherwise = setIdType old_id (substTy subst id_ty)
374 -- id2 has its fragile IdInfo zapped
375 id2 = maybeModifyIdInfo zapFragileIdInfo id1
377 -- new_id is cloned if necessary
378 new_id = uniqAway in_scope id2
381 Now a variant that unconditionally allocates a new unique.
384 substAndCloneIds :: Subst -> UniqSupply -> [Id] -> (Subst, UniqSupply, [Id])
385 substAndCloneIds subst us [] = (subst, us, [])
386 substAndCloneIds subst us (b:bs) = case substAndCloneId subst us b of { (subst1, us1, b') ->
387 case substAndCloneIds subst1 us1 bs of { (subst2, us2, bs') ->
388 (subst2, us2, (b':bs')) }}
390 substAndCloneId :: Subst -> UniqSupply -> Id -> (Subst, UniqSupply, Id)
391 substAndCloneId subst@(Subst in_scope env) us old_id
392 = (Subst (in_scope `extendVarSet` new_id)
393 (extendSubstEnv env old_id (DoneEx (Var new_id))),
397 id_ty = idType old_id
398 id1 | noTypeSubst env || isEmptyVarSet (tyVarsOfType id_ty) = old_id
399 | otherwise = setIdType old_id (substTy subst id_ty)
401 id2 = maybeModifyIdInfo zapFragileIdInfo id1
402 new_id = setVarUnique id2 (uniqFromSupply us1)
403 (us1,new_us) = splitUniqSupply us
407 %************************************************************************
409 \section{IdInfo substitution}
411 %************************************************************************
415 -> IdInfo -- Get un-substituted ones from here
416 -> IdInfo -- Substitute it and add it to here
417 -> IdInfo -- To give this
418 -- Seq'ing on the returned IdInfo is enough to cause all the
419 -- substitutions to happen completely
421 substIdInfo subst old_info new_info
424 info1 | isEmptyCoreRules old_rules = new_info
425 | otherwise = new_info `setSpecInfo` new_rules
426 -- setSpecInfo does a seq
428 new_rules = substRules subst old_rules
430 info2 | not (workerExists old_wrkr) = info1
431 | otherwise = info1 `setWorkerInfo` new_wrkr
432 -- setWorkerInfo does a seq
434 new_wrkr = substWorker subst old_wrkr
436 old_rules = specInfo old_info
437 old_wrkr = workerInfo old_info
439 substWorker :: Subst -> WorkerInfo -> WorkerInfo
440 -- Seq'ing on the returned WorkerInfo is enough to cause all the
441 -- substitutions to happen completely
443 substWorker subst Nothing
445 substWorker subst (Just w)
446 = case lookupSubst subst w of
448 Just (DoneEx (Var w1)) -> Just w1
449 Just (DoneEx other) -> WARN( True, text "substWorker: DoneEx" <+> ppr w )
450 Nothing -- Worker has got substituted away altogether
451 Just (ContEx se1 e) -> WARN( True, text "substWorker: ContEx" <+> ppr w )
454 substRules :: Subst -> CoreRules -> CoreRules
455 -- Seq'ing on the returned CoreRules is enough to cause all the
456 -- substitutions to happen completely
458 substRules subst rules
459 | isEmptySubst subst = rules
461 substRules subst (Rules rules rhs_fvs)
462 = seqRules new_rules `seq` new_rules
464 new_rules = Rules (map do_subst rules)
465 (subst_fvs (substEnv subst) rhs_fvs)
467 do_subst (Rule name tpl_vars lhs_args rhs)
468 = Rule name tpl_vars'
469 (map (substExpr subst') lhs_args)
470 (substExpr subst' rhs)
472 (subst', tpl_vars') = substBndrs subst tpl_vars
475 = foldVarSet (unionVarSet . subst_fv) emptyVarSet rhs_fvs
477 subst_fv fv = case lookupSubstEnv se fv of
478 Nothing -> unitVarSet fv
479 Just (DoneEx expr) -> exprFreeVars expr
480 Just (DoneTy ty) -> tyVarsOfType ty
481 Just (ContEx se' expr) -> subst_fvs se' (exprFreeVars expr)