VarEnv.lhs 12.5 KB
Newer Older
Simon Marlow's avatar
Simon Marlow committed
1
2
%
% (c) The University of Glasgow 2006
3
4
5
6
7
8
9
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%

\begin{code}
module VarEnv (
	VarEnv, IdEnv, TyVarEnv,
	emptyVarEnv, unitVarEnv, mkVarEnv,
10
	elemVarEnv, varEnvElts, varEnvKeys,
11
	extendVarEnv, extendVarEnv_C, extendVarEnvList,
12
13
	plusVarEnv, plusVarEnv_C,
	delVarEnvList, delVarEnv,
14
	lookupVarEnv, lookupVarEnv_NF, lookupWithDefaultVarEnv,
15
16
	mapVarEnv, zipVarEnv,
	modifyVarEnv, modifyVarEnv_Directly,
17
	isEmptyVarEnv, foldVarEnv, 
18
	elemVarEnvByKey, lookupVarEnv_Directly,
19
	filterVarEnv_Directly,
20

21
22
	-- InScopeSet
	InScopeSet, emptyInScopeSet, mkInScopeSet, delInScopeSet,
23
24
	extendInScopeSet, extendInScopeSetList, extendInScopeSetSet, 
	modifyInScopeSet,
25
	getInScopeVars, lookupInScope, elemInScopeSet, uniqAway, 
26
	mapInScopeSet,
27

28
29
	-- RnEnv2 and its operations
	RnEnv2, mkRnEnv2, rnBndr2, rnBndrs2, rnOccL, rnOccR, inRnEnvL, inRnEnvR,
30
		rnBndrL, rnBndrR, nukeRnEnvL, nukeRnEnvR, extendRnInScopeList,
31
		rnInScope, lookupRnInScope,
32

33
34
	-- TidyEnvs
	TidyEnv, emptyTidyEnv
35
36
37
38
    ) where

#include "HsVersions.h"

Simon Marlow's avatar
Simon Marlow committed
39
40
import OccName
import Var
41
import VarSet
42
import UniqFM  
Simon Marlow's avatar
Simon Marlow committed
43
44
45
46
import Unique
import Util
import Maybes
import StaticFlags
47
48
import Outputable
import FastTypes
49
50
51
\end{code}


52
53
%************************************************************************
%*									*
54
		In-scope sets
55
56
57
58
%*									*
%************************************************************************

\begin{code}
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
data InScopeSet = InScope (VarEnv Var) FastInt
	-- The Int# is a kind of hash-value used by uniqAway
	-- For example, it might be the size of the set
	-- INVARIANT: it's not zero; we use it as a multiplier in uniqAway

instance Outputable InScopeSet where
  ppr (InScope s i) = ptext SLIT("InScope") <+> ppr s

emptyInScopeSet :: InScopeSet
emptyInScopeSet = InScope emptyVarSet 1#

getInScopeVars ::  InScopeSet -> VarEnv Var
getInScopeVars (InScope vs _) = vs

mkInScopeSet :: VarEnv Var -> InScopeSet
mkInScopeSet in_scope = InScope in_scope 1#

