Comments only
[ghc-hetmet.git] / compiler / coreSyn / CoreSubst.lhs
1 %
2 % (c) The University of Glasgow 2006
3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 %
5
6 Utility functions on @Core@ syntax
7
8 \begin{code}
9 module CoreSubst (
10         -- * Main data types
11         Subst, TvSubstEnv, IdSubstEnv, InScopeSet,
12
13         -- ** Substituting into expressions and related types
14         deShadowBinds,
15         substTy, substExpr, substSpec, substWorker,
16         lookupIdSubst, lookupTvSubst, 
17
18         -- ** Operations on substitutions
19         emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst, 
20         extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
21         extendSubst, extendSubstList, zapSubstEnv,
22         extendInScope, extendInScopeList, extendInScopeIds, 
23         isInScope,
24
25         -- ** Substituting and cloning binders
26         substBndr, substBndrs, substRecBndrs,
27         cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
28     ) where
29
30 #include "HsVersions.h"
31
32 import CoreSyn
33 import CoreFVs
34 import CoreUtils
35
36 import qualified Type
37 import Type     ( Type, TvSubst(..), TvSubstEnv )
38 import VarSet
39 import VarEnv
40 import Id
41 import Var      ( Var, TyVar, setVarUnique )
42 import IdInfo
43 import Unique
44 import UniqSupply
45 import Maybes
46 import Outputable
47 import PprCore          ()              -- Instances
48 import FastString
49
50 import Data.List
51 \end{code}
52
53
54 %************************************************************************
55 %*                                                                      *
56 \subsection{Substitutions}
57 %*                                                                      *
58 %************************************************************************
59
60 \begin{code}
61 -- | A substitution environment, containing both 'Id' and 'TyVar' substitutions.
62 --
63 -- Some invariants apply to how you use the substitution:
64 --
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.
68 --
69 -- 2. #apply_once# You may apply the substitution only /once/
70 --
71 -- There are various ways of setting up the in-scope set such that the first of these invariants hold:
72 --
73 -- * Arrange that the in-scope set really is all the things in scope
74 --
75 -- * Arrange that it's the free vars of the range of the substitution
76 --
77 -- * Make it empty, if you know that all the free vars of the substitution are fresh, and hence can't possibly clash
78 data Subst 
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
83
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...
87         --
88         -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
89         --              Types.TvSubstEnv
90         --
91         -- INVARIANT 3: See Note [Extending the Subst]
92
93 {-
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.  
98
99 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
100
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
104
105 In consequence:
106
107 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
108
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.)
113
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.
120
121 Why do we make a different choice for the IdSubstEnv than the TvSubstEnv?
122
123 * For Ids, we change the IdInfo all the time (e.g. deleting the
124   unfolding), and adding it back later, so using the TyVar convention
125   would entail extending the substitution almost all the time
126
127 * The simplifier wants to look up in the in-scope set anyway, in case it 
128   can see a better unfolding from an enclosing case expression
129
130 * For TyVars, only coercion variables can possibly change, and they are 
131   easy to spot
132 -}
133
134 -- | An environment for substituting for 'Id's
135 type IdSubstEnv = IdEnv CoreExpr
136
137 ----------------------------
138 isEmptySubst :: Subst -> Bool
139 isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env
140
141 emptySubst :: Subst
142 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv
143
144 mkEmptySubst :: InScopeSet -> Subst
145 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv
146
147 mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
148 mkSubst in_scope tvs ids = Subst in_scope ids tvs
149
150 -- getTvSubst :: Subst -> TvSubst
151 -- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env
152
153 -- getTvSubstEnv :: Subst -> TvSubstEnv
154 -- getTvSubstEnv (Subst _ _ tv_env) = tv_env
155 -- 
156 -- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
157 -- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs
158
159 -- | Find the in-scope set: see "CoreSubst#in_scope_invariant"
160 substInScope :: Subst -> InScopeSet
161 substInScope (Subst in_scope _ _) = in_scope
162
163 -- | Remove all substitutions for 'Id's and 'Var's that might have been built up
164 -- while preserving the in-scope set
165 zapSubstEnv :: Subst -> Subst
166 zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv
167
168 -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the in-scope set is
169 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
170 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
171 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
172 extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs
173
174 -- | Adds multiple 'Id' substitutions to the 'Subst': see also 'extendIdSubst'
175 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
176 extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs
177
178 -- | Add a substitution for a 'TyVar' to the 'Subst': you must ensure that the in-scope set is
179 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
180 extendTvSubst :: Subst -> TyVar -> Type -> Subst
181 extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r) 
182
183 -- | Adds multiple 'TyVar' substitutions to the 'Subst': see also 'extendTvSubst'
184 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
185 extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)
186
187 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to the 'Var' being added. See also
188 -- 'extendIdSubst' and 'extendTvSubst'
189 extendSubst :: Subst -> Var -> CoreArg -> Subst
190 extendSubst (Subst in_scope ids tvs) tv (Type ty)
191   = ASSERT( isTyVar tv ) Subst in_scope ids (extendVarEnv tvs tv ty)
192 extendSubst (Subst in_scope ids tvs) id expr
193   = ASSERT( isId id ) Subst in_scope (extendVarEnv ids id expr) tvs
194
195 -- | Add a substitution for a 'TyVar' or 'Id' as appropriate to all the 'Var's being added. See also 'extendSubst'
196 extendSubstList :: Subst -> [(Var,CoreArg)] -> Subst
197 extendSubstList subst []              = subst
198 extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var rhs) prs
199
200 -- | Find the substitution for an 'Id' in the 'Subst'
201 lookupIdSubst :: Subst -> Id -> CoreExpr
202 lookupIdSubst (Subst in_scope ids _) v
203   | not (isLocalId v) = Var v
204   | Just e  <- lookupVarEnv ids       v = e
205   | Just v' <- lookupInScope in_scope v = Var v'
206         -- Vital! See Note [Extending the Subst]
207   | otherwise = WARN( True, ptext (sLit "CoreSubst.lookupIdSubst") <+> ppr v ) 
208                 Var v
209
210 -- | Find the substitution for a 'TyVar' in the 'Subst'
211 lookupTvSubst :: Subst -> TyVar -> Type
212 lookupTvSubst (Subst _ _ tvs) v = lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
213
214 ------------------------------
215 isInScope :: Var -> Subst -> Bool
216 isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope
217
218 -- | Add the 'Var' to the in-scope set: as a side effect, removes any existing substitutions for it
219 extendInScope :: Subst -> Var -> Subst
220 extendInScope (Subst in_scope ids tvs) v
221   = Subst (in_scope `extendInScopeSet` v) 
222           (ids `delVarEnv` v) (tvs `delVarEnv` v)
223
224 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
225 extendInScopeList :: Subst -> [Var] -> Subst
226 extendInScopeList (Subst in_scope ids tvs) vs
227   = Subst (in_scope `extendInScopeSetList` vs) 
228           (ids `delVarEnvList` vs) (tvs `delVarEnvList` vs)
229
230 -- | Optimized version of 'extendInScopeList' that can be used if you are certain 
231 -- all the things being added are 'Id's and hence none are 'TyVar's
232 extendInScopeIds :: Subst -> [Id] -> Subst
233 extendInScopeIds (Subst in_scope ids tvs) vs 
234   = Subst (in_scope `extendInScopeSetList` vs) 
235           (ids `delVarEnvList` vs) tvs
236 \end{code}
237
238 Pretty printing, for debugging only
239
240 \begin{code}
241 instance Outputable Subst where
242   ppr (Subst in_scope ids tvs) 
243         =  ptext (sLit "<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
244         $$ ptext (sLit " IdSubst   =") <+> ppr ids
245         $$ ptext (sLit " TvSubst   =") <+> ppr tvs
246          <> char '>'
247 \end{code}
248
249
250 %************************************************************************
251 %*                                                                      *
252         Substituting expressions
253 %*                                                                      *
254 %************************************************************************
255
256 \begin{code}
257 -- | Apply a substititon to an entire 'CoreExpr'. Rememeber, you may only 
258 -- apply the substitution /once/: see "CoreSubst#apply_once"
259 substExpr :: Subst -> CoreExpr -> CoreExpr
260 substExpr subst expr
261   = go expr
262   where
263     go (Var v)         = lookupIdSubst subst v 
264     go (Type ty)       = Type (substTy subst ty)
265     go (Lit lit)       = Lit lit
266     go (App fun arg)   = App (go fun) (go arg)
267     go (Note note e)   = Note (go_note note) (go e)
268     go (Cast e co)     = Cast (go e) (substTy subst co)
269     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
270                        where
271                          (subst', bndr') = substBndr subst bndr
272
273     go (Let bind body) = Let bind' (substExpr subst' body)
274                        where
275                          (subst', bind') = substBind subst bind
276
277     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
278                                  where
279                                  (subst', bndr') = substBndr subst bndr
280
281     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
282                                  where
283                                    (subst', bndrs') = substBndrs subst bndrs
284
285     go_note note             = note
286
287 -- | Apply a substititon to an entire 'CoreBind', additionally returning an updated 'Subst'
288 -- that should be used by subsequent substitutons.
289 substBind :: Subst -> CoreBind -> (Subst, CoreBind)
290 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (substExpr subst rhs))
291                                   where
292                                     (subst', bndr') = substBndr subst bndr
293
294 substBind subst (Rec pairs) = (subst', Rec pairs')
295                             where
296                                 (subst', bndrs') = substRecBndrs subst (map fst pairs)
297                                 pairs'  = bndrs' `zip` rhss'
298                                 rhss'   = map (substExpr subst' . snd) pairs
299 \end{code}
300
301 \begin{code}
302 -- | De-shadowing the program is sometimes a useful pre-pass. It can be done simply
303 -- by running over the bindings with an empty substitution, becuase substitution
304 -- returns a result that has no-shadowing guaranteed.
305 --
306 -- (Actually, within a single /type/ there might still be shadowing, because 
307 -- 'substTy' is a no-op for the empty substitution, but that's probably OK.)
308 deShadowBinds :: [CoreBind] -> [CoreBind]
309 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
310 \end{code}
311
312
313 %************************************************************************
314 %*                                                                      *
315         Substituting binders
316 %*                                                                      *
317 %************************************************************************
318
319 Remember that substBndr and friends are used when doing expression
320 substitution only.  Their only business is substitution, so they
321 preserve all IdInfo (suitably substituted).  For example, we *want* to
322 preserve occ info in rules.
323
324 \begin{code}
325 -- | Substitutes a 'Var' for another one according to the 'Subst' given, returning
326 -- the result and an updated 'Subst' that should be used by subsequent substitutons.
327 -- 'IdInfo' is preserved by this process, although it is substituted into appropriately.
328 substBndr :: Subst -> Var -> (Subst, Var)
329 substBndr subst bndr
330   | isTyVar bndr  = substTyVarBndr subst bndr
331   | otherwise     = substIdBndr subst subst bndr
332
333 -- | Applies 'substBndr' to a number of 'Var's, accumulating a new 'Subst' left-to-right
334 substBndrs :: Subst -> [Var] -> (Subst, [Var])
335 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
336
337 -- | Substitute in a mutually recursive group of 'Id's
338 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
339 substRecBndrs subst bndrs 
340   = (new_subst, new_bndrs)
341   where         -- Here's the reason we need to pass rec_subst to subst_id
342     (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
343 \end{code}
344
345
346 \begin{code}
347 substIdBndr :: Subst            -- ^ Substitution to use for the IdInfo
348             -> Subst -> Id      -- ^ Substitition and Id to transform
349             -> (Subst, Id)      -- ^ Transformed pair
350                                 -- NB: unfolding may be zapped
351
352 substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
353   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
354   where
355     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
356     id2 | no_type_change = id1
357         | otherwise      = setIdType id1 (substTy subst old_ty)
358
359     old_ty = idType old_id
360     no_type_change = isEmptyVarEnv tvs || 
361                      isEmptyVarSet (Type.tyVarsOfType old_ty)
362
363         -- new_id has the right IdInfo
364         -- The lazy-set is because we're in a loop here, with 
365         -- rec_subst, when dealing with a mutually-recursive group
366     new_id = maybeModifyIdInfo mb_new_info id2
367     mb_new_info = substIdInfo rec_subst id2 (idInfo id2)
368         -- NB: unfolding info may be zapped
369
370         -- Extend the substitution if the unique has changed
371         -- See the notes with substTyVarBndr for the delVarEnv
372     new_env | no_change = delVarEnv env old_id
373             | otherwise = extendVarEnv env old_id (Var new_id)
374
375     no_change = id1 == old_id
376         -- See Note [Extending the Subst]
377         -- it's /not/ necessary to check mb_new_info and no_type_change
378 \end{code}
379
380 Now a variant that unconditionally allocates a new unique.
381 It also unconditionally zaps the OccInfo.
382
383 \begin{code}
384 -- | Very similar to 'substBndr', but it always allocates a new 'Unique' for
385 -- each variable in its output and removes all 'IdInfo'
386 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
387 cloneIdBndr subst us old_id
388   = clone_id subst subst (old_id, uniqFromSupply us)
389
390 -- | Applies 'cloneIdBndr' to a number of 'Id's, accumulating a final
391 -- substitution from left to right
392 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
393 cloneIdBndrs subst us ids
394   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
395
396 -- | Clone a mutually recursive group of 'Id's
397 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
398 cloneRecIdBndrs subst us ids
399   = (subst', ids')
400   where
401     (subst', ids') = mapAccumL (clone_id subst') subst
402                                (ids `zip` uniqsFromSupply us)
403
404 -- Just like substIdBndr, except that it always makes a new unique
405 -- It is given the unique to use
406 clone_id    :: Subst                    -- Substitution for the IdInfo
407             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
408             -> (Subst, Id)              -- Transformed pair
409
410 clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
411   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
412   where
413     id1     = setVarUnique old_id uniq
414     id2     = substIdType subst id1
415     new_id  = maybeModifyIdInfo (substIdInfo rec_subst id2 (idInfo old_id)) id2
416     new_env = extendVarEnv env old_id (Var new_id)
417 \end{code}
418
419
420 %************************************************************************
421 %*                                                                      *
422                 Types
423 %*                                                                      *
424 %************************************************************************
425
426 For types we just call the corresponding function in Type, but we have
427 to repackage the substitution, from a Subst to a TvSubst
428
429 \begin{code}
430 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
431 substTyVarBndr (Subst in_scope id_env tv_env) tv
432   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
433         (TvSubst in_scope' tv_env', tv') 
434            -> (Subst in_scope' id_env tv_env', tv')
435
436 -- | See 'Type.substTy'
437 substTy :: Subst -> Type -> Type 
438 substTy (Subst in_scope _id_env tv_env) ty
439   = Type.substTy (TvSubst in_scope tv_env) ty
440 \end{code}
441
442
443 %************************************************************************
444 %*                                                                      *
445 \section{IdInfo substitution}
446 %*                                                                      *
447 %************************************************************************
448
449 \begin{code}
450 substIdType :: Subst -> Id -> Id
451 substIdType subst@(Subst _ _ tv_env) id
452   | isEmptyVarEnv tv_env || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
453   | otherwise   = setIdType id (substTy subst old_ty)
454                 -- The tyVarsOfType is cheaper than it looks
455                 -- because we cache the free tyvars of the type
456                 -- in a Note in the id's type itself
457   where
458     old_ty = idType id
459
460 ------------------
461 -- | Substitute into some 'IdInfo' with regard to the supplied new 'Id'.
462 -- Always zaps the unfolding, to save substitution work
463 substIdInfo :: Subst -> Id -> IdInfo -> Maybe IdInfo
464 substIdInfo subst new_id info
465   | nothing_to_do = Nothing
466   | otherwise     = Just (info `setSpecInfo`      substSpec subst new_id old_rules
467                                `setWorkerInfo`    substWorker subst old_wrkr
468                                `setUnfoldingInfo` noUnfolding)
469   where
470     old_rules     = specInfo info
471     old_wrkr      = workerInfo info
472     nothing_to_do = isEmptySpecInfo old_rules &&
473                     not (workerExists old_wrkr) &&
474                     not (hasUnfolding (unfoldingInfo info))
475     
476
477 ------------------
478 -- | Substitutes for the 'Id's within the 'WorkerInfo'
479 substWorker :: Subst -> WorkerInfo -> WorkerInfo
480         -- Seq'ing on the returned WorkerInfo is enough to cause all the 
481         -- substitutions to happen completely
482
483 substWorker _ NoWorker
484   = NoWorker
485 substWorker subst (HasWorker w a)
486   = case lookupIdSubst subst w of
487         Var w1 -> HasWorker w1 a
488         other  -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
489                   NoWorker      -- Worker has got substituted away altogether
490                                 -- (This can happen if it's trivial, 
491                                 --  via postInlineUnconditionally, hence warning)
492
493 ------------------
494 -- | Substitutes for the 'Id's within the 'WorkerInfo' given the new function 'Id'
495 substSpec :: Subst -> Id -> SpecInfo -> SpecInfo
496
497 substSpec subst new_fn spec@(SpecInfo rules rhs_fvs)
498   | isEmptySubst subst
499   = spec
500   | otherwise
501   = seqSpecInfo new_rules `seq` new_rules
502   where
503     new_name = idName new_fn
504     new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
505
506     do_subst rule@(BuiltinRule {}) = rule
507     do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
508         = rule { ru_bndrs = bndrs', 
509                  ru_fn = new_name,      -- Important: the function may have changed its name!
510                  ru_args  = map (substExpr subst') args,
511                  ru_rhs   = substExpr subst' rhs }
512         where
513           (subst', bndrs') = substBndrs subst bndrs
514
515 ------------------
516 substVarSet :: Subst -> VarSet -> VarSet
517 substVarSet subst fvs 
518   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
519   where
520     subst_fv subst fv 
521         | isId fv   = exprFreeVars (lookupIdSubst subst fv)
522         | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)
523 \end{code}