ca0fbd5a52a6e66e697eacda932b16fc1a2ff3c1
[ghc.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(..), -- Implementation exported for supercompiler's Renaming.hs only
12         TvSubstEnv, IdSubstEnv, InScopeSet,
13
14         -- ** Substituting into expressions and related types
15         deShadowBinds, substSpec, substRulesForImportedIds,
16         substTy, substCo, substExpr, substExprSC, substBind, substBindSC,
17         substUnfolding, substUnfoldingSC,
18         substUnfoldingSource, lookupIdSubst, lookupTvSubst, lookupCvSubst, substIdOcc,
19
20         -- ** Operations on substitutions
21         emptySubst, mkEmptySubst, mkSubst, mkOpenSubst, substInScope, isEmptySubst, 
22         extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
23         extendCvSubst, extendCvSubstList,
24         extendSubst, extendSubstList, extendSubstWithVar, zapSubstEnv,
25         addInScopeSet, extendInScope, extendInScopeList, extendInScopeIds,
26         isInScope, setInScope,
27         delBndr, delBndrs,
28
29         -- ** Substituting and cloning binders
30         substBndr, substBndrs, substRecBndrs,
31         cloneBndrs, cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs,
32
33         -- ** Simple expression optimiser
34         simpleOptPgm, simpleOptExpr, simpleOptExprWith
35     ) where
36
37 #include "HsVersions.h"
38
39 import CoreSyn
40 import CoreFVs
41 import CoreUtils
42 import OccurAnal( occurAnalyseExpr, occurAnalysePgm )
43
44 import qualified Type
45 import qualified Coercion
46
47         -- We are defining local versions
48 import Type     hiding ( substTy, extendTvSubst, extendTvSubstList
49                        , isInScope, substTyVarBndr, cloneTyVarBndr )
50 import Coercion hiding ( substTy, substCo, extendTvSubst, substTyVarBndr, substCoVarBndr )
51
52 import OptCoercion ( optCoercion )
53 import PprCore     ( pprCoreBindings, pprRules )
54 import Module      ( Module )
55 import VarSet
56 import VarEnv
57 import Id
58 import Name     ( Name )
59 import Var
60 import IdInfo
61 import Unique
62 import UniqSupply
63 import Maybes
64 import ErrUtils
65 import DynFlags   ( DynFlags, DynFlag(..) )
66 import BasicTypes ( isAlwaysActive )
67 import Outputable
68 import PprCore          ()              -- Instances
69 import FastString
70
71 import Data.List
72 \end{code}
73
74
75 %************************************************************************
76 %*                                                                      *
77 \subsection{Substitutions}
78 %*                                                                      *
79 %************************************************************************
80
81 \begin{code}
82 -- | A substitution environment, containing both 'Id' and 'TyVar' substitutions.
83 --
84 -- Some invariants apply to how you use the substitution:
85 --
86 -- 1. #in_scope_invariant# The in-scope set contains at least those 'Id's and 'TyVar's that will be in scope /after/
87 -- applying the substitution to a term. Precisely, the in-scope set must be a superset of the free vars of the
88 -- substitution range that might possibly clash with locally-bound variables in the thing being substituted in.
89 --
90 -- 2. #apply_once# You may apply the substitution only /once/
91 --
92 -- There are various ways of setting up the in-scope set such that the first of these invariants hold:
93 --
94 -- * Arrange that the in-scope set really is all the things in scope
95 --
96 -- * Arrange that it's the free vars of the range of the substitution
97 --
98 -- * Make it empty, if you know that all the free vars of the substitution are fresh, and hence can't possibly clash
99 data Subst 
100   = Subst InScopeSet  -- Variables in in scope (both Ids and TyVars) /after/
101                       -- applying the substitution
102           IdSubstEnv  -- Substitution for Ids
103           TvSubstEnv  -- Substitution from TyVars to Types
104           CvSubstEnv  -- Substitution from CoVars to Coercions
105
106         -- INVARIANT 1: See #in_scope_invariant#
107         -- This is what lets us deal with name capture properly
108         -- It's a hard invariant to check...
109         --
110         -- INVARIANT 2: The substitution is apply-once; see Note [Apply once] with
111         --              Types.TvSubstEnv
112         --
113         -- INVARIANT 3: See Note [Extending the Subst]
114 \end{code}
115
116 Note [Extending the Subst]
117 ~~~~~~~~~~~~~~~~~~~~~~~~~~
118 For a core Subst, which binds Ids as well, we make a different choice for Ids
119 than we do for TyVars.  
120
121 For TyVars, see Note [Extending the TvSubst] with Type.TvSubstEnv
122
123 For Ids, we have a different invariant
124         The IdSubstEnv is extended *only* when the Unique on an Id changes
125         Otherwise, we just extend the InScopeSet
126
127 In consequence:
128
129 * If the TvSubstEnv and IdSubstEnv are both empty, substExpr would be a
130   no-op, so substExprSC ("short cut") does nothing.
131
132   However, substExpr still goes ahead and substitutes.  Reason: we may
133   want to replace existing Ids with new ones from the in-scope set, to
134   avoid space leaks.
135
136 * In substIdBndr, we extend the IdSubstEnv only when the unique changes
137
138 * If the CvSubstEnv, TvSubstEnv and IdSubstEnv are all empty,
139   substExpr does nothing (Note that the above rule for substIdBndr
140   maintains this property.  If the incoming envts are both empty, then
141   substituting the type and IdInfo can't change anything.)
142
143 * In lookupIdSubst, we *must* look up the Id in the in-scope set, because
144   it may contain non-trivial changes.  Example:
145         (/\a. \x:a. ...x...) Int
146   We extend the TvSubstEnv with [a |-> Int]; but x's unique does not change
147   so we only extend the in-scope set.  Then we must look up in the in-scope
148   set when we find the occurrence of x.
149
150 * The requirement to look up the Id in the in-scope set means that we
151   must NOT take no-op short cut when the IdSubst is empty.
152   We must still look up every Id in the in-scope set.
153
154 * (However, we don't need to do so for expressions found in the IdSubst
155   itself, whose range is assumed to be correct wrt the in-scope set.)
156
157 Why do we make a different choice for the IdSubstEnv than the
158 TvSubstEnv and CvSubstEnv?
159
160 * For Ids, we change the IdInfo all the time (e.g. deleting the
161   unfolding), and adding it back later, so using the TyVar convention
162   would entail extending the substitution almost all the time
163
164 * The simplifier wants to look up in the in-scope set anyway, in case it 
165   can see a better unfolding from an enclosing case expression
166
167 * For TyVars, only coercion variables can possibly change, and they are 
168   easy to spot
169
170 \begin{code}
171 -- | An environment for substituting for 'Id's
172 type IdSubstEnv = IdEnv CoreExpr
173
174 ----------------------------
175 isEmptySubst :: Subst -> Bool
176 isEmptySubst (Subst _ id_env tv_env cv_env) 
177   = isEmptyVarEnv id_env && isEmptyVarEnv tv_env && isEmptyVarEnv cv_env
178
179 emptySubst :: Subst
180 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv emptyVarEnv
181
182 mkEmptySubst :: InScopeSet -> Subst
183 mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv emptyVarEnv
184
185 mkSubst :: InScopeSet -> TvSubstEnv -> CvSubstEnv -> IdSubstEnv -> Subst
186 mkSubst in_scope tvs cvs ids = Subst in_scope ids tvs cvs
187
188 -- | Find the in-scope set: see "CoreSubst#in_scope_invariant"
189 substInScope :: Subst -> InScopeSet
190 substInScope (Subst in_scope _ _ _) = in_scope
191
192 -- | Remove all substitutions for 'Id's and 'Var's that might have been built up
193 -- while preserving the in-scope set
194 zapSubstEnv :: Subst -> Subst
195 zapSubstEnv (Subst in_scope _ _ _) = Subst in_scope emptyVarEnv emptyVarEnv emptyVarEnv
196
197 -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the in-scope set is
198 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
199 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
200 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
201 extendIdSubst (Subst in_scope ids tvs cvs) v r = Subst in_scope (extendVarEnv ids v r) tvs cvs
202
203 -- | Adds multiple 'Id' substitutions to the 'Subst': see also 'extendIdSubst'
204 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
205 extendIdSubstList (Subst in_scope ids tvs cvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs cvs
206
207 -- | Add a substitution for a 'TyVar' to the 'Subst': you must ensure that the in-scope set is
208 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
209 extendTvSubst :: Subst -> TyVar -> Type -> Subst
210 extendTvSubst (Subst in_scope ids tvs cvs) v r = Subst in_scope ids (extendVarEnv tvs v r) cvs
211
212 -- | Adds multiple 'TyVar' substitutions to the 'Subst': see also 'extendTvSubst'
213 extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
214 extendTvSubstList (Subst in_scope ids tvs cvs) prs = Subst in_scope ids (extendVarEnvList tvs prs) cvs
215
216 -- | Add a substitution from a 'CoVar' to a 'Coercion' to the 'Subst': you must ensure that the in-scope set is
217 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
218 extendCvSubst :: Subst -> CoVar -> Coercion -> Subst
219 extendCvSubst (Subst in_scope ids tvs cvs) v r = Subst in_scope ids tvs (extendVarEnv cvs v r)
220
221 -- | Adds multiple 'CoVar' -> 'Coercion' substitutions to the
222 -- 'Subst': see also 'extendCvSubst'
223 extendCvSubstList :: Subst -> [(CoVar,Coercion)] -> Subst
224 extendCvSubstList (Subst in_scope ids tvs cvs) prs = Subst in_scope ids tvs (extendVarEnvList cvs prs)
225
226 -- | Add a substitution appropriate to the thing being substituted
227 --   (whether an expression, type, or coercion). See also
228 --   'extendIdSubst', 'extendTvSubst', and 'extendCvSubst'.
229 extendSubst :: Subst -> Var -> CoreArg -> Subst
230 extendSubst subst var arg
231   = case arg of
232       Type ty     -> ASSERT( isTyVar var ) extendTvSubst subst var ty
233       Coercion co -> ASSERT( isCoVar var ) extendCvSubst subst var co
234       _           -> ASSERT( isId    var ) extendIdSubst subst var arg
235
236 extendSubstWithVar :: Subst -> Var -> Var -> Subst
237 extendSubstWithVar subst v1 v2
238   | isTyVar v1 = ASSERT( isTyVar v2 ) extendTvSubst subst v1 (mkTyVarTy v2)
239   | isCoVar v1 = ASSERT( isCoVar v2 ) extendCvSubst subst v1 (mkCoVarCo v2)
240   | otherwise  = ASSERT( isId    v2 ) extendIdSubst subst v1 (Var v2)
241
242 -- | Add a substitution as appropriate to each of the terms being
243 --   substituted (whether expressions, types, or coercions). See also
244 --   'extendSubst'.
245 extendSubstList :: Subst -> [(Var,CoreArg)] -> Subst
246 extendSubstList subst []              = subst
247 extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var rhs) prs
248
249 -- | Find the substitution for an 'Id' in the 'Subst'
250 lookupIdSubst :: SDoc -> Subst -> Id -> CoreExpr
251 lookupIdSubst doc (Subst in_scope ids _ _) v
252   | not (isLocalId v) = Var v
253   | Just e  <- lookupVarEnv ids       v = e
254   | Just v' <- lookupInScope in_scope v = Var v'
255         -- Vital! See Note [Extending the Subst]
256   | otherwise = WARN( True, ptext (sLit "CoreSubst.lookupIdSubst") <+> doc <+> ppr v 
257                             $$ ppr in_scope) 
258                 Var v
259
260 -- | Find the substitution for a 'TyVar' in the 'Subst'
261 lookupTvSubst :: Subst -> TyVar -> Type
262 lookupTvSubst (Subst _ _ tvs _) v = ASSERT( isTyVar v) lookupVarEnv tvs v `orElse` Type.mkTyVarTy v
263
264 -- | Find the coercion substitution for a 'CoVar' in the 'Subst'
265 lookupCvSubst :: Subst -> CoVar -> Coercion
266 lookupCvSubst (Subst _ _ _ cvs) v = ASSERT( isCoVar v ) lookupVarEnv cvs v `orElse` mkCoVarCo v
267
268 delBndr :: Subst -> Var -> Subst
269 delBndr (Subst in_scope ids tvs cvs) v
270   | isCoVar v = Subst in_scope ids tvs (delVarEnv cvs v)
271   | isTyVar v = Subst in_scope ids (delVarEnv tvs v) cvs
272   | otherwise = Subst in_scope (delVarEnv ids v) tvs cvs
273
274 delBndrs :: Subst -> [Var] -> Subst
275 delBndrs (Subst in_scope ids tvs cvs) vs
276   = Subst in_scope (delVarEnvList ids vs) (delVarEnvList tvs vs) (delVarEnvList cvs vs)
277       -- Easist thing is just delete all from all!
278
279 -- | Simultaneously substitute for a bunch of variables
280 --   No left-right shadowing
281 --   ie the substitution for   (\x \y. e) a1 a2
282 --      so neither x nor y scope over a1 a2
283 mkOpenSubst :: InScopeSet -> [(Var,CoreArg)] -> Subst
284 mkOpenSubst in_scope pairs = Subst in_scope
285                                    (mkVarEnv [(id,e)  | (id, e) <- pairs, isId id])
286                                    (mkVarEnv [(tv,ty) | (tv, Type ty) <- pairs])
287                                    (mkVarEnv [(v,co)  | (v, Coercion co) <- pairs])
288
289 ------------------------------
290 isInScope :: Var -> Subst -> Bool
291 isInScope v (Subst in_scope _ _ _) = v `elemInScopeSet` in_scope
292
293 -- | Add the 'Var' to the in-scope set, but do not remove
294 -- any existing substitutions for it
295 addInScopeSet :: Subst -> VarSet -> Subst
296 addInScopeSet (Subst in_scope ids tvs cvs) vs
297   = Subst (in_scope `extendInScopeSetSet` vs) ids tvs cvs
298
299 -- | Add the 'Var' to the in-scope set: as a side effect,
300 -- and remove any existing substitutions for it
301 extendInScope :: Subst -> Var -> Subst
302 extendInScope (Subst in_scope ids tvs cvs) v
303   = Subst (in_scope `extendInScopeSet` v) 
304           (ids `delVarEnv` v) (tvs `delVarEnv` v) (cvs `delVarEnv` v)
305
306 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
307 extendInScopeList :: Subst -> [Var] -> Subst
308 extendInScopeList (Subst in_scope ids tvs cvs) vs
309   = Subst (in_scope `extendInScopeSetList` vs) 
310           (ids `delVarEnvList` vs) (tvs `delVarEnvList` vs) (cvs `delVarEnvList` vs)
311
312 -- | Optimized version of 'extendInScopeList' that can be used if you are certain 
313 -- all the things being added are 'Id's and hence none are 'TyVar's or 'CoVar's
314 extendInScopeIds :: Subst -> [Id] -> Subst
315 extendInScopeIds (Subst in_scope ids tvs cvs) vs 
316   = Subst (in_scope `extendInScopeSetList` vs) 
317           (ids `delVarEnvList` vs) tvs cvs
318
319 setInScope :: Subst -> InScopeSet -> Subst
320 setInScope (Subst _ ids tvs cvs) in_scope = Subst in_scope ids tvs cvs
321 \end{code}
322
323 Pretty printing, for debugging only
324
325 \begin{code}
326 instance Outputable Subst where
327   ppr (Subst in_scope ids tvs cvs) 
328         =  ptext (sLit "<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
329         $$ ptext (sLit " IdSubst   =") <+> ppr ids
330         $$ ptext (sLit " TvSubst   =") <+> ppr tvs
331         $$ ptext (sLit " CvSubst   =") <+> ppr cvs   
332          <> char '>'
333 \end{code}
334
335
336 %************************************************************************
337 %*                                                                      *
338         Substituting expressions
339 %*                                                                      *
340 %************************************************************************
341
342 \begin{code}
343 -- | Apply a substititon to an entire 'CoreExpr'. Rememeber, you may only 
344 -- apply the substitution /once/: see "CoreSubst#apply_once"
345 --
346 -- Do *not* attempt to short-cut in the case of an empty substitution!
347 -- See Note [Extending the Subst]
348 substExprSC :: SDoc -> Subst -> CoreExpr -> CoreExpr
349 substExprSC _doc subst orig_expr
350   | isEmptySubst subst = orig_expr
351   | otherwise          = -- pprTrace "enter subst-expr" (doc $$ ppr orig_expr) $
352                          subst_expr subst orig_expr
353
354 substExpr :: SDoc -> Subst -> CoreExpr -> CoreExpr
355 substExpr _doc subst orig_expr = subst_expr subst orig_expr
356
357 subst_expr :: Subst -> CoreExpr -> CoreExpr
358 subst_expr subst expr
359   = go expr
360   where
361     go (Var v)         = lookupIdSubst (text "subst_expr") subst v 
362     go (Type ty)       = Type (substTy subst ty)
363     go (Coercion co)   = Coercion (substCo subst co)
364     go (Lit lit)       = Lit lit
365     go (App fun arg)   = App (go fun) (go arg)
366     go (Note note e)   = Note (go_note note) (go e)
367     go (Cast e co)     = Cast (go e) (substCo subst co)
368        -- Do not optimise even identity coercions
369        -- Reason: substitution applies to the LHS of RULES, and
370        --         if you "optimise" an identity coercion, you may
371        --         lose a binder. We optimise the LHS of rules at
372        --         construction time
373
374     go (Lam bndr body) = Lam bndr' (subst_expr subst' body)
375                        where
376                          (subst', bndr') = substBndr subst bndr
377
378     go (Let bind body) = Let bind' (subst_expr subst' body)
379                        where
380                          (subst', bind') = substBind subst bind
381
382     go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
383                                  where
384                                  (subst', bndr') = substBndr subst bndr
385
386     go_alt subst (con, bndrs, rhs) = (con, bndrs', subst_expr subst' rhs)
387                                  where
388                                    (subst', bndrs') = substBndrs subst bndrs
389
390     go_note note             = note
391
392 -- | Apply a substititon to an entire 'CoreBind', additionally returning an updated 'Subst'
393 -- that should be used by subsequent substitutons.
394 substBind, substBindSC :: Subst -> CoreBind -> (Subst, CoreBind)
395
396 substBindSC subst bind    -- Short-cut if the substitution is empty
397   | not (isEmptySubst subst)
398   = substBind subst bind
399   | otherwise
400   = case bind of
401        NonRec bndr rhs -> (subst', NonRec bndr' rhs)
402           where
403             (subst', bndr') = substBndr subst bndr
404        Rec pairs -> (subst', Rec (bndrs' `zip` rhss'))
405           where
406             (bndrs, rhss)    = unzip pairs
407             (subst', bndrs') = substRecBndrs subst bndrs
408             rhss' | isEmptySubst subst' = rhss
409                   | otherwise           = map (subst_expr subst') rhss
410
411 substBind subst (NonRec bndr rhs) = (subst', NonRec bndr' (subst_expr subst rhs))
412                                   where
413                                     (subst', bndr') = substBndr subst bndr
414
415 substBind subst (Rec pairs) = (subst', Rec (bndrs' `zip` rhss'))
416                             where
417                                 (bndrs, rhss)    = unzip pairs
418                                 (subst', bndrs') = substRecBndrs subst bndrs
419                                 rhss' = map (subst_expr subst') rhss
420 \end{code}
421
422 \begin{code}
423 -- | De-shadowing the program is sometimes a useful pre-pass. It can be done simply
424 -- by running over the bindings with an empty substitution, becuase substitution
425 -- returns a result that has no-shadowing guaranteed.
426 --
427 -- (Actually, within a single /type/ there might still be shadowing, because 
428 -- 'substTy' is a no-op for the empty substitution, but that's probably OK.)
429 --
430 -- [Aug 09] This function is not used in GHC at the moment, but seems so 
431 --          short and simple that I'm going to leave it here
432 deShadowBinds :: [CoreBind] -> [CoreBind]
433 deShadowBinds binds = snd (mapAccumL substBind emptySubst binds)
434 \end{code}
435
436
437 %************************************************************************
438 %*                                                                      *
439         Substituting binders
440 %*                                                                      *
441 %************************************************************************
442
443 Remember that substBndr and friends are used when doing expression
444 substitution only.  Their only business is substitution, so they
445 preserve all IdInfo (suitably substituted).  For example, we *want* to
446 preserve occ info in rules.
447
448 \begin{code}
449 -- | Substitutes a 'Var' for another one according to the 'Subst' given, returning
450 -- the result and an updated 'Subst' that should be used by subsequent substitutons.
451 -- 'IdInfo' is preserved by this process, although it is substituted into appropriately.
452 substBndr :: Subst -> Var -> (Subst, Var)
453 substBndr subst bndr
454   | isTyVar bndr  = substTyVarBndr subst bndr
455   | isCoVar bndr  = substCoVarBndr subst bndr
456   | otherwise     = substIdBndr (text "var-bndr") subst subst bndr
457
458 -- | Applies 'substBndr' to a number of 'Var's, accumulating a new 'Subst' left-to-right
459 substBndrs :: Subst -> [Var] -> (Subst, [Var])
460 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
461
462 -- | Substitute in a mutually recursive group of 'Id's
463 substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
464 substRecBndrs subst bndrs 
465   = (new_subst, new_bndrs)
466   where         -- Here's the reason we need to pass rec_subst to subst_id
467     (new_subst, new_bndrs) = mapAccumL (substIdBndr (text "rec-bndr") new_subst) subst bndrs
468 \end{code}
469
470
471 \begin{code}
472 substIdBndr :: SDoc 
473             -> Subst            -- ^ Substitution to use for the IdInfo
474             -> Subst -> Id      -- ^ Substitition and Id to transform
475             -> (Subst, Id)      -- ^ Transformed pair
476                                 -- NB: unfolding may be zapped
477
478 substIdBndr _doc rec_subst subst@(Subst in_scope env tvs cvs) old_id
479   = -- pprTrace "substIdBndr" (doc $$ ppr old_id $$ ppr in_scope) $
480     (Subst (in_scope `extendInScopeSet` new_id) new_env tvs cvs, new_id)
481   where
482     id1 = uniqAway in_scope old_id      -- id1 is cloned if necessary
483     id2 | no_type_change = id1
484         | otherwise      = setIdType id1 (substTy subst old_ty)
485
486     old_ty = idType old_id
487     no_type_change = isEmptyVarEnv tvs || 
488                      isEmptyVarSet (Type.tyVarsOfType old_ty)
489
490         -- new_id has the right IdInfo
491         -- The lazy-set is because we're in a loop here, with 
492         -- rec_subst, when dealing with a mutually-recursive group
493     new_id = maybeModifyIdInfo mb_new_info id2
494     mb_new_info = substIdInfo rec_subst id2 (idInfo id2)
495         -- NB: unfolding info may be zapped
496
497         -- Extend the substitution if the unique has changed
498         -- See the notes with substTyVarBndr for the delVarEnv
499     new_env | no_change = delVarEnv env old_id
500             | otherwise = extendVarEnv env old_id (Var new_id)
501
502     no_change = id1 == old_id
503         -- See Note [Extending the Subst]
504         -- it's /not/ necessary to check mb_new_info and no_type_change
505 \end{code}
506
507 Now a variant that unconditionally allocates a new unique.
508 It also unconditionally zaps the OccInfo.
509
510 \begin{code}
511 -- | Very similar to 'substBndr', but it always allocates a new 'Unique' for
512 -- each variable in its output.  It substitutes the IdInfo though.
513 cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
514 cloneIdBndr subst us old_id
515   = clone_id subst subst (old_id, uniqFromSupply us)
516
517 -- | Applies 'cloneIdBndr' to a number of 'Id's, accumulating a final
518 -- substitution from left to right
519 cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
520 cloneIdBndrs subst us ids
521   = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)
522
523 cloneBndrs :: Subst -> UniqSupply -> [Var] -> (Subst, [Var])
524 -- Works for all kinds of variables (typically case binders)
525 -- not just Ids
526 cloneBndrs subst us vs
527   = mapAccumL clone subst (vs `zip` uniqsFromSupply us)
528   where
529     clone subst (v,uniq) 
530       | isTyVar v = cloneTyVarBndr subst v uniq
531       | otherwise = clone_id subst subst (v,uniq)  -- Works for coercion variables too
532
533 -- | Clone a mutually recursive group of 'Id's
534 cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
535 cloneRecIdBndrs subst us ids
536   = (subst', ids')
537   where
538     (subst', ids') = mapAccumL (clone_id subst') subst
539                                (ids `zip` uniqsFromSupply us)
540
541 -- Just like substIdBndr, except that it always makes a new unique
542 -- It is given the unique to use
543 clone_id    :: Subst                    -- Substitution for the IdInfo
544             -> Subst -> (Id, Unique)    -- Substitition and Id to transform
545             -> (Subst, Id)              -- Transformed pair
546
547 clone_id rec_subst subst@(Subst in_scope env tvs cvs) (old_id, uniq)
548   = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs cvs, new_id)
549   where
550     id1     = setVarUnique old_id uniq
551     id2     = substIdType subst id1
552     new_id  = maybeModifyIdInfo (substIdInfo rec_subst id2 (idInfo old_id)) id2
553     new_env = extendVarEnv env old_id (Var new_id)
554 \end{code}
555
556
557 %************************************************************************
558 %*                                                                      *
559                 Types and Coercions
560 %*                                                                      *
561 %************************************************************************
562
563 For types and coercions we just call the corresponding functions in
564 Type and Coercion, but we have to repackage the substitution, from a
565 Subst to a TvSubst.
566
567 \begin{code}
568 substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
569 substTyVarBndr (Subst in_scope id_env tv_env cv_env) tv
570   = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
571         (TvSubst in_scope' tv_env', tv') 
572            -> (Subst in_scope' id_env tv_env' cv_env, tv')
573
574 cloneTyVarBndr :: Subst -> TyVar -> Unique -> (Subst, TyVar)
575 cloneTyVarBndr (Subst in_scope id_env tv_env cv_env) tv uniq
576   = case Type.cloneTyVarBndr (TvSubst in_scope tv_env) tv uniq of
577         (TvSubst in_scope' tv_env', tv') 
578            -> (Subst in_scope' id_env tv_env' cv_env, tv')
579
580 substCoVarBndr :: Subst -> TyVar -> (Subst, TyVar)
581 substCoVarBndr (Subst in_scope id_env tv_env cv_env) cv
582   = case Coercion.substCoVarBndr (CvSubst in_scope tv_env cv_env) cv of
583         (CvSubst in_scope' tv_env' cv_env', cv') 
584            -> (Subst in_scope' id_env tv_env' cv_env', cv')
585
586 -- | See 'Type.substTy'
587 substTy :: Subst -> Type -> Type 
588 substTy subst ty = Type.substTy (getTvSubst subst) ty
589
590 getTvSubst :: Subst -> TvSubst
591 getTvSubst (Subst in_scope _ tenv _) = TvSubst in_scope tenv
592
593 getCvSubst :: Subst -> CvSubst
594 getCvSubst (Subst in_scope _ tenv cenv) = CvSubst in_scope tenv cenv
595
596 -- | See 'Coercion.substCo'
597 substCo :: Subst -> Coercion -> Coercion
598 substCo subst co = Coercion.substCo (getCvSubst subst) co
599 \end{code}
600
601
602 %************************************************************************
603 %*                                                                      *
604 \section{IdInfo substitution}
605 %*                                                                      *
606 %************************************************************************
607
608 \begin{code}
609 substIdType :: Subst -> Id -> Id
610 substIdType subst@(Subst _ _ tv_env cv_env) id
611   | (isEmptyVarEnv tv_env && isEmptyVarEnv cv_env) || isEmptyVarSet (Type.tyVarsOfType old_ty) = id
612   | otherwise   = setIdType id (substTy subst old_ty)
613                 -- The tyVarsOfType is cheaper than it looks
614                 -- because we cache the free tyvars of the type
615                 -- in a Note in the id's type itself
616   where
617     old_ty = idType id
618
619 ------------------
620 -- | Substitute into some 'IdInfo' with regard to the supplied new 'Id'.
621 substIdInfo :: Subst -> Id -> IdInfo -> Maybe IdInfo
622 substIdInfo subst new_id info
623   | nothing_to_do = Nothing
624   | otherwise     = Just (info `setSpecInfo`      substSpec subst new_id old_rules
625                                `setUnfoldingInfo` substUnfolding subst old_unf)
626   where
627     old_rules     = specInfo info
628     old_unf       = unfoldingInfo info
629     nothing_to_do = isEmptySpecInfo old_rules && isClosedUnfolding old_unf
630     
631
632 ------------------
633 -- | Substitutes for the 'Id's within an unfolding
634 substUnfolding, substUnfoldingSC :: Subst -> Unfolding -> Unfolding
635         -- Seq'ing on the returned Unfolding is enough to cause
636         -- all the substitutions to happen completely
637
638 substUnfoldingSC subst unf       -- Short-cut version
639   | isEmptySubst subst = unf
640   | otherwise          = substUnfolding subst unf
641
642 substUnfolding subst (DFunUnfolding ar con args)
643   = DFunUnfolding ar con (map subst_arg args)
644   where
645     subst_arg = substExpr (text "dfun-unf") subst
646
647 substUnfolding subst unf@(CoreUnfolding { uf_tmpl = tmpl, uf_src = src })
648         -- Retain an InlineRule!
649   | not (isStableSource src)  -- Zap an unstable unfolding, to save substitution work
650   = NoUnfolding
651   | otherwise                 -- But keep a stable one!
652   = seqExpr new_tmpl `seq` 
653     new_src `seq`
654     unf { uf_tmpl = new_tmpl, uf_src = new_src }
655   where
656     new_tmpl = substExpr (text "subst-unf") subst tmpl
657     new_src  = substUnfoldingSource subst src
658
659 substUnfolding _ unf = unf      -- NoUnfolding, OtherCon
660
661 -------------------
662 substUnfoldingSource :: Subst -> UnfoldingSource -> UnfoldingSource
663 substUnfoldingSource (Subst in_scope ids _ _) (InlineWrapper wkr)
664   | Just wkr_expr <- lookupVarEnv ids wkr 
665   = case wkr_expr of
666       Var w1 -> InlineWrapper w1
667       _other -> -- WARN( True, text "Interesting! CoreSubst.substWorker1:" <+> ppr wkr 
668                 --             <+> ifPprDebug (equals <+> ppr wkr_expr) )   
669                               -- Note [Worker inlining]
670                 InlineStable  -- It's not a wrapper any more, but still inline it!
671
672   | Just w1  <- lookupInScope in_scope wkr = InlineWrapper w1
673   | otherwise = -- WARN( True, text "Interesting! CoreSubst.substWorker2:" <+> ppr wkr )
674                 -- This can legitimately happen.  The worker has been inlined and
675                 -- dropped as dead code, because we don't treat the UnfoldingSource
676                 -- as an "occurrence".
677                 -- Note [Worker inlining]
678                 InlineStable
679
680 substUnfoldingSource _ src = src
681
682 ------------------
683 substIdOcc :: Subst -> Id -> Id
684 -- These Ids should not be substituted to non-Ids
685 substIdOcc subst v = case lookupIdSubst (text "substIdOcc") subst v of
686                         Var v' -> v'
687                         other  -> pprPanic "substIdOcc" (vcat [ppr v <+> ppr other, ppr subst])
688
689 ------------------
690 -- | Substitutes for the 'Id's within the 'WorkerInfo' given the new function 'Id'
691 substSpec :: Subst -> Id -> SpecInfo -> SpecInfo
692 substSpec subst new_id (SpecInfo rules rhs_fvs)
693   = seqSpecInfo new_spec `seq` new_spec
694   where
695     subst_ru_fn = const (idName new_id)
696     new_spec = SpecInfo (map (substRule subst subst_ru_fn) rules)
697                         (substVarSet subst rhs_fvs)
698
699 ------------------
700 substRulesForImportedIds :: Subst -> [CoreRule] -> [CoreRule]
701 substRulesForImportedIds subst rules 
702   = map (substRule subst not_needed) rules
703   where
704     not_needed name = pprPanic "substRulesForImportedIds" (ppr name)
705
706 ------------------
707 substRule :: Subst -> (Name -> Name) -> CoreRule -> CoreRule
708
709 -- The subst_ru_fn argument is applied to substitute the ru_fn field
710 -- of the rule:
711 --    - Rules for *imported* Ids never change ru_fn
712 --    - Rules for *local* Ids are in the IdInfo for that Id,
713 --      and the ru_fn field is simply replaced by the new name 
714 --      of the Id
715 substRule _ _ rule@(BuiltinRule {}) = rule
716 substRule subst subst_ru_fn rule@(Rule { ru_bndrs = bndrs, ru_args = args
717                                        , ru_fn = fn_name, ru_rhs = rhs
718                                        , ru_local = is_local })
719   = rule { ru_bndrs = bndrs', 
720            ru_fn    = if is_local 
721                         then subst_ru_fn fn_name 
722                         else fn_name,
723            ru_args  = map (substExpr (text "subst-rule" <+> ppr fn_name) subst') args,
724            ru_rhs   = simpleOptExprWith subst' rhs }
725            -- Do simple optimisation on RHS, in case substitution lets
726            -- you improve it.  The real simplifier never gets to look at it.
727   where
728     (subst', bndrs') = substBndrs subst bndrs
729
730 ------------------
731 substVects :: Subst -> [CoreVect] -> [CoreVect]
732 substVects subst = map (substVect subst)
733
734 ------------------
735 substVect :: Subst -> CoreVect -> CoreVect
736 substVect _subst (Vect   v Nothing)    = Vect   v Nothing
737 substVect subst  (Vect   v (Just rhs)) = Vect   v (Just (simpleOptExprWith subst rhs))
738 substVect _subst (NoVect v)            = NoVect v
739
740 ------------------
741 substVarSet :: Subst -> VarSet -> VarSet
742 substVarSet subst fvs
743   = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
744   where
745     subst_fv subst fv 
746         | isId fv   = exprFreeVars (lookupIdSubst (text "substVarSet") subst fv)
747         | otherwise = Type.tyVarsOfType (lookupTvSubst subst fv)
748 \end{code}
749
750 Note [Worker inlining]
751 ~~~~~~~~~~~~~~~~~~~~~~
752 A worker can get sustituted away entirely.
753         - it might be trivial
754         - it might simply be very small
755 We do not treat an InlWrapper as an 'occurrence' in the occurence 
756 analyser, so it's possible that the worker is not even in scope any more.
757
758 In all all these cases we simply drop the special case, returning to
759 InlVanilla.  The WARN is just so I can see if it happens a lot.
760
761
762 %************************************************************************
763 %*                                                                      *
764         The Very Simple Optimiser
765 %*                                                                      *
766 %************************************************************************
767
768 \begin{code}
769 simpleOptExpr :: CoreExpr -> CoreExpr
770 -- Do simple optimisation on an expression
771 -- The optimisation is very straightforward: just
772 -- inline non-recursive bindings that are used only once, 
773 -- or where the RHS is trivial
774 --
775 -- The result is NOT guaranteed occurence-analysed, becuase
776 -- in  (let x = y in ....) we substitute for x; so y's occ-info
777 -- may change radically
778
779 simpleOptExpr expr
780   = -- pprTrace "simpleOptExpr" (ppr init_subst $$ ppr expr)
781     simpleOptExprWith init_subst expr
782   where
783     init_subst = mkEmptySubst (mkInScopeSet (exprFreeVars expr))
784         -- It's potentially important to make a proper in-scope set
785         -- Consider  let x = ..y.. in \y. ...x...
786         -- Then we should remember to clone y before substituting
787         -- for x.  It's very unlikely to occur, because we probably
788         -- won't *be* substituting for x if it occurs inside a
789         -- lambda.  
790         --
791         -- It's a bit painful to call exprFreeVars, because it makes
792         -- three passes instead of two (occ-anal, and go)
793
794 simpleOptExprWith :: Subst -> InExpr -> OutExpr
795 simpleOptExprWith subst expr = simple_opt_expr subst (occurAnalyseExpr expr)
796
797 ----------------------
798 simpleOptPgm :: DynFlags -> Module 
799              -> [CoreBind] -> [CoreRule] -> [CoreVect] 
800              -> IO ([CoreBind], [CoreRule], [CoreVect])
801 simpleOptPgm dflags this_mod binds rules vects
802   = do { dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
803                        (pprCoreBindings occ_anald_binds $$ pprRules rules );
804
805        ; return (reverse binds', substRulesForImportedIds subst' rules, substVects subst' vects) }
806   where
807     occ_anald_binds  = occurAnalysePgm this_mod (\_ -> False) {- No rules active -}
808                                        rules vects binds
809     (subst', binds') = foldl do_one (emptySubst, []) occ_anald_binds
810                        
811     do_one (subst, binds') bind 
812       = case simple_opt_bind subst bind of
813           (subst', Nothing)    -> (subst', binds')
814           (subst', Just bind') -> (subst', bind':binds')
815
816 ----------------------
817 type InVar   = Var
818 type OutVar  = Var
819 type InId    = Id
820 type OutId   = Id
821 type InExpr  = CoreExpr
822 type OutExpr = CoreExpr
823
824 -- In these functions the substitution maps InVar -> OutExpr
825
826 ----------------------
827 simple_opt_expr, simple_opt_expr' :: Subst -> InExpr -> OutExpr
828 simple_opt_expr s e = simple_opt_expr' s e
829
830 simple_opt_expr' subst expr
831   = go expr
832   where
833     go (Var v)          = lookupIdSubst (text "simpleOptExpr") subst v
834     go (App e1 e2)      = simple_app subst e1 [go e2]
835     go (Type ty)        = Type     (substTy subst ty)
836     go (Coercion co)    = Coercion (optCoercion (getCvSubst subst) co)
837     go (Lit lit)        = Lit lit
838     go (Note note e)    = Note note (go e)
839     go (Cast e co)      | isReflCo co' = go e
840                         | otherwise    = Cast (go e) co' 
841                         where
842                           co' = optCoercion (getCvSubst subst) co
843
844     go (Let bind body) = case simple_opt_bind subst bind of
845                            (subst', Nothing)   -> simple_opt_expr subst' body
846                            (subst', Just bind) -> Let bind (simple_opt_expr subst' body)
847
848     go lam@(Lam {})     = go_lam [] subst lam
849     go (Case e b ty as) = Case (go e) b' (substTy subst ty)
850                                (map (go_alt subst') as)
851                         where
852                           (subst', b') = subst_opt_bndr subst b
853
854     ----------------------
855     go_alt subst (con, bndrs, rhs) 
856       = (con, bndrs', simple_opt_expr subst' rhs)
857       where
858         (subst', bndrs') = subst_opt_bndrs subst bndrs
859
860     ----------------------
861     -- go_lam tries eta reduction
862     go_lam bs' subst (Lam b e) 
863        = go_lam (b':bs') subst' e
864        where
865          (subst', b') = subst_opt_bndr subst b
866     go_lam bs' subst e 
867        | Just etad_e <- tryEtaReduce bs e' = etad_e
868        | otherwise                         = mkLams bs e'
869        where
870          bs = reverse bs'
871          e' = simple_opt_expr subst e
872
873 ----------------------
874 -- simple_app collects arguments for beta reduction
875 simple_app :: Subst -> InExpr -> [OutExpr] -> CoreExpr
876 simple_app subst (App e1 e2) as   
877   = simple_app subst e1 (simple_opt_expr subst e2 : as)
878 simple_app subst (Lam b e) (a:as) 
879   = case maybe_substitute subst b a of
880       Just ext_subst -> simple_app ext_subst e as
881       Nothing        -> Let (NonRec b2 a) (simple_app subst' e as)
882   where
883     (subst', b') = subst_opt_bndr subst b
884     b2 = add_info subst' b b'
885 simple_app subst e as
886   = foldl App (simple_opt_expr subst e) as
887
888 ----------------------
889 simple_opt_bind,simple_opt_bind' :: Subst -> CoreBind -> (Subst, Maybe CoreBind)
890 simple_opt_bind s b               -- Can add trace stuff here
891   = simple_opt_bind' s b
892
893 simple_opt_bind' subst (Rec prs)
894   = (subst'', res_bind)
895   where
896     res_bind            = Just (Rec (reverse rev_prs'))
897     (subst', bndrs')    = subst_opt_bndrs subst (map fst prs)
898     (subst'', rev_prs') = foldl do_pr (subst', []) (prs `zip` bndrs')
899     do_pr (subst, prs) ((b,r), b') 
900        = case maybe_substitute subst b r2 of
901            Just subst' -> (subst', prs)
902            Nothing     -> (subst,  (b2,r2):prs)
903        where
904          b2 = add_info subst b b'
905          r2 = simple_opt_expr subst r
906
907 simple_opt_bind' subst (NonRec b r)
908   = case maybe_substitute subst b r' of
909       Just ext_subst -> (ext_subst, Nothing)
910       Nothing        -> (subst', Just (NonRec b2 r'))
911   where
912     r' = simple_opt_expr subst r
913     (subst', b') = subst_opt_bndr subst b
914     b2 = add_info subst' b b'
915
916 ----------------------
917 maybe_substitute :: Subst -> InVar -> OutExpr -> Maybe Subst
918     -- (maybe_substitute subst in_var out_rhs)  
919     --   either extends subst with (in_var -> out_rhs)
920     --   or     returns Nothing
921 maybe_substitute subst b r
922   | Type ty <- r        -- let a::* = TYPE ty in <body>
923   = ASSERT( isTyVar b )
924     Just (extendTvSubst subst b ty)
925
926   | Coercion co <- r
927   = ASSERT( isCoVar b )
928     Just (extendCvSubst subst b co)
929
930   | isId b              -- let x = e in <body>
931   , safe_to_inline (idOccInfo b) 
932   , isAlwaysActive (idInlineActivation b)       -- Note [Inline prag in simplOpt]
933   , not (isStableUnfolding (idUnfolding b))
934   , not (isExportedId b)
935   = Just (extendIdSubst subst b r)
936   
937   | otherwise
938   = Nothing
939   where
940         -- Unconditionally safe to inline
941     safe_to_inline :: OccInfo -> Bool
942     safe_to_inline (IAmALoopBreaker {})     = False
943     safe_to_inline IAmDead                  = True
944     safe_to_inline (OneOcc in_lam one_br _) = (not in_lam && one_br) || exprIsTrivial r
945     safe_to_inline NoOccInfo                = exprIsTrivial r
946
947 ----------------------
948 subst_opt_bndr :: Subst -> InVar -> (Subst, OutVar)
949 subst_opt_bndr subst bndr
950   | isTyVar bndr  = substTyVarBndr subst bndr
951   | isCoVar bndr  = substCoVarBndr subst bndr
952   | otherwise     = subst_opt_id_bndr subst bndr
953
954 subst_opt_id_bndr :: Subst -> InId -> (Subst, OutId)
955 -- Nuke all fragile IdInfo, unfolding, and RULES; 
956 --    it gets added back later by add_info
957 -- Rather like SimplEnv.substIdBndr
958 --
959 -- It's important to zap fragile OccInfo (which CoreSubst.substIdBndr 
960 -- carefully does not do) because simplOptExpr invalidates it
961
962 subst_opt_id_bndr subst@(Subst in_scope id_subst tv_subst cv_subst) old_id
963   = (Subst new_in_scope new_id_subst tv_subst cv_subst, new_id)
964   where
965     id1    = uniqAway in_scope old_id
966     id2    = setIdType id1 (substTy subst (idType old_id))
967     new_id = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
968                                         -- and fragile OccInfo
969     new_in_scope = in_scope `extendInScopeSet` new_id
970
971         -- Extend the substitution if the unique has changed,
972         -- or there's some useful occurrence information
973         -- See the notes with substTyVarBndr for the delSubstEnv
974     new_id_subst | new_id /= old_id
975                  = extendVarEnv id_subst old_id (Var new_id)
976                  | otherwise 
977                  = delVarEnv id_subst old_id
978
979 ----------------------
980 subst_opt_bndrs :: Subst -> [InVar] -> (Subst, [OutVar])
981 subst_opt_bndrs subst bndrs
982   = mapAccumL subst_opt_bndr subst bndrs
983
984 ----------------------
985 add_info :: Subst -> InVar -> OutVar -> OutVar
986 add_info subst old_bndr new_bndr
987  | isTyVar old_bndr = new_bndr
988  | otherwise        = maybeModifyIdInfo mb_new_info new_bndr
989  where
990    mb_new_info = substIdInfo subst new_bndr (idInfo old_bndr)
991 \end{code}
992
993 Note [Inline prag in simplOpt]
994 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
995 If there's an INLINE/NOINLINE pragma that restricts the phase in 
996 which the binder can be inlined, we don't inline here; after all,
997 we don't know what phase we're in.  Here's an example
998
999   foo :: Int -> Int -> Int
1000   {-# INLINE foo #-}
1001   foo m n = inner m
1002      where
1003        {-# INLINE [1] inner #-}
1004        inner m = m+n
1005
1006   bar :: Int -> Int
1007   bar n = foo n 1
1008
1009 When inlining 'foo' in 'bar' we want the let-binding for 'inner' 
1010 to remain visible until Phase 1
1011
1012