[project @ 2004-09-30 10:35:15 by simonpj]
[ghc-hetmet.git] / ghc / compiler / coreSyn / CoreFVs.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 Taken quite directly from the Peyton Jones/Lester paper.
5
6 \begin{code}
7 module CoreFVs (
8         exprFreeVars,   -- CoreExpr -> VarSet   -- Find all locally-defined free Ids or tyvars
9         exprsFreeVars,  -- [CoreExpr] -> VarSet
10
11         exprSomeFreeVars, exprsSomeFreeVars,
12
13         idRuleVars, idFreeVars, idFreeTyVars,
14         ruleRhsFreeVars, ruleLhsFreeNames, ruleLhsFreeIds, 
15
16         CoreExprWithFVs,        -- = AnnExpr Id VarSet
17         CoreBindWithFVs,        -- = AnnBind Id VarSet
18         freeVars,               -- CoreExpr -> CoreExprWithFVs
19         freeVarsOf              -- CoreExprWithFVs -> IdSet
20     ) where
21
22 #include "HsVersions.h"
23
24 import CoreSyn
25 import Id               ( Id, idType, idSpecialisation )
26 import NameSet
27 import VarSet
28 import Var              ( Var, isId, isLocalVar, varName )
29 import Type             ( tyVarsOfType )
30 import TcType           ( tyClsNamesOfType )
31 import Util             ( mapAndUnzip )
32 import Outputable
33 \end{code}
34
35
36 %************************************************************************
37 %*                                                                      *
38 \section{Finding the free variables of an expression}
39 %*                                                                      *
40 %************************************************************************
41
42 This function simply finds the free variables of an expression.
43 So far as type variables are concerned, it only finds tyvars that are
44
45         * free in type arguments, 
46         * free in the type of a binder,
47
48 but not those that are free in the type of variable occurrence.
49
50 \begin{code}
51 exprFreeVars :: CoreExpr -> VarSet      -- Find all locally-defined free Ids or tyvars
52 exprFreeVars = exprSomeFreeVars isLocalVar
53
54 exprsFreeVars :: [CoreExpr] -> VarSet
55 exprsFreeVars = foldr (unionVarSet . exprFreeVars) emptyVarSet
56
57 exprSomeFreeVars :: InterestingVarFun   -- Says which Vars are interesting
58                  -> CoreExpr
59                  -> VarSet
60 exprSomeFreeVars fv_cand e = expr_fvs e fv_cand emptyVarSet
61
62 exprsSomeFreeVars :: InterestingVarFun  -- Says which Vars are interesting
63                   -> [CoreExpr]
64                   -> VarSet
65 exprsSomeFreeVars fv_cand = foldr (unionVarSet . exprSomeFreeVars fv_cand) emptyVarSet
66
67 type InterestingVarFun = Var -> Bool    -- True <=> interesting
68 \end{code}
69
70
71 \begin{code}
72 type FV = InterestingVarFun 
73           -> VarSet             -- In scope
74           -> VarSet             -- Free vars
75
76 union :: FV -> FV -> FV
77 union fv1 fv2 fv_cand in_scope = fv1 fv_cand in_scope `unionVarSet` fv2 fv_cand in_scope
78
79 noVars :: FV
80 noVars fv_cand in_scope = emptyVarSet
81
82 -- At a variable occurrence, add in any free variables of its rule rhss
83 -- Curiously, we gather the Id's free *type* variables from its binding
84 -- site, but its free *rule-rhs* variables from its usage sites.  This
85 -- is a little weird.  The reason is that the former is more efficient,
86 -- but the latter is more fine grained, and a makes a difference when
87 -- a variable mentions itself one of its own rule RHSs
88 oneVar :: Id -> FV
89 oneVar var fv_cand in_scope
90   = ASSERT( isId var ) 
91     foldVarSet add_rule_var var_itself_set (idRuleVars var)
92   where
93     var_itself_set | keep_it fv_cand in_scope var = unitVarSet var
94                    | otherwise                = emptyVarSet
95     add_rule_var var set | keep_it fv_cand in_scope var = extendVarSet set var
96                          | otherwise                    = set
97
98 someVars :: VarSet -> FV
99 someVars vars fv_cand in_scope
100   = filterVarSet (keep_it fv_cand in_scope) vars
101
102 keep_it fv_cand in_scope var
103   | var `elemVarSet` in_scope = False
104   | fv_cand var               = True
105   | otherwise                 = False
106
107
108 addBndr :: CoreBndr -> FV -> FV
109 addBndr bndr fv fv_cand in_scope
110   | isId bndr = inside_fvs `unionVarSet` someVars (idFreeTyVars bndr) fv_cand in_scope
111   | otherwise = inside_fvs
112   where
113     inside_fvs = fv fv_cand (in_scope `extendVarSet` bndr) 
114
115 addBndrs :: [CoreBndr] -> FV -> FV
116 addBndrs bndrs fv = foldr addBndr fv bndrs
117 \end{code}
118
119
120 \begin{code}
121 expr_fvs :: CoreExpr -> FV
122
123 expr_fvs (Type ty)       = someVars (tyVarsOfType ty)
124 expr_fvs (Var var)       = oneVar var
125 expr_fvs (Lit lit)       = noVars
126 expr_fvs (Note _ expr)   = expr_fvs expr
127 expr_fvs (App fun arg)   = expr_fvs fun `union` expr_fvs arg
128 expr_fvs (Lam bndr body) = addBndr bndr (expr_fvs body)
129
130 -- gaw 2004
131 expr_fvs (Case scrut bndr ty alts)
132   = expr_fvs scrut `union` someVars (tyVarsOfType ty) `union` addBndr bndr  
133       (foldr (union . alt_fvs) noVars alts)
134   where
135     alt_fvs (con, bndrs, rhs) = addBndrs bndrs (expr_fvs rhs)
136
137 expr_fvs (Let (NonRec bndr rhs) body)
138   = expr_fvs rhs `union` addBndr bndr (expr_fvs body)
139
140 expr_fvs (Let (Rec pairs) body)
141   = addBndrs bndrs (foldr (union . expr_fvs) (expr_fvs body) rhss)
142   where
143     (bndrs,rhss) = unzip pairs
144 \end{code}
145
146
147 %************************************************************************
148 %*                                                                      *
149 \section{Free names}
150 %*                                                                      *
151 %************************************************************************
152
153 exprFreeNames finds the free *names* of an expression, notably
154 including the names of type constructors (which of course do not show
155 up in exprFreeVars).  Similarly ruleLhsFreeNames.  The latter is used
156 when deciding whether a rule is an orphan.  In particular, suppose that
157 T is defined in this module; we want to avoid declaring that a rule like
158         fromIntegral T = fromIntegral_T
159 is an orphan.  Of course it isn't, an declaring it an orphan would
160 make the whole module an orphan module, which is bad.
161
162 \begin{code}
163 ruleLhsFreeNames :: IdCoreRule -> NameSet
164 ruleLhsFreeNames (fn, BuiltinRule _ _) = unitNameSet (varName fn)
165 ruleLhsFreeNames (fn, Rule _ _ tpl_vars tpl_args rhs)
166   = addOneToNameSet (exprsFreeNames tpl_args `del_binders` tpl_vars) (varName fn)
167
168 exprFreeNames :: CoreExpr -> NameSet
169 exprFreeNames (Var v)     = unitNameSet (varName v)
170 exprFreeNames (Lit _)     = emptyNameSet
171 exprFreeNames (Type ty)   = tyClsNamesOfType ty -- Don't need free tyvars
172 exprFreeNames (App e1 e2) = exprFreeNames e1 `unionNameSets` exprFreeNames e2
173 exprFreeNames (Lam v e)   = exprFreeNames e `delFromNameSet` varName v
174 exprFreeNames (Note n e)  = exprFreeNames e
175
176 exprFreeNames (Let (NonRec b r) e) = (exprFreeNames e `delFromNameSet` varName b)
177                                      `unionNameSets` exprFreeNames r
178
179 exprFreeNames (Let (Rec prs) e) = (exprsFreeNames rs `unionNameSets` exprFreeNames e)
180                                   `del_binders` bs
181                                 where
182                                   (bs, rs) = unzip prs
183
184 -- gaw 2004
185 exprFreeNames (Case e b ty as) = exprFreeNames e `unionNameSets` tyClsNamesOfType ty 
186                                  `unionNameSets`
187                                  (unionManyNameSets (map altFreeNames as) `delFromNameSet` varName b)
188
189 -- Helpers
190 altFreeNames (_,bs,r) = exprFreeNames r `del_binders` bs
191
192 exprsFreeNames es = foldr (unionNameSets . exprFreeNames) emptyNameSet es
193
194 del_binders :: NameSet -> [Var] -> NameSet
195 del_binders names bndrs = foldl (\s b -> delFromNameSet s (varName b)) names bndrs
196 \end{code}
197
198 %************************************************************************
199 %*                                                                      *
200 \section[freevars-everywhere]{Attaching free variables to every sub-expression}
201 %*                                                                      *
202 %************************************************************************
203
204
205 \begin{code}
206 ruleRhsFreeVars :: CoreRule -> VarSet
207 ruleRhsFreeVars (BuiltinRule _ _) = noFVs
208 ruleRhsFreeVars (Rule str _ tpl_vars tpl_args rhs)
209   = rule_fvs isLocalVar emptyVarSet
210   where
211     rule_fvs = addBndrs tpl_vars (expr_fvs rhs)
212
213 ruleLhsFreeIds :: CoreRule -> VarSet
214 -- This finds all the free Ids on the LHS of the rule
215 -- *including* imported ids
216 ruleLhsFreeIds (BuiltinRule _ _) = noFVs
217 ruleLhsFreeIds (Rule _ _ tpl_vars tpl_args rhs)
218   = foldl delVarSet (exprsSomeFreeVars isId tpl_args) tpl_vars
219 \end{code}
220
221
222 %************************************************************************
223 %*                                                                      *
224 \section[freevars-everywhere]{Attaching free variables to every sub-expression}
225 %*                                                                      *
226 %************************************************************************
227
228 The free variable pass annotates every node in the expression with its
229 NON-GLOBAL free variables and type variables.
230
231 \begin{code}
232 type CoreBindWithFVs = AnnBind Id VarSet
233 type CoreExprWithFVs = AnnExpr Id VarSet
234         -- Every node annotated with its free variables,
235         -- both Ids and TyVars
236
237 freeVarsOf :: CoreExprWithFVs -> IdSet
238 freeVarsOf (free_vars, _) = free_vars
239
240 noFVs    = emptyVarSet
241 aFreeVar = unitVarSet
242 unionFVs = unionVarSet
243
244 delBindersFV :: [Var] -> VarSet -> VarSet
245 delBindersFV bs fvs = foldr delBinderFV fvs bs
246
247 delBinderFV :: Var -> VarSet -> VarSet
248 -- This way round, so we can do it multiple times using foldr
249
250 -- (b `delBinderFV` s) removes the binder b from the free variable set s,
251 -- but *adds* to s
252 --      (a) the free variables of b's type
253 --      (b) the idSpecVars of b
254 --
255 -- This is really important for some lambdas:
256 --      In (\x::a -> x) the only mention of "a" is in the binder.
257 --
258 -- Also in
259 --      let x::a = b in ...
260 -- we should really note that "a" is free in this expression.
261 -- It'll be pinned inside the /\a by the binding for b, but
262 -- it seems cleaner to make sure that a is in the free-var set 
263 -- when it is mentioned.
264 --
265 -- This also shows up in recursive bindings.  Consider:
266 --      /\a -> letrec x::a = x in E
267 -- Now, there are no explicit free type variables in the RHS of x,
268 -- but nevertheless "a" is free in its definition.  So we add in
269 -- the free tyvars of the types of the binders, and include these in the
270 -- free vars of the group, attached to the top level of each RHS.
271 --
272 -- This actually happened in the defn of errorIO in IOBase.lhs:
273 --      errorIO (ST io) = case (errorIO# io) of
274 --                          _ -> bottom
275 --                        where
276 --                          bottom = bottom -- Never evaluated
277
278 delBinderFV b s | isId b    = (s `delVarSet` b) `unionFVs` idFreeVars b
279                 | otherwise = s `delVarSet` b
280
281 idFreeVars :: Id -> VarSet
282 idFreeVars id = ASSERT( isId id) idRuleVars id `unionVarSet` idFreeTyVars id
283
284 idFreeTyVars :: Id -> TyVarSet
285 -- Only local Ids conjured up locally, can have free type variables.
286 -- (During type checking top-level Ids can have free tyvars)
287 idFreeTyVars id = tyVarsOfType (idType id)
288 -- | isLocalId id = tyVarsOfType (idType id)
289 -- | otherwise    = emptyVarSet
290
291 idRuleVars ::Id -> VarSet
292 idRuleVars id = ASSERT( isId id) rulesRhsFreeVars (idSpecialisation id)
293 \end{code}
294
295
296 %************************************************************************
297 %*                                                                      *
298 \subsection{Free variables (and types)}
299 %*                                                                      *
300 %************************************************************************
301
302 \begin{code}
303 freeVars :: CoreExpr -> CoreExprWithFVs
304
305 freeVars (Var v)
306   = (fvs, AnnVar v)
307   where
308         -- ToDo: insert motivating example for why we *need*
309         -- to include the idSpecVars in the FV list.
310         --      Actually [June 98] I don't think it's necessary
311         -- fvs = fvs_v `unionVarSet` idSpecVars v
312
313     fvs | isLocalVar v = aFreeVar v
314         | otherwise    = noFVs
315
316 freeVars (Lit lit) = (noFVs, AnnLit lit)
317 freeVars (Lam b body)
318   = (b `delBinderFV` freeVarsOf body', AnnLam b body')
319   where
320     body' = freeVars body
321
322 freeVars (App fun arg)
323   = (freeVarsOf fun2 `unionFVs` freeVarsOf arg2, AnnApp fun2 arg2)
324   where
325     fun2 = freeVars fun
326     arg2 = freeVars arg
327
328 freeVars (Case scrut bndr ty alts)
329 -- gaw 2004
330   = ((bndr `delBinderFV` alts_fvs) `unionFVs` freeVarsOf scrut2 `unionFVs` tyVarsOfType ty,
331      AnnCase scrut2 bndr ty alts2)
332   where
333     scrut2 = freeVars scrut
334
335     (alts_fvs_s, alts2) = mapAndUnzip fv_alt alts
336     alts_fvs            = foldr1 unionFVs alts_fvs_s
337
338     fv_alt (con,args,rhs) = (delBindersFV args (freeVarsOf rhs2),
339                              (con, args, rhs2))
340                           where
341                              rhs2 = freeVars rhs
342
343 freeVars (Let (NonRec binder rhs) body)
344   = (freeVarsOf rhs2 `unionFVs` body_fvs,
345      AnnLet (AnnNonRec binder rhs2) body2)
346   where
347     rhs2     = freeVars rhs
348     body2    = freeVars body
349     body_fvs = binder `delBinderFV` freeVarsOf body2
350
351 freeVars (Let (Rec binds) body)
352   = (foldl delVarSet group_fvs binders,
353         -- The "delBinderFV" part may have added one of the binders
354         -- via the idSpecVars part, so we must delete it again
355      AnnLet (AnnRec (binders `zip` rhss2)) body2)
356   where
357     (binders, rhss) = unzip binds
358
359     rhss2     = map freeVars rhss
360     all_fvs   = foldr (unionFVs . fst) body_fvs rhss2
361     group_fvs = delBindersFV binders all_fvs
362
363     body2     = freeVars body
364     body_fvs  = freeVarsOf body2
365
366 freeVars (Note (Coerce to_ty from_ty) expr)
367   = (freeVarsOf expr2 `unionFVs` tfvs1 `unionFVs` tfvs2,
368      AnnNote (Coerce to_ty from_ty) expr2)
369   where
370     expr2  = freeVars expr
371     tfvs1  = tyVarsOfType from_ty
372     tfvs2  = tyVarsOfType to_ty
373
374 freeVars (Note other_note expr)
375   = (freeVarsOf expr2, AnnNote other_note expr2)
376   where
377     expr2 = freeVars expr
378
379 freeVars (Type ty) = (tyVarsOfType ty, AnnType ty)
380 \end{code}
381