extendInScopeSet :: InScopeSet -> Var -> InScopeSet
extendInScopeSet (InScope in_scope n) v = InScope (extendVarEnv in_scope v v) (n +# 1#)

extendInScopeSetList :: InScopeSet -> [Var] -> InScopeSet
extendInScopeSetList (InScope in_scope n) vs
   = InScope (foldl (\s v -> extendVarEnv s v v) in_scope vs)
		    (n +# iUnbox (length vs))

84
85
86
87
extendInScopeSetSet :: InScopeSet -> VarEnv Var -> InScopeSet
extendInScopeSetSet (InScope in_scope n) vs
   = InScope (in_scope `plusVarEnv` vs) (n +# iUnbox (sizeUFM vs))

88
89
90
91
92
93
94
95
modifyInScopeSet :: InScopeSet -> Var -> Var -> InScopeSet
-- Exploit the fact that the in-scope "set" is really a map
-- 	Make old_v map to new_v
modifyInScopeSet (InScope in_scope n) old_v new_v = InScope (extendVarEnv in_scope old_v new_v) (n +# 1#)

delInScopeSet :: InScopeSet -> Var -> InScopeSet
delInScopeSet (InScope in_scope n) v = InScope (in_scope `delVarEnv` v) n

96
97
98
mapInScopeSet :: (Var -> Var) -> InScopeSet -> InScopeSet
mapInScopeSet f (InScope in_scope n) = InScope (mapVarEnv f in_scope) n

99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
elemInScopeSet :: Var -> InScopeSet -> Bool
elemInScopeSet v (InScope in_scope n) = v `elemVarEnv` in_scope

lookupInScope :: InScopeSet -> Var -> Maybe Var
-- It's important to look for a fixed point
-- When we see (case x of y { I# v -> ... })
-- we add  [x -> y] to the in-scope set (Simplify.simplCaseBinder).
-- When we lookup up an occurrence of x, we map to y, but then
-- we want to look up y in case it has acquired more evaluation information by now.
lookupInScope (InScope in_scope n) v 
  = go v
  where
    go v = case lookupVarEnv in_scope v of
		Just v' | v == v'   -> Just v'	-- Reached a fixed point
			| otherwise -> go v'
		Nothing		    -> Nothing
\end{code}
116

117
118
119
120
121
122
\begin{code}
uniqAway :: InScopeSet -> Var -> Var
-- (uniqAway in_scope v) finds a unique that is not used in the
-- in-scope set, and gives that to v.  It starts with v's current unique, of course,
-- in the hope that it won't have to change it, and thereafter uses a combination
-- of that and the hash-code found in the in-scope set
123
124
125
126
127
128
129
130
uniqAway in_scope var
  | var `elemInScopeSet` in_scope = uniqAway' in_scope var	-- Make a new one
  | otherwise 			  = var				-- Nothing to do

uniqAway' :: InScopeSet -> Var -> Var
-- This one *always* makes up a new variable
uniqAway' (InScope set n) var
  = try 1#
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
  where
    orig_unique = getUnique var
    try k 
#ifdef DEBUG
	  | k ># 1000#
	  = pprPanic "uniqAway loop:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n)) 
#endif			    
	  | uniq `elemVarSetByKey` set = try (k +# 1#)
#ifdef DEBUG
	  | opt_PprStyle_Debug && k ># 3#
	  = pprTrace "uniqAway:" (ppr (iBox k) <+> text "tries" <+> ppr var <+> int (iBox n)) 
	    setVarUnique var uniq
#endif			    
	  | otherwise = setVarUnique var uniq
	  where
	    uniq = deriveUnique orig_unique (iBox (n *# k))
147
148
149
\end{code}


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
%************************************************************************
%*									*
		Dual renaming
%*									*
%************************************************************************

When we are comparing (or matching) types or terms, we are faced with 
"going under" corresponding binders.  E.g. when comparing
	\x. e1	~   \y. e2

Basically we want to rename [x->y] or [y->x], but there are lots of 
things we must be careful of.  In particular, x might be free in e2, or
y in e1.  So the idea is that we come up with a fresh binder that is free
in neither, and rename x and y respectively.  That means we must maintain
	a) a renaming for the left-hand expression
	b) a renaming for the right-hand expressions
	c) an in-scope set

Furthermore, when matching, we want to be able to have an 'occurs check',
to prevent
	\x. f   ~   \y. y
matching with f->y.  So for each expression we want to know that set of
locally-bound variables. That is precisely the domain of the mappings (a)
and (b), but we must ensure that we always extend the mappings as we go in.


\begin{code}
data RnEnv2 
  = RV2 { envL 	   :: VarEnv Var	-- Renaming for Left term
	, envR 	   :: VarEnv Var	-- Renaming for Right term
	, in_scope :: InScopeSet }	-- In scope in left or right terms

-- The renamings envL and envR are *guaranteed* to contain a binding
-- for every variable bound as we go into the term, even if it is not
-- renamed.  That way we can ask what variables are locally bound
-- (inRnEnvL, inRnEnvR)

mkRnEnv2 :: InScopeSet -> RnEnv2
mkRnEnv2 vars = RV2	{ envL 	   = emptyVarEnv 
			, envR 	   = emptyVarEnv
			, in_scope = vars }

192
193
194
195
196
197
198
extendRnInScopeList :: RnEnv2 -> [Var] -> RnEnv2
extendRnInScopeList env vs
  = env { in_scope = extendInScopeSetList (in_scope env) vs }

rnInScope :: Var -> RnEnv2 -> Bool
rnInScope x env = x `elemInScopeSet` in_scope env

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
rnBndrs2 :: RnEnv2 -> [Var] -> [Var] -> RnEnv2
-- Arg lists must be of equal length
rnBndrs2 env bsL bsR = foldl2 rnBndr2 env bsL bsR 

rnBndr2 :: RnEnv2 -> Var -> Var -> RnEnv2
-- (rnBndr2 env bL bR) go under a binder bL in the Left term 1, 
-- 		       and binder bR in the Right term
-- It finds a new binder, new_b,
-- and returns an environment mapping bL->new_b and bR->new_b resp.
rnBndr2 (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL bR
  = RV2 { envL 	   = extendVarEnv envL bL new_b	  -- See Note
	, envR 	   = extendVarEnv envR bR new_b	  -- [Rebinding]
	, in_scope = extendInScopeSet in_scope new_b }
  where
	-- Find a new binder not in scope in either term
    new_b | not (bL `elemInScopeSet` in_scope) = bL
      	  | not (bR `elemInScopeSet` in_scope) = bR
      	  | otherwise			       = uniqAway' in_scope bL

	-- Note [Rebinding]
	-- If the new var is the same as the old one, note that
	-- the extendVarEnv *deletes* any current renaming
	-- E.g.	  (\x. \x. ...)	 ~  (\y. \z. ...)
	--
	--   Inside \x  \y	{ [x->y], [y->y],       {y} }
	-- 	 \x  \z	  	{ [x->x], [y->y, z->x], {y,x} }

rnBndrL, rnBndrR :: RnEnv2 -> Var -> (RnEnv2, Var)
-- Used when there's a binder on one side or the other only
-- Useful when eta-expanding
229
-- 
230
231
rnBndrL (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bL
  = (RV2 { envL     = extendVarEnv envL bL new_b
232
	 , envR     = extendVarEnv envR new_b new_b 	-- Note [rnBndrLR]
233
234
	 , in_scope = extendInScopeSet in_scope new_b }, new_b)
  where
235
    new_b = uniqAway in_scope bL
236
237

rnBndrR (RV2 { envL = envL, envR = envR, in_scope = in_scope }) bR
238
  = (RV2 { envL     = extendVarEnv envL new_b new_b	-- Note [rnBndrLR]
239
240
241
	 , envR     = extendVarEnv envR bR new_b
	 , in_scope = extendInScopeSet in_scope new_b }, new_b)
  where
242
    new_b = uniqAway in_scope bR
243

244
245
246
247
248
249
250
251
252
253
254
255
256
-- Note [rnBndrLR] 
-- ~~~~~~~~~~~~~~~
-- Notice that in rnBndrL, rnBndrR, we extend envR, envL respectively
-- with a binding [new_b -> new_b], where new_b is the new binder.
-- This is important when doing eta expansion; e.g. matching (\x.M) ~ N
-- In effect we switch to (\x'.M) ~ (\x'.N x'), where x' is new_b
-- So we must add x' to the env of both L and R.  (x' is fresh, so it
-- can't capture anything in N.)  
--
-- If we don't do this, we can get silly matches like
--	forall a.  \y.a  ~   v
-- succeeding with [x -> v y], which is bogus of course 

257
258
259
260
261
262
263
rnOccL, rnOccR :: RnEnv2 -> Var -> Var
-- Look up the renaming of an occurrence in the left or right term
rnOccL (RV2 { envL = env }) v = lookupVarEnv env v `orElse` v
rnOccR (RV2 { envR = env }) v = lookupVarEnv env v `orElse` v

inRnEnvL, inRnEnvR :: RnEnv2 -> Var -> Bool
-- Tells whether a variable is locally bound
264
265
inRnEnvL (RV2 { envL = env }) v = v `elemVarEnv` env
inRnEnvR (RV2 { envR = env }) v = v `elemVarEnv` env
266

267
268
269
lookupRnInScope :: RnEnv2 -> Var -> Var
lookupRnInScope env v = lookupInScope (in_scope env) v `orElse` v

270
271
272
273
274
275
nukeRnEnvL, nukeRnEnvR :: RnEnv2 -> RnEnv2
nukeRnEnvL env = env { envL = emptyVarEnv }
nukeRnEnvR env = env { envR = emptyVarEnv }
\end{code}


276
277
%************************************************************************
%*									*
278
		Tidying
279
280
281
%*									*
%************************************************************************

282
283
When tidying up print names, we keep a mapping of in-scope occ-names
(the TidyOccEnv) and a Var-to-Var of the current renamings.
284

285
286
\begin{code}
type TidyEnv = (TidyOccEnv, VarEnv Var)
287

288
289
emptyTidyEnv :: TidyEnv
emptyTidyEnv = (emptyTidyOccEnv, emptyVarEnv)
290
291
292
\end{code}


293
294
295
296
297
298
299
300
301
302
303
304
%************************************************************************
%*									*
\subsection{@VarEnv@s}
%*									*
%************************************************************************

\begin{code}
type VarEnv elt   = UniqFM elt
type IdEnv elt    = VarEnv elt
type TyVarEnv elt = VarEnv elt

emptyVarEnv	  :: VarEnv a
305
306
307
308
mkVarEnv	  :: [(Var, a)] -> VarEnv a
zipVarEnv	  :: [Var] -> [a] -> VarEnv a
unitVarEnv	  :: Var -> a -> VarEnv a
extendVarEnv	  :: VarEnv a -> Var -> a -> VarEnv a
309
extendVarEnv_C	  :: (a->a->a) -> VarEnv a -> Var -> a -> VarEnv a
310
plusVarEnv	  :: VarEnv a -> VarEnv a -> VarEnv a
311
extendVarEnvList  :: VarEnv a -> [(Var, a)] -> VarEnv a
312
		  
313
314
lookupVarEnv_Directly :: VarEnv a -> Unique -> Maybe a
filterVarEnv_Directly :: (Unique -> a -> Bool) -> VarEnv a -> VarEnv a
315
316
delVarEnvList     :: VarEnv a -> [Var] -> VarEnv a
delVarEnv	  :: VarEnv a -> Var -> VarEnv a
317
318
plusVarEnv_C	  :: (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
mapVarEnv	  :: (a -> b) -> VarEnv a -> VarEnv b
319
modifyVarEnv	  :: (a -> a) -> VarEnv a -> Var -> VarEnv a
320
varEnvElts	  :: VarEnv a -> [a]
321
varEnvKeys	  :: VarEnv a -> [Unique]
322
323
		  
isEmptyVarEnv	  :: VarEnv a -> Bool
324
325
lookupVarEnv	  :: VarEnv a -> Var -> Maybe a
lookupVarEnv_NF   :: VarEnv a -> Var -> a
326
lookupWithDefaultVarEnv :: VarEnv a -> a -> Var -> a
327
elemVarEnv	  :: Var -> VarEnv a -> Bool
328
elemVarEnvByKey   :: Unique -> VarEnv a -> Bool
329
330
331
332
333
foldVarEnv	  :: (a -> b -> b) -> b -> VarEnv a -> b
\end{code}

\begin{code}
elemVarEnv       = elemUFM
334
elemVarEnvByKey  = elemUFM_Directly
335
extendVarEnv	 = addToUFM
336
extendVarEnv_C	 = addToUFM_C
337
extendVarEnvList = addListToUFM
338
339
340
341
342
plusVarEnv_C	 = plusUFM_C
delVarEnvList	 = delListFromUFM
delVarEnv	 = delFromUFM
plusVarEnv	 = plusUFM
lookupVarEnv	 = lookupUFM
343
lookupWithDefaultVarEnv = lookupWithDefaultUFM
344
345
346
mapVarEnv	 = mapUFM
mkVarEnv	 = listToUFM
emptyVarEnv	 = emptyUFM
347
varEnvElts	 = eltsUFM
348
varEnvKeys	 = keysUFM
349
350
351
unitVarEnv	 = unitUFM
isEmptyVarEnv	 = isNullUFM
foldVarEnv	 = foldUFM
352
353
lookupVarEnv_Directly = lookupUFM_Directly
filterVarEnv_Directly = filterUFM_Directly
354

355
356
zipVarEnv tyvars tys   = mkVarEnv (zipEqual "zipVarEnv" tyvars tys)
lookupVarEnv_NF env id = case (lookupVarEnv env id) of { Just xx -> xx }
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
\end{code}

@modifyVarEnv@: Look up a thing in the VarEnv, 
then mash it with the modify function, and put it back.

\begin{code}
modifyVarEnv mangle_fn env key
  = case (lookupVarEnv env key) of
      Nothing -> env
      Just xx -> extendVarEnv env key (mangle_fn xx)

modifyVarEnv_Directly mangle_fn env key
  = case (lookupUFM_Directly env key) of
      Nothing -> env
      Just xx -> addToUFM_Directly env key (mangle_fn xx)
\end{code}