CoreSubst.lhs 13 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
%
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section[CoreUtils]{Utility functions on @Core@ syntax}

\begin{code}
module CoreSubst (
	-- Substitution stuff
	Subst, TvSubstEnv, IdSubstEnv, InScopeSet,

11
	substTy, substExpr, substSpec, substWorker,
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
	lookupIdSubst, lookupTvSubst, 

	emptySubst, mkEmptySubst, mkSubst, substInScope, isEmptySubst, 
 	extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
	extendInScope, extendInScopeIds,
	isInScope,

	-- Binders
	substBndr, substBndrs, substRecBndrs,
	cloneIdBndr, cloneIdBndrs, cloneRecIdBndrs
    ) where

#include "HsVersions.h"

import CoreSyn		( Expr(..), Bind(..), Note(..), CoreExpr,
27
			  CoreRule(..), hasUnfolding, noUnfolding
28
29
30
31
32
33
34
35
36
37
			)
import CoreFVs		( exprFreeVars )
import CoreUtils	( exprIsTrivial )

import qualified Type	( substTy, substTyVarBndr )
import Type		( Type, tyVarsOfType, TvSubstEnv, TvSubst(..), mkTyVarTy )
import VarSet
import VarEnv
import Var		( setVarUnique, isId )
import Id		( idType, setIdType, maybeModifyIdInfo, isLocalId )
38
39
import IdInfo		( IdInfo, SpecInfo(..), specInfo, setSpecInfo, isEmptySpecInfo,
			  unfoldingInfo, setUnfoldingInfo, seqSpecInfo,
40
			  WorkerInfo(..), workerExists, workerInfo, setWorkerInfo
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
			)
import Unique		( Unique )
import UniqSupply	( UniqSupply, uniqFromSupply, uniqsFromSupply )
import Var		( Var, Id, TyVar, isTyVar )
import Maybes		( orElse )
import Outputable
import PprCore		()		-- Instances
import Util		( mapAccumL )
import FastTypes
\end{code}


%************************************************************************
%*									*
\subsection{Substitutions}
%*									*
%************************************************************************

\begin{code}
data Subst 
  = Subst InScopeSet	-- Variables in in scope (both Ids and TyVars)
	  IdSubstEnv	-- Substitution for Ids
	  TvSubstEnv	-- Substitution for TyVars

	-- INVARIANT 1: The (domain of the) in-scope set is a superset
	-- 	        of the free vars of the range of the substitution
	--		that might possibly clash with locally-bound variables
	--		in the thing being substituted in.
	-- This is what lets us deal with name capture properly
	-- It's a hard invariant to check...
	-- There are various ways of causing it to happen:
	--	- arrange that the in-scope set really is all the things in scope
	--	- arrange that it's the free vars of the range of the substitution
	--	- make it empty because all the free vars of the subst are fresh,
	--		and hence can't possibly clash.a
	--
	-- INVARIANT 2: The substitution is apply-once; see notes with
	--		Types.TvSubstEnv

type IdSubstEnv = IdEnv CoreExpr

----------------------------
isEmptySubst :: Subst -> Bool
isEmptySubst (Subst _ id_env tv_env) = isEmptyVarEnv id_env && isEmptyVarEnv tv_env

emptySubst :: Subst
emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv

mkEmptySubst :: InScopeSet -> Subst
mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv

mkSubst :: InScopeSet -> TvSubstEnv -> IdSubstEnv -> Subst
mkSubst in_scope tvs ids = Subst in_scope ids tvs

-- getTvSubst :: Subst -> TvSubst
-- getTvSubst (Subst in_scope _ tv_env) = TvSubst in_scope tv_env

-- getTvSubstEnv :: Subst -> TvSubstEnv
-- getTvSubstEnv (Subst _ _ tv_env) = tv_env
-- 
-- setTvSubstEnv :: Subst -> TvSubstEnv -> Subst
-- setTvSubstEnv (Subst in_scope ids _) tvs = Subst in_scope ids tvs

substInScope :: Subst -> InScopeSet
substInScope (Subst in_scope _ _) = in_scope

-- zapSubstEnv :: Subst -> Subst
-- zapSubstEnv (Subst in_scope _ _) = Subst in_scope emptyVarEnv emptyVarEnv

-- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
extendIdSubst (Subst in_scope ids tvs) v r = Subst in_scope (extendVarEnv ids v r) tvs

extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
extendIdSubstList (Subst in_scope ids tvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs

extendTvSubst :: Subst -> TyVar -> Type -> Subst
extendTvSubst (Subst in_scope ids tvs) v r = Subst in_scope ids (extendVarEnv tvs v r) 

extendTvSubstList :: Subst -> [(TyVar,Type)] -> Subst
extendTvSubstList (Subst in_scope ids tvs) prs = Subst in_scope ids (extendVarEnvList tvs prs)

lookupIdSubst :: Subst -> Id -> CoreExpr
lookupIdSubst (Subst in_scope ids tvs) v 
  | not (isLocalId v) = Var v
  | otherwise
  = case lookupVarEnv ids v of {
	Just e  -> e ;
	Nothing -> 	
    case lookupInScope in_scope v of {
	-- Watch out!  Must get the Id from the in-scope set,
	-- because its type there may differ
	Just v  -> Var v ;
	Nothing -> WARN( True, ptext SLIT("CoreSubst.lookupIdSubst") <+> ppr v ) 
		   Var v
    }}

lookupTvSubst :: Subst -> TyVar -> Type
lookupTvSubst (Subst _ ids tvs) v = lookupVarEnv tvs v `orElse` mkTyVarTy v

------------------------------
isInScope :: Var -> Subst -> Bool
isInScope v (Subst in_scope _ _) = v `elemInScopeSet` in_scope

extendInScope :: Subst -> Var -> Subst
extendInScope (Subst in_scope ids tvs) v
  = Subst (in_scope `extendInScopeSet` v) 
	  (ids `delVarEnv` v) (tvs `delVarEnv` v)

extendInScopeIds :: Subst -> [Id] -> Subst
extendInScopeIds (Subst in_scope ids tvs) vs 
  = Subst (in_scope `extendInScopeSetList` vs) 
	  (ids `delVarEnvList` vs) tvs
\end{code}

Pretty printing, for debugging only

\begin{code}
instance Outputable Subst where
  ppr (Subst in_scope ids tvs) 
	=  ptext SLIT("<InScope =") <+> braces (fsep (map ppr (varEnvElts (getInScopeVars in_scope))))
	$$ ptext SLIT(" IdSubst   =") <+> ppr ids
	$$ ptext SLIT(" TvSubst   =") <+> ppr tvs
 	 <> char '>'
\end{code}


%************************************************************************
%*									*
	Substituting expressions
%*									*
%************************************************************************

\begin{code}
substExpr :: Subst -> CoreExpr -> CoreExpr
substExpr subst expr
  = go expr
  where
    go (Var v)	       = lookupIdSubst subst v 
    go (Type ty)       = Type (substTy subst ty)
    go (Lit lit)       = Lit lit
    go (App fun arg)   = App (go fun) (go arg)
    go (Note note e)   = Note (go_note note) (go e)
    go (Lam bndr body) = Lam bndr' (substExpr subst' body)
		       where
			 (subst', bndr') = substBndr subst bndr

    go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
				    where
				      (subst', bndr') = substBndr subst bndr

    go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
			      where
				(subst', bndrs') = substRecBndrs subst (map fst pairs)
				pairs'	= bndrs' `zip` rhss'
				rhss'	= map (substExpr subst' . snd) pairs

    go (Case scrut bndr ty alts) = Case (go scrut) bndr' (substTy subst ty) (map (go_alt subst') alts)
			         where
			  	 (subst', bndr') = substBndr subst bndr

    go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
				 where
				   (subst', bndrs') = substBndrs subst bndrs

    go_note (Coerce ty1 ty2) = Coerce (substTy subst ty1) (substTy subst ty2)
    go_note note	     = note
\end{code}


%************************************************************************
%*									*
	Substituting binders
%*									*
%************************************************************************

Remember that substBndr and friends are used when doing expression
substitution only.  Their only business is substitution, so they
preserve all IdInfo (suitably substituted).  For example, we *want* to
preserve occ info in rules.

\begin{code}
substBndr :: Subst -> Var -> (Subst, Var)
substBndr subst bndr
  | isTyVar bndr  = substTyVarBndr subst bndr
  | otherwise     = substIdBndr subst subst bndr

substBndrs :: Subst -> [Var] -> (Subst, [Var])
substBndrs subst bndrs = mapAccumL substBndr subst bndrs

substRecBndrs :: Subst -> [Id] -> (Subst, [Id])
-- Substitute a mutually recursive group
substRecBndrs subst bndrs 
  = (new_subst, new_bndrs)
  where		-- Here's the reason we need to pass rec_subst to subst_id
    (new_subst, new_bndrs) = mapAccumL (substIdBndr new_subst) subst bndrs
\end{code}


\begin{code}
substIdBndr :: Subst		-- Substitution to use for the IdInfo
	    -> Subst -> Id 	-- Substitition and Id to transform
	    -> (Subst, Id)	-- Transformed pair

substIdBndr rec_subst subst@(Subst in_scope env tvs) old_id
  = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
  where
    id1 = uniqAway in_scope old_id	-- id1 is cloned if necessary
    id2 = substIdType subst id1		-- id2 has its type zapped

	-- new_id has the right IdInfo
	-- The lazy-set is because we're in a loop here, with 
	-- rec_subst, when dealing with a mutually-recursive group
    new_id = maybeModifyIdInfo (substIdInfo rec_subst) id2

	-- Extend the substitution if the unique has changed
	-- See the notes with substTyVarBndr for the delSubstEnv
    new_env | new_id /= old_id  = extendVarEnv env old_id (Var new_id)
	    | otherwise         = delVarEnv env old_id
\end{code}

Now a variant that unconditionally allocates a new unique.
It also unconditionally zaps the OccInfo.

