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