[project @ 2000-03-23 17:45:17 by simonpj]
[ghc-hetmet.git] / ghc / compiler / coreSyn / Subst.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 \section[CoreUtils]{Utility functions on @Core@ syntax}
5
6 \begin{code}
7 module Subst (
8         -- In-scope set
9         InScopeSet, emptyInScopeSet,
10         lookupInScope, setInScope, extendInScope, extendInScopes, isInScope, modifyInScope,
11
12         -- Substitution stuff
13         Subst, TyVarSubst, IdSubst,
14         emptySubst, mkSubst, substEnv, substInScope,
15         lookupSubst, lookupIdSubst, isEmptySubst, extendSubst, extendSubstList,
16         zapSubstEnv, setSubstEnv, 
17
18         bindSubst, unBindSubst, bindSubstList, unBindSubstList,
19
20         -- Binders
21         substBndr, substBndrs, substTyVar, substId, substIds,
22         substAndCloneId, substAndCloneIds,
23
24         -- Type stuff
25         mkTyVarSubst, mkTopTyVarSubst, 
26         substTy, substClasses, substTheta,
27
28         -- Expression stuff
29         substExpr, substIdInfo
30     ) where
31
32 #include "HsVersions.h"
33
34 import CoreSyn          ( Expr(..), Bind(..), Note(..), CoreExpr, CoreBndr,
35                           CoreRules(..), CoreRule(..), 
36                           emptyCoreRules, isEmptyCoreRules, seqRules
37                         )
38 import CoreFVs          ( exprFreeVars )
39 import TypeRep          ( Type(..), TyNote(..), 
40                         )  -- friend
41 import Type             ( ThetaType, PredType(..), ClassContext,
42                           tyVarsOfType, tyVarsOfTypes, mkAppTy
43                         )
44 import VarSet
45 import VarEnv
46 import Var              ( setVarUnique, isId )
47 import Id               ( idType, setIdType, idOccInfo, zapFragileIdInfo )
48 import Name             ( isLocallyDefined )
49 import IdInfo           ( IdInfo, isFragileOccInfo,
50                           specInfo, setSpecInfo, 
51                           WorkerInfo(..), workerExists, workerInfo, setWorkerInfo, WorkerInfo
52                         )
53 import BasicTypes       ( OccInfo(..) )
54 import UniqSupply       ( UniqSupply, uniqFromSupply, splitUniqSupply )
55 import Var              ( Var, Id, TyVar, isTyVar )
56 import Outputable
57 import PprCore          ()      -- Instances
58 import Util             ( mapAccumL, foldl2, seqList, ($!) )
59 \end{code}
60
61 %************************************************************************
62 %*                                                                      *
63 \subsection{Substitutions}
64 %*                                                                      *
65 %************************************************************************
66
67 \begin{code}
68 type InScopeSet = VarEnv Var
69
70 data Subst = Subst InScopeSet           -- In scope
71                    SubstEnv             -- Substitution itself
72         -- INVARIANT 1: The (domain of the) in-scope set is a superset
73         --              of the free vars of the range of the substitution
74         --              that might possibly clash with locally-bound variables
75         --              in the thing being substituted in.
76         -- This is what lets us deal with name capture properly
77         -- It's a hard invariant to check...
78         -- There are various ways of causing it to happen:
79         --      - arrange that the in-scope set really is all the things in scope
80         --      - arrange that it's the free vars of the range of the substitution
81         --      - make it empty because all the free vars of the subst are fresh,
82         --              and hence can't possibly clash.a
83         --
84         -- INVARIANT 2: No variable is both in scope and in the domain of the substitution
85         --              Equivalently, the substitution is idempotent
86         --
87
88 type IdSubst    = Subst
89 \end{code}
90
91 The general plan about the substitution and in-scope set for Ids is as follows
92
93 * substId always adds new_id to the in-scope set.
94   new_id has a correctly-substituted type, but all its fragile IdInfo has been zapped.
95   That is added back in later.  So new_id is the minimal thing it's 
96   correct to substitute.
97
98 * substId adds a binding (DoneVar new_id occ) to the substitution if 
99         EITHER the Id's unique has changed
100         OR     the Id has interesting occurrence information
101   So in effect you can only get to interesting occurrence information
102   by looking up the *old* Id; it's not really attached to the new id
103   at all.
104
105   Note, though that the substitution isn't necessarily extended
106   if the type changes.  Why not?  Because of the next point:
107
108 * We *always, always* finish by looking up in the in-scope set 
109   any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
110   Reason: so that we never finish up with a "old" Id in the result.  
111   An old Id might point to an old unfolding and so on... which gives a space leak.
112
113   [The DoneEx and DoneVar hits map to "new" stuff.]
114
115 * It follows that substExpr must not do a no-op if the substitution is empty.
116   substType is free to do so, however.
117
118 * When we come to a let-binding (say) we generate new IdInfo, including an
119   unfolding, attach it to the binder, and add this newly adorned binder to
120   the in-scope set.  So all subsequent occurrences of the binder will get mapped
121   to the full-adorned binder, which is also the one put in the binding site.
122
123 * The in-scope "set" usually maps x->x; we use it simply for its domain.
124   But sometimes we have two in-scope Ids that are synomyms, and should
125   map to the same target:  x->x, y->x.  Notably:
126         case y of x { ... }
127   That's why the "set" is actually a VarEnv Var
128
129 \begin{code}
130 emptyInScopeSet :: InScopeSet
131 emptyInScopeSet = emptyVarSet
132
133 add_in_scope :: InScopeSet -> Var -> InScopeSet
134 add_in_scope in_scope v = extendVarEnv in_scope v v
135 \end{code}
136
137
138
139 \begin{code}
140 isEmptySubst :: Subst -> Bool
141 isEmptySubst (Subst _ env) = isEmptySubstEnv env
142
143 emptySubst :: Subst
144 emptySubst = Subst emptyInScopeSet emptySubstEnv
145
146 mkSubst :: InScopeSet -> SubstEnv -> Subst
147 mkSubst in_scope env = Subst in_scope env
148
149 substEnv :: Subst -> SubstEnv
150 substEnv (Subst _ env) = env
151
152 substInScope :: Subst -> InScopeSet
153 substInScope (Subst in_scope _) = in_scope
154
155 zapSubstEnv :: Subst -> Subst
156 zapSubstEnv (Subst in_scope env) = Subst in_scope emptySubstEnv
157
158 extendSubst :: Subst -> Var -> SubstResult -> Subst
159 extendSubst (Subst in_scope env) v r = Subst in_scope (extendSubstEnv env v r)
160
161 extendSubstList :: Subst -> [Var] -> [SubstResult] -> Subst
162 extendSubstList (Subst in_scope env) v r = Subst in_scope (extendSubstEnvList env v r)
163
164 lookupSubst :: Subst -> Var -> Maybe SubstResult
165 lookupSubst (Subst _ env) v = lookupSubstEnv env v
166
167 lookupIdSubst :: Subst -> Id -> SubstResult
168 -- Does the lookup in the in-scope set too
169 lookupIdSubst (Subst in_scope env) v
170   = case lookupSubstEnv env v of
171         Just (DoneId v' occ) -> DoneId (lookupInScope in_scope v') occ
172         Just res             -> res
173         Nothing              -> DoneId v' (idOccInfo v')
174                                 -- We don't use DoneId for LoopBreakers, so the idOccInfo is
175                                 -- very important!  If isFragileOccInfo returned True for
176                                 -- loop breakers we could avoid this call, but at the expense
177                                 -- of adding more to the substitution, and building new Ids
178                                 -- in substId a bit more often than really necessary
179                              where
180                                     v' = lookupInScope in_scope v
181
182 lookupInScope :: InScopeSet -> Var -> Var
183 -- It's important to look for a fixed point
184 -- When we see (case x of y { I# v -> ... })
185 -- we add  [x -> y] to the in-scope set (Simplify.simplCaseBinder).
186 -- When we lookup up an occurrence of x, we map to y, but then
187 -- we want to look up y in case it has acquired more evaluation information by now.
188 lookupInScope in_scope v 
189   = case lookupVarEnv in_scope v of
190         Just v' | v == v'   -> v'       -- Reached a fixed point
191                 | otherwise -> lookupInScope in_scope v'
192         Nothing             -> v
193
194 isInScope :: Var -> Subst -> Bool
195 isInScope v (Subst in_scope _) = v `elemVarEnv` in_scope
196
197 extendInScope :: Subst -> Var -> Subst
198 extendInScope (Subst in_scope env) v = Subst (in_scope `add_in_scope` v) env
199
200 modifyInScope :: Subst -> Var -> Var -> Subst
201 modifyInScope (Subst in_scope env) old_v new_v = Subst (extendVarEnv in_scope old_v new_v) env
202         -- make old_v map to new_v
203
204 extendInScopes :: Subst -> [Var] -> Subst
205 extendInScopes (Subst in_scope env) vs = Subst (foldl add_in_scope in_scope vs) env
206
207 -------------------------------
208 bindSubst :: Subst -> Var -> Var -> Subst
209 -- Extend with a substitution, v1 -> Var v2
210 -- and extend the in-scopes with v2
211 bindSubst (Subst in_scope env) old_bndr new_bndr
212   = Subst (in_scope `add_in_scope` new_bndr)
213           (extendSubstEnv env old_bndr subst_result)
214   where
215     subst_result | isId old_bndr = DoneEx (Var new_bndr)
216                  | otherwise     = DoneTy (TyVarTy new_bndr)
217
218 unBindSubst :: Subst -> Var -> Var -> Subst
219 -- Reverse the effect of bindSubst
220 -- If old_bndr was already in the substitution, this doesn't quite work
221 unBindSubst (Subst in_scope env) old_bndr new_bndr
222   = Subst (in_scope `delVarEnv` new_bndr) (delSubstEnv env old_bndr)
223
224 -- And the "List" forms
225 bindSubstList :: Subst -> [Var] -> [Var] -> Subst
226 bindSubstList subst old_bndrs new_bndrs
227   = foldl2 bindSubst subst old_bndrs new_bndrs
228
229 unBindSubstList :: Subst -> [Var] -> [Var] -> Subst
230 unBindSubstList subst old_bndrs new_bndrs
231   = foldl2 unBindSubst subst old_bndrs new_bndrs
232
233
234 -------------------------------
235 setInScope :: Subst     -- Take env part from here
236            -> InScopeSet
237            -> Subst
238 setInScope (Subst in_scope1 env1) in_scope2
239   = Subst in_scope2 env1
240
241 setSubstEnv :: Subst            -- Take in-scope part from here
242             -> SubstEnv         -- ... and env part from here
243             -> Subst
244 setSubstEnv (Subst in_scope1 _) env2 = Subst in_scope1 env2
245 \end{code}
246
247
248 %************************************************************************
249 %*                                                                      *
250 \subsection{Type substitution}
251 %*                                                                      *
252 %************************************************************************
253
254 \begin{code}
255 type TyVarSubst    = Subst      -- TyVarSubst are expected to have range elements
256         -- (We could have a variant of Subst, but it doesn't seem worth it.)
257
258 -- mkTyVarSubst generates the in-scope set from
259 -- the types given; but it's just a thunk so with a bit of luck
260 -- it'll never be evaluated
261 mkTyVarSubst :: [TyVar] -> [Type] -> Subst
262 mkTyVarSubst tyvars tys = Subst (tyVarsOfTypes tys) (zip_ty_env tyvars tys emptySubstEnv)
263
264 -- mkTopTyVarSubst is called when doing top-level substitutions.
265 -- Here we expect that the free vars of the range of the
266 -- substitution will be empty.
267 mkTopTyVarSubst :: [TyVar] -> [Type] -> Subst
268 mkTopTyVarSubst tyvars tys = Subst emptyInScopeSet (zip_ty_env tyvars tys emptySubstEnv)
269
270 zip_ty_env []       []       env = env
271 zip_ty_env (tv:tvs) (ty:tys) env = zip_ty_env tvs tys (extendSubstEnv env tv (DoneTy ty))
272 \end{code}
273
274 substTy works with general Substs, so that it can be called from substExpr too.
275
276 \begin{code}
277 substTy :: Subst -> Type  -> Type
278 substTy subst ty | isEmptySubst subst = ty
279                  | otherwise          = subst_ty subst ty
280
281 substClasses :: TyVarSubst -> ClassContext -> ClassContext
282 substClasses subst theta
283   | isEmptySubst subst = theta
284   | otherwise          = [(clas, map (subst_ty subst) tys) | (clas, tys) <- theta]
285
286 substTheta :: TyVarSubst -> ThetaType -> ThetaType
287 substTheta subst theta
288   | isEmptySubst subst = theta
289   | otherwise          = map (substPred subst) theta
290
291 substPred :: TyVarSubst -> PredType -> PredType
292 substPred subst (Class clas tys) = Class clas (map (subst_ty subst) tys)
293 substPred subst (IParam n ty)    = IParam n (subst_ty subst ty)
294
295 subst_ty subst ty
296    = go ty
297   where
298     go (TyConApp tc tys)          = let args = map go tys
299                                     in  args `seqList` TyConApp tc args
300     go (NoteTy (SynNote ty1) ty2) = NoteTy (SynNote $! (go ty1)) $! (go ty2)
301     go (NoteTy (FTVNote _) ty2)   = go ty2              -- Discard the free tyvar note
302     go (FunTy arg res)            = (FunTy $! (go arg)) $! (go res)
303     go (NoteTy (UsgNote usg)  ty2) = (NoteTy $! UsgNote usg) $! go ty2          -- Keep usage annot
304     go (NoteTy (UsgForAll uv) ty2) = (NoteTy $! UsgForAll uv) $! go ty2         -- Keep uvar bdr
305     go (NoteTy (IPNote nm) ty2)    = (NoteTy $! IPNote nm) $! go ty2            -- Keep ip note
306     go (AppTy fun arg)            = mkAppTy (go fun) $! (go arg)
307     go ty@(TyVarTy tv)            = case (lookupSubst subst tv) of
308                                         Nothing            -> ty
309                                         Just (DoneTy ty')  -> ty'
310                                         
311     go (ForAllTy tv ty)           = case substTyVar subst tv of
312                                         (subst', tv') -> ForAllTy tv' $! (subst_ty subst' ty)
313 \end{code}
314
315 Here is where we invent a new binder if necessary.
316
317 \begin{code}
318 substTyVar :: Subst -> TyVar -> (Subst, TyVar)  
319 substTyVar subst@(Subst in_scope env) old_var
320   | old_var == new_var  -- No need to clone
321                         -- But we *must* zap any current substitution for the variable.
322                         --  For example:
323                         --      (\x.e) with id_subst = [x |-> e']
324                         -- Here we must simply zap the substitution for x
325                         --
326                         -- The new_id isn't cloned, but it may have a different type
327                         -- etc, so we must return it, not the old id
328   = (Subst (in_scope `add_in_scope` new_var)
329            (delSubstEnv env old_var),
330      new_var)
331
332   | otherwise   -- The new binder is in scope so
333                 -- we'd better rename it away from the in-scope variables
334                 -- Extending the substitution to do this renaming also
335                 -- has the (correct) effect of discarding any existing
336                 -- substitution for that variable
337   = (Subst (in_scope `add_in_scope` new_var) 
338            (extendSubstEnv env old_var (DoneTy (TyVarTy new_var))),
339      new_var)
340   where
341     new_var = uniqAway in_scope old_var
342         -- The uniqAway part makes sure the new variable is not already in scope
343 \end{code}
344
345
346 %************************************************************************
347 %*                                                                      *
348 \section{Expression substitution}
349 %*                                                                      *
350 %************************************************************************
351
352 This expression substituter deals correctly with name capture.
353
354 BUT NOTE that substExpr silently discards the
355         unfolding, and
356         spec env
357 IdInfo attached to any binders in the expression.  It's quite
358 tricky to do them 'right' in the case of mutually recursive bindings,
359 and so far has proved unnecessary.
360
361 \begin{code}
362 substExpr :: Subst -> CoreExpr -> CoreExpr
363 substExpr subst expr
364         -- NB: we do not do a no-op when the substitution is empty,
365         -- because we always want to substitute the variables in the
366         -- in-scope set for their occurrences.  Why?
367         --      (a) because they may contain more information
368         --      (b) because leaving an un-substituted Id might cause
369         --          a space leak (its unfolding might point to an old version
370         --          of its right hand side).
371
372   = go expr
373   where
374     go (Var v) = -- See the notes at the top, with the Subst data type declaration
375                  case lookupIdSubst subst v of
376         
377                     ContEx env' e' -> substExpr (setSubstEnv subst env') e'
378                     DoneId v _     -> Var v
379                     DoneEx e'      -> e'
380
381     go (Type ty)      = Type (go_ty ty)
382     go (Lit lit)      = Lit lit
383     go (App fun arg)  = App (go fun) (go arg)
384     go (Note note e)  = Note (go_note note) (go e)
385
386     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
387                        where
388                          (subst', bndr') = substBndr subst bndr
389
390     go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
391                                     where
392                                       (subst', bndr') = substBndr subst bndr
393
394     go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
395                               where
396                                 (subst', bndrs') = substBndrs subst (map fst pairs)
397                                 pairs'  = bndrs' `zip` rhss'
398                                 rhss'   = map (substExpr subst' . snd) pairs
399
400     go (Case scrut bndr alts) = Case (go scrut) bndr' (map (go_alt subst') alts)
401                               where
402                                 (subst', bndr') = substBndr subst bndr
403
404     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
405                                  where
406                                    (subst', bndrs') = substBndrs subst bndrs
407
408     go_note (Coerce ty1 ty2) = Coerce (go_ty ty1) (go_ty ty2)
409     go_note note             = note
410
411     go_ty ty = substTy subst ty
412
413 \end{code}
414
415 Substituting in binders is a rather tricky part of the whole compiler.
416
417 When we hit a binder we may need to
418   (a) apply the the type envt (if non-empty) to its type
419   (c) give it a new unique to avoid name clashes
420
421 \begin{code}
422 substBndr :: Subst -> Var -> (Subst, Var)
423 substBndr subst bndr
424   | isTyVar bndr  = substTyVar subst bndr
425   | otherwise     = substId    subst bndr
426
427 substBndrs :: Subst -> [Var] -> (Subst, [Var])
428 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
429
430
431 substIds :: Subst -> [Id] -> (Subst, [Id])
432 substIds subst bndrs = mapAccumL substId subst bndrs
433
434 substId :: Subst -> Id -> (Subst, Id)
435         -- Returns an Id with empty IdInfo
436         -- See the notes with the Subst data type decl at the
437         -- top of this module
438
439 substId subst@(Subst in_scope env) old_id
440   = (Subst (in_scope `add_in_scope` new_id) new_env, new_id)
441   where
442     id_ty    = idType old_id
443     occ_info = idOccInfo old_id
444
445        -- id1 has its type zapped
446     id1 |  noTypeSubst env
447         || isEmptyVarSet (tyVarsOfType id_ty) = old_id
448                         -- The tyVarsOfType is cheaper than it looks
449                         -- because we cache the free tyvars of the type
450                         -- in a Note in the id's type itself
451         | otherwise  = setIdType old_id (substTy subst id_ty)
452
453         -- id2 has its IdInfo zapped
454     id2 = zapFragileIdInfo id1
455
456         -- new_id is cloned if necessary
457     new_id = uniqAway in_scope id2
458
459         -- Extend the substitution if the unique has changed,
460         -- or there's some useful occurrence information
461         -- See the notes with substTyVar for the delSubstEnv
462     new_env | new_id /= old_id || isFragileOccInfo occ_info 
463             = extendSubstEnv env old_id (DoneId new_id occ_info)
464             | otherwise 
465             = delSubstEnv env old_id
466 \end{code}
467
468 Now a variant that unconditionally allocates a new unique.
469
470 \begin{code}
471 substAndCloneIds :: Subst -> UniqSupply -> [Id] -> (Subst, UniqSupply, [Id])
472 substAndCloneIds subst us [] = (subst, us, [])
473 substAndCloneIds subst us (b:bs) = case substAndCloneId  subst  us  b  of { (subst1, us1, b') ->
474                                    case substAndCloneIds subst1 us1 bs of { (subst2, us2, bs') ->
475                                    (subst2, us2, (b':bs')) }}
476                                         
477 substAndCloneId :: Subst -> UniqSupply -> Id -> (Subst, UniqSupply, Id)
478 substAndCloneId subst@(Subst in_scope env) us old_id
479   = (Subst (in_scope `add_in_scope` new_id) 
480            (extendSubstEnv env old_id (DoneEx (Var new_id))),
481      new_us,
482      new_id)
483   where
484     id_ty    = idType old_id
485     id1 | noTypeSubst env || isEmptyVarSet (tyVarsOfType id_ty) = old_id
486         | otherwise                                             = setIdType old_id (substTy subst id_ty)
487
488     id2          = zapFragileIdInfo id1
489     new_id       = setVarUnique id2 (uniqFromSupply us1)
490     (us1,new_us) = splitUniqSupply us
491 \end{code}
492
493
494 %************************************************************************
495 %*                                                                      *
496 \section{IdInfo substitution}
497 %*                                                                      *
498 %************************************************************************
499
500 \begin{code}
501 substIdInfo :: Subst 
502             -> IdInfo           -- Get un-substituted ones from here
503             -> IdInfo           -- Substitute it and add it to here
504             -> IdInfo           -- To give this
505         -- Seq'ing on the returned IdInfo is enough to cause all the 
506         -- substitutions to happen completely
507
508 substIdInfo subst old_info new_info
509   = info2
510   where 
511     info1 | isEmptyCoreRules old_rules = new_info
512           | otherwise                  = new_info `setSpecInfo` new_rules
513                         -- setSpecInfo does a seq
514           where
515             new_rules = substRules subst old_rules
516  
517     info2 | not (workerExists old_wrkr) = info1
518           | otherwise                   = info1 `setWorkerInfo` new_wrkr
519                         -- setWorkerInfo does a seq
520           where
521             new_wrkr = substWorker subst old_wrkr
522
523     old_rules = specInfo   old_info
524     old_wrkr  = workerInfo old_info
525
526 substWorker :: Subst -> WorkerInfo -> WorkerInfo
527         -- Seq'ing on the returned WorkerInfo is enough to cause all the 
528         -- substitutions to happen completely
529
530 substWorker subst NoWorker
531   = NoWorker
532 substWorker subst (HasWorker w a)
533   = case lookupSubst subst w of
534         Nothing                -> HasWorker w a
535         Just (DoneId w1 _)     -> HasWorker w1 a
536         Just (DoneEx (Var w1)) -> HasWorker w1 a
537         Just (DoneEx other)    -> WARN( True, text "substWorker: DoneEx" <+> ppr w )
538                                   NoWorker      -- Worker has got substituted away altogether
539         Just (ContEx se1 e)    -> WARN( True, text "substWorker: ContEx" <+> ppr w <+> ppr e)
540                                   NoWorker      -- Ditto
541                         
542 substRules :: Subst -> CoreRules -> CoreRules
543         -- Seq'ing on the returned CoreRules is enough to cause all the 
544         -- substitutions to happen completely
545
546 substRules subst rules
547  | isEmptySubst subst = rules
548
549 substRules subst (Rules rules rhs_fvs)
550   = seqRules new_rules `seq` new_rules
551   where
552     new_rules = Rules (map do_subst rules)
553                       (subst_fvs (substEnv subst) rhs_fvs)
554
555     do_subst rule@(BuiltinRule _) = rule
556     do_subst (Rule name tpl_vars lhs_args rhs)
557         = Rule name tpl_vars' 
558                (map (substExpr subst') lhs_args)
559                (substExpr subst' rhs)
560         where
561           (subst', tpl_vars') = substBndrs subst tpl_vars
562
563     subst_fvs se fvs
564         = foldVarSet (unionVarSet . subst_fv) emptyVarSet rhs_fvs
565         where
566           subst_fv fv = case lookupSubstEnv se fv of
567                                 Nothing                   -> unitVarSet fv
568                                 Just (DoneId fv' _)       -> unitVarSet fv'
569                                 Just (DoneEx expr)        -> exprFreeVars expr
570                                 Just (DoneTy ty)          -> tyVarsOfType ty 
571                                 Just (ContEx se' expr) -> subst_fvs se' (exprFreeVars expr)
572 \end{code}