\begin{code}
cloneIdBndr :: Subst -> UniqSupply -> Id -> (Subst, Id)
cloneIdBndr subst us old_id
  = clone_id subst subst (old_id, uniqFromSupply us)

cloneIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
cloneIdBndrs subst us ids
  = mapAccumL (clone_id subst) subst (ids `zip` uniqsFromSupply us)

cloneRecIdBndrs :: Subst -> UniqSupply -> [Id] -> (Subst, [Id])
cloneRecIdBndrs subst us ids
  = (subst', ids')
  where
    (subst', ids') = mapAccumL (clone_id subst') subst
			       (ids `zip` uniqsFromSupply us)

-- Just like substIdBndr, except that it always makes a new unique
-- It is given the unique to use
clone_id    :: Subst			-- Substitution for the IdInfo
	    -> Subst -> (Id, Unique)	-- Substitition and Id to transform
	    -> (Subst, Id)		-- Transformed pair

clone_id rec_subst subst@(Subst in_scope env tvs) (old_id, uniq)
  = (Subst (in_scope `extendInScopeSet` new_id) new_env tvs, new_id)
  where
    id1	    = setVarUnique old_id uniq
    id2     = substIdType subst id1
    new_id  = maybeModifyIdInfo (substIdInfo rec_subst) id2
    new_env = extendVarEnv env old_id (Var new_id)
\end{code}


%************************************************************************
%*									*
		Types
%*									*
%************************************************************************

For types we just call the corresponding function in Type, but we have
to repackage the substitution, from a Subst to a TvSubst

\begin{code}
substTyVarBndr :: Subst -> TyVar -> (Subst, TyVar)
substTyVarBndr (Subst in_scope id_env tv_env) tv
  = case Type.substTyVarBndr (TvSubst in_scope tv_env) tv of
	(TvSubst in_scope' tv_env', tv') 
	   -> (Subst in_scope' id_env tv_env', tv')

substTy :: Subst -> Type -> Type 
substTy (Subst in_scope id_env tv_env) ty 
  = Type.substTy (TvSubst in_scope tv_env) ty
\end{code}


%************************************************************************
%*									*
\section{IdInfo substitution}
%*									*
%************************************************************************

\begin{code}
substIdType :: Subst -> Id -> Id
substIdType subst@(Subst in_scope id_env tv_env) id
  | isEmptyVarEnv tv_env || isEmptyVarSet (tyVarsOfType old_ty) = id
  | otherwise	= setIdType id (substTy subst old_ty)
		-- The tyVarsOfType is cheaper than it looks
		-- because we cache the free tyvars of the type
		-- in a Note in the id's type itself
  where
    old_ty = idType id

------------------
substIdInfo :: Subst -> IdInfo -> Maybe IdInfo
-- Always zaps the unfolding, to save substitution work
substIdInfo  subst info
  | nothing_to_do = Nothing
341
  | otherwise     = Just (info `setSpecInfo`   	  substSpec  subst old_rules
342
343
344
345
346
			       `setWorkerInfo` 	  substWorker subst old_wrkr
			       `setUnfoldingInfo` noUnfolding)
  where
    old_rules 	  = specInfo info
    old_wrkr  	  = workerInfo info
347
    nothing_to_do = isEmptySpecInfo old_rules &&
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
		    not (workerExists old_wrkr) &&
		    not (hasUnfolding (unfoldingInfo info))
    

------------------
substWorker :: Subst -> WorkerInfo -> WorkerInfo
	-- Seq'ing on the returned WorkerInfo is enough to cause all the 
	-- substitutions to happen completely

substWorker subst NoWorker
  = NoWorker
substWorker subst (HasWorker w a)
  = case lookupIdSubst subst w of
	Var w1 -> HasWorker w1 a
	other  -> WARN( not (exprIsTrivial other), text "CoreSubst.substWorker:" <+> ppr w )
		  NoWorker	-- Worker has got substituted away altogether
				-- (This can happen if it's trivial, 
				--  via postInlineUnconditionally, hence warning)

------------------
368
substSpec :: Subst -> SpecInfo -> SpecInfo
369

370
371
372
373
374
substSpec subst spec@(SpecInfo rules rhs_fvs)
  | isEmptySubst subst
  = spec
  | otherwise
  = seqSpecInfo new_rules `seq` new_rules
375
  where
376
    new_rules = SpecInfo (map do_subst rules) (substVarSet subst rhs_fvs)
377

378
379
380
381
382
    do_subst rule@(BuiltinRule {}) = rule
    do_subst rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs })
	= rule { ru_bndrs = bndrs',
		 ru_args  = map (substExpr subst') args,
		 ru_rhs   = substExpr subst' rhs }
383
	where
384
	  (subst', bndrs') = substBndrs subst bndrs
385
386
387
388
389
390
391
392
393

------------------
substVarSet subst fvs 
  = foldVarSet (unionVarSet . subst_fv subst) emptyVarSet fvs
  where
    subst_fv subst fv 
	| isId fv   = exprFreeVars (lookupIdSubst subst fv)
	| otherwise = tyVarsOfType (lookupTvSubst subst fv)
\end{code}