Add kind equalities to GHC.
[ghc.git] / compiler / coreSyn / CoreFVs.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4
5 Taken quite directly from the Peyton Jones/Lester paper.
6 -}
7
8 {-# LANGUAGE CPP #-}
9
10 -- | A module concerned with finding the free variables of an expression.
11 module CoreFVs (
12 -- * Free variables of expressions and binding groups
13 exprFreeVars,
14 exprFreeVarsDSet,
15 exprFreeIds,
16 exprsFreeVars,
17 exprsFreeVarsList,
18 bindFreeVars,
19
20 -- * Selective free variables of expressions
21 InterestingVarFun,
22 exprSomeFreeVars, exprsSomeFreeVars,
23
24 -- * Free variables of Rules, Vars and Ids
25 varTypeTyCoVars,
26 varTypeTyCoVarsAcc,
27 idUnfoldingVars, idFreeVars, dIdFreeVars,
28 idRuleAndUnfoldingVars, idRuleAndUnfoldingVarsDSet,
29 idFreeVarsAcc,
30 idRuleVars, idRuleRhsVars, stableUnfoldingVars,
31 ruleRhsFreeVars, ruleFreeVars, rulesFreeVars,
32 rulesFreeVarsDSet,
33 ruleLhsFreeIds, exprsOrphNames,
34 vectsFreeVars,
35
36 expr_fvs,
37
38 -- * Core syntax tree annotation with free variables
39 FVAnn, -- annotation, abstract
40 CoreExprWithFVs, -- = AnnExpr Id FVAnn
41 CoreExprWithFVs', -- = AnnExpr' Id FVAnn
42 CoreBindWithFVs, -- = AnnBind Id FVAnn
43 CoreAltWithFVs, -- = AnnAlt Id FVAnn
44 freeVars, -- CoreExpr -> CoreExprWithFVs
45 freeVarsOf, -- CoreExprWithFVs -> DIdSet
46 freeVarsOfType, -- CoreExprWithFVs -> TyCoVarSet
47 freeVarsOfAnn, freeVarsOfTypeAnn,
48 exprTypeFV -- CoreExprWithFVs -> Type
49 ) where
50
51 #include "HsVersions.h"
52
53 import CoreSyn
54 import Id
55 import IdInfo
56 import NameSet
57 import UniqFM
58 import Literal ( literalType )
59 import Name
60 import VarSet
61 import Var
62 import TcType
63 import Type
64 import Coercion
65 import Maybes( orElse )
66 import Util
67 import BasicTypes( Activation )
68 import Outputable
69 import FV
70
71 {-
72 ************************************************************************
73 * *
74 \section{Finding the free variables of an expression}
75 * *
76 ************************************************************************
77
78 This function simply finds the free variables of an expression.
79 So far as type variables are concerned, it only finds tyvars that are
80
81 * free in type arguments,
82 * free in the type of a binder,
83
84 but not those that are free in the type of variable occurrence.
85 -}
86
87 -- | Find all locally-defined free Ids or type variables in an expression
88 -- returning a non-deterministic set.
89 exprFreeVars :: CoreExpr -> VarSet
90 exprFreeVars = runFVSet . exprFreeVarsAcc
91
92 -- | Find all locally-defined free Ids or type variables in an expression
93 -- returning a composable FV computation. See Note [FV naming coventions] in FV
94 -- for why export it.
95 exprFreeVarsAcc :: CoreExpr -> FV
96 exprFreeVarsAcc = filterFV isLocalVar . expr_fvs
97
98 -- | Find all locally-defined free Ids or type variables in an expression
99 -- returning a deterministic set.
100 exprFreeVarsDSet :: CoreExpr -> DVarSet
101 exprFreeVarsDSet = runFVDSet . exprFreeVarsAcc
102
103 -- | Find all locally-defined free Ids in an expression
104 exprFreeIds :: CoreExpr -> IdSet -- Find all locally-defined free Ids
105 exprFreeIds = exprSomeFreeVars isLocalId
106
107 -- | Find all locally-defined free Ids or type variables in several expressions
108 -- returning a non-deterministic set.
109 exprsFreeVars :: [CoreExpr] -> VarSet
110 exprsFreeVars = runFVSet . exprsFreeVarsAcc
111
112 -- | Find all locally-defined free Ids or type variables in several expressions
113 -- returning a composable FV computation. See Note [FV naming coventions] in FV
114 -- for why export it.
115 exprsFreeVarsAcc :: [CoreExpr] -> FV
116 exprsFreeVarsAcc exprs = mapUnionFV exprFreeVarsAcc exprs
117
118 -- | Find all locally-defined free Ids or type variables in several expressions
119 -- returning a deterministically ordered list.
120 exprsFreeVarsList :: [CoreExpr] -> [Var]
121 exprsFreeVarsList = runFVList . exprsFreeVarsAcc
122
123 -- | Find all locally defined free Ids in a binding group
124 bindFreeVars :: CoreBind -> VarSet
125 bindFreeVars (NonRec b r) = runFVSet $ filterFV isLocalVar $ rhs_fvs (b,r)
126 bindFreeVars (Rec prs) = runFVSet $ filterFV isLocalVar $
127 addBndrs (map fst prs)
128 (mapUnionFV rhs_fvs prs)
129
130 -- | Finds free variables in an expression selected by a predicate
131 exprSomeFreeVars :: InterestingVarFun -- ^ Says which 'Var's are interesting
132 -> CoreExpr
133 -> VarSet
134 exprSomeFreeVars fv_cand e = runFVSet $ filterFV fv_cand $ expr_fvs e
135
136 -- | Finds free variables in several expressions selected by a predicate
137 exprsSomeFreeVars :: InterestingVarFun -- Says which 'Var's are interesting
138 -> [CoreExpr]
139 -> VarSet
140 exprsSomeFreeVars fv_cand es =
141 runFVSet $ filterFV fv_cand $ mapUnionFV expr_fvs es
142
143 -- Comment about obselete code
144 -- We used to gather the free variables the RULES at a variable occurrence
145 -- with the following cryptic comment:
146 -- "At a variable occurrence, add in any free variables of its rule rhss
147 -- Curiously, we gather the Id's free *type* variables from its binding
148 -- site, but its free *rule-rhs* variables from its usage sites. This
149 -- is a little weird. The reason is that the former is more efficient,
150 -- but the latter is more fine grained, and a makes a difference when
151 -- a variable mentions itself one of its own rule RHSs"
152 -- Not only is this "weird", but it's also pretty bad because it can make
153 -- a function seem more recursive than it is. Suppose
154 -- f = ...g...
155 -- g = ...
156 -- RULE g x = ...f...
157 -- Then f is not mentioned in its own RHS, and needn't be a loop breaker
158 -- (though g may be). But if we collect the rule fvs from g's occurrence,
159 -- it looks as if f mentions itself. (This bites in the eftInt/eftIntFB
160 -- code in GHC.Enum.)
161 --
162 -- Anyway, it seems plain wrong. The RULE is like an extra RHS for the
163 -- function, so its free variables belong at the definition site.
164 --
165 -- Deleted code looked like
166 -- foldVarSet add_rule_var var_itself_set (idRuleVars var)
167 -- add_rule_var var set | keep_it fv_cand in_scope var = extendVarSet set var
168 -- | otherwise = set
169 -- SLPJ Feb06
170
171 addBndr :: CoreBndr -> FV -> FV
172 addBndr bndr fv fv_cand in_scope acc
173 = (varTypeTyCoVarsAcc bndr `unionFV`
174 -- Include type varibles in the binder's type
175 -- (not just Ids; coercion variables too!)
176 FV.delFV bndr fv) fv_cand in_scope acc
177
178 addBndrs :: [CoreBndr] -> FV -> FV
179 addBndrs bndrs fv = foldr addBndr fv bndrs
180
181 expr_fvs :: CoreExpr -> FV
182 expr_fvs (Type ty) fv_cand in_scope acc =
183 tyCoVarsOfTypeAcc ty fv_cand in_scope acc
184 expr_fvs (Coercion co) fv_cand in_scope acc =
185 tyCoVarsOfCoAcc co fv_cand in_scope acc
186 expr_fvs (Var var) fv_cand in_scope acc = oneVar var fv_cand in_scope acc
187 expr_fvs (Lit _) fv_cand in_scope acc = noVars fv_cand in_scope acc
188 expr_fvs (Tick t expr) fv_cand in_scope acc =
189 (tickish_fvs t `unionFV` expr_fvs expr) fv_cand in_scope acc
190 expr_fvs (App fun arg) fv_cand in_scope acc =
191 (expr_fvs fun `unionFV` expr_fvs arg) fv_cand in_scope acc
192 expr_fvs (Lam bndr body) fv_cand in_scope acc =
193 addBndr bndr (expr_fvs body) fv_cand in_scope acc
194 expr_fvs (Cast expr co) fv_cand in_scope acc =
195 (expr_fvs expr `unionFV` tyCoVarsOfCoAcc co) fv_cand in_scope acc
196
197 expr_fvs (Case scrut bndr ty alts) fv_cand in_scope acc
198 = (expr_fvs scrut `unionFV` tyCoVarsOfTypeAcc ty `unionFV` addBndr bndr
199 (mapUnionFV alt_fvs alts)) fv_cand in_scope acc
200 where
201 alt_fvs (_, bndrs, rhs) = addBndrs bndrs (expr_fvs rhs)
202
203 expr_fvs (Let (NonRec bndr rhs) body) fv_cand in_scope acc
204 = (rhs_fvs (bndr, rhs) `unionFV` addBndr bndr (expr_fvs body))
205 fv_cand in_scope acc
206
207 expr_fvs (Let (Rec pairs) body) fv_cand in_scope acc
208 = addBndrs (map fst pairs)
209 (mapUnionFV rhs_fvs pairs `unionFV` expr_fvs body)
210 fv_cand in_scope acc
211
212 ---------
213 rhs_fvs :: (Id, CoreExpr) -> FV
214 rhs_fvs (bndr, rhs) = expr_fvs rhs `unionFV`
215 bndrRuleAndUnfoldingVarsAcc bndr
216 -- Treat any RULES as extra RHSs of the binding
217
218 ---------
219 exprs_fvs :: [CoreExpr] -> FV
220 exprs_fvs exprs = mapUnionFV expr_fvs exprs
221
222 tickish_fvs :: Tickish Id -> FV
223 tickish_fvs (Breakpoint _ ids) = someVars ids
224 tickish_fvs _ = noVars
225
226 {-
227 ************************************************************************
228 * *
229 \section{Free names}
230 * *
231 ************************************************************************
232 -}
233
234 -- | Finds the free /external/ names of an expression, notably
235 -- including the names of type constructors (which of course do not show
236 -- up in 'exprFreeVars').
237 exprOrphNames :: CoreExpr -> NameSet
238 -- There's no need to delete local binders, because they will all
239 -- be /internal/ names.
240 exprOrphNames e
241 = go e
242 where
243 go (Var v)
244 | isExternalName n = unitNameSet n
245 | otherwise = emptyNameSet
246 where n = idName v
247 go (Lit _) = emptyNameSet
248 go (Type ty) = orphNamesOfType ty -- Don't need free tyvars
249 go (Coercion co) = orphNamesOfCo co
250 go (App e1 e2) = go e1 `unionNameSet` go e2
251 go (Lam v e) = go e `delFromNameSet` idName v
252 go (Tick _ e) = go e
253 go (Cast e co) = go e `unionNameSet` orphNamesOfCo co
254 go (Let (NonRec _ r) e) = go e `unionNameSet` go r
255 go (Let (Rec prs) e) = exprsOrphNames (map snd prs) `unionNameSet` go e
256 go (Case e _ ty as) = go e `unionNameSet` orphNamesOfType ty
257 `unionNameSet` unionNameSets (map go_alt as)
258
259 go_alt (_,_,r) = go r
260
261 -- | Finds the free /external/ names of several expressions: see 'exprOrphNames' for details
262 exprsOrphNames :: [CoreExpr] -> NameSet
263 exprsOrphNames es = foldr (unionNameSet . exprOrphNames) emptyNameSet es
264
265 {-
266 ************************************************************************
267 * *
268 \section[freevars-everywhere]{Attaching free variables to every sub-expression}
269 * *
270 ************************************************************************
271 -}
272
273 -- | Those variables free in the right hand side of a rule returned as a
274 -- non-deterministic set
275 ruleRhsFreeVars :: CoreRule -> VarSet
276 ruleRhsFreeVars (BuiltinRule {}) = noFVs
277 ruleRhsFreeVars (Rule { ru_fn = _, ru_bndrs = bndrs, ru_rhs = rhs })
278 = runFVSet $ filterFV isLocalVar $ addBndrs bndrs (expr_fvs rhs)
279 -- See Note [Rule free var hack]
280
281 -- | Those variables free in the both the left right hand sides of a rule
282 -- returned as a non-deterministic set
283 ruleFreeVars :: CoreRule -> VarSet
284 ruleFreeVars = runFVSet . ruleFreeVarsAcc
285
286 -- | Those variables free in the both the left right hand sides of a rule
287 -- returned as FV computation
288 ruleFreeVarsAcc :: CoreRule -> FV
289 ruleFreeVarsAcc (BuiltinRule {}) = noVars
290 ruleFreeVarsAcc (Rule { ru_fn = _do_not_include
291 -- See Note [Rule free var hack]
292 , ru_bndrs = bndrs
293 , ru_rhs = rhs, ru_args = args })
294 = filterFV isLocalVar $ addBndrs bndrs (exprs_fvs (rhs:args))
295
296 -- | Those variables free in the both the left right hand sides of rules
297 -- returned as FV computation
298 rulesFreeVarsAcc :: [CoreRule] -> FV
299 rulesFreeVarsAcc = mapUnionFV ruleFreeVarsAcc
300
301 -- | Those variables free in the both the left right hand sides of rules
302 -- returned as a deterministic set
303 rulesFreeVarsDSet :: [CoreRule] -> DVarSet
304 rulesFreeVarsDSet rules = runFVDSet $ rulesFreeVarsAcc rules
305
306 idRuleRhsVars :: (Activation -> Bool) -> Id -> VarSet
307 -- Just the variables free on the *rhs* of a rule
308 idRuleRhsVars is_active id
309 = mapUnionVarSet get_fvs (idCoreRules id)
310 where
311 get_fvs (Rule { ru_fn = fn, ru_bndrs = bndrs
312 , ru_rhs = rhs, ru_act = act })
313 | is_active act
314 -- See Note [Finding rule RHS free vars] in OccAnal.hs
315 = delFromUFM fvs fn -- Note [Rule free var hack]
316 where
317 fvs = runFVSet $ filterFV isLocalVar $ addBndrs bndrs (expr_fvs rhs)
318 get_fvs _ = noFVs
319
320 -- | Those variables free in the right hand side of several rules
321 rulesFreeVars :: [CoreRule] -> VarSet
322 rulesFreeVars rules = mapUnionVarSet ruleFreeVars rules
323
324 ruleLhsFreeIds :: CoreRule -> VarSet
325 -- ^ This finds all locally-defined free Ids on the left hand side of a rule
326 ruleLhsFreeIds (BuiltinRule {}) = noFVs
327 ruleLhsFreeIds (Rule { ru_bndrs = bndrs, ru_args = args })
328 = runFVSet $ filterFV isLocalId $ addBndrs bndrs (exprs_fvs args)
329
330 {-
331 Note [Rule free var hack] (Not a hack any more)
332 ~~~~~~~~~~~~~~~~~~~~~~~~~
333 We used not to include the Id in its own rhs free-var set.
334 Otherwise the occurrence analyser makes bindings recursive:
335 f x y = x+y
336 RULE: f (f x y) z ==> f x (f y z)
337 However, the occurrence analyser distinguishes "non-rule loop breakers"
338 from "rule-only loop breakers" (see BasicTypes.OccInfo). So it will
339 put this 'f' in a Rec block, but will mark the binding as a non-rule loop
340 breaker, which is perfectly inlinable.
341 -}
342
343 -- |Free variables of a vectorisation declaration
344 vectsFreeVars :: [CoreVect] -> VarSet
345 vectsFreeVars = mapUnionVarSet vectFreeVars
346 where
347 vectFreeVars (Vect _ rhs) = runFVSet $ filterFV isLocalId $ expr_fvs rhs
348 vectFreeVars (NoVect _) = noFVs
349 vectFreeVars (VectType _ _ _) = noFVs
350 vectFreeVars (VectClass _) = noFVs
351 vectFreeVars (VectInst _) = noFVs
352 -- this function is only concerned with values, not types
353
354 {-
355 ************************************************************************
356 * *
357 \section[freevars-everywhere]{Attaching free variables to every sub-expression}
358 * *
359 ************************************************************************
360
361 The free variable pass annotates every node in the expression with its
362 NON-GLOBAL free variables and type variables.
363 -}
364
365 data FVAnn = FVAnn { fva_fvs :: DVarSet -- free in expression
366 , fva_ty_fvs :: DVarSet -- free only in expression's type
367 , fva_ty :: Type -- expression's type
368 }
369
370 -- | Every node in a binding group annotated with its
371 -- (non-global) free variables, both Ids and TyVars, and type.
372 type CoreBindWithFVs = AnnBind Id FVAnn
373 -- | Every node in an expression annotated with its
374 -- (non-global) free variables, both Ids and TyVars, and type.
375 type CoreExprWithFVs = AnnExpr Id FVAnn
376 type CoreExprWithFVs' = AnnExpr' Id FVAnn
377
378 -- | Every node in an expression annotated with its
379 -- (non-global) free variables, both Ids and TyVars, and type.
380 type CoreAltWithFVs = AnnAlt Id FVAnn
381
382 freeVarsOf :: CoreExprWithFVs -> DIdSet
383 -- ^ Inverse function to 'freeVars'
384 freeVarsOf (FVAnn { fva_fvs = fvs }, _) = fvs
385
386 -- | Extract the vars free in an annotated expression's type
387 freeVarsOfType :: CoreExprWithFVs -> DTyCoVarSet
388 freeVarsOfType (FVAnn { fva_ty_fvs = ty_fvs }, _) = ty_fvs
389
390 -- | Extract the type of an annotated expression. (This is cheap.)
391 exprTypeFV :: CoreExprWithFVs -> Type
392 exprTypeFV (FVAnn { fva_ty = ty }, _) = ty
393
394 -- | Extract the vars reported in a FVAnn
395 freeVarsOfAnn :: FVAnn -> DIdSet
396 freeVarsOfAnn = fva_fvs
397
398 -- | Extract the type-level vars reported in a FVAnn
399 freeVarsOfTypeAnn :: FVAnn -> DTyCoVarSet
400 freeVarsOfTypeAnn = fva_ty_fvs
401
402 noFVs :: VarSet
403 noFVs = emptyVarSet
404
405 aFreeVar :: Var -> DVarSet
406 aFreeVar = unitDVarSet
407
408 unionFVs :: DVarSet -> DVarSet -> DVarSet
409 unionFVs = unionDVarSet
410
411 unionFVss :: [DVarSet] -> DVarSet
412 unionFVss = unionDVarSets
413
414 delBindersFV :: [Var] -> DVarSet -> DVarSet
415 delBindersFV bs fvs = foldr delBinderFV fvs bs
416
417 delBinderFV :: Var -> DVarSet -> DVarSet
418 -- This way round, so we can do it multiple times using foldr
419
420 -- (b `delBinderFV` s) removes the binder b from the free variable set s,
421 -- but *adds* to s
422 --
423 -- the free variables of b's type
424 --
425 -- This is really important for some lambdas:
426 -- In (\x::a -> x) the only mention of "a" is in the binder.
427 --
428 -- Also in
429 -- let x::a = b in ...
430 -- we should really note that "a" is free in this expression.
431 -- It'll be pinned inside the /\a by the binding for b, but
432 -- it seems cleaner to make sure that a is in the free-var set
433 -- when it is mentioned.
434 --
435 -- This also shows up in recursive bindings. Consider:
436 -- /\a -> letrec x::a = x in E
437 -- Now, there are no explicit free type variables in the RHS of x,
438 -- but nevertheless "a" is free in its definition. So we add in
439 -- the free tyvars of the types of the binders, and include these in the
440 -- free vars of the group, attached to the top level of each RHS.
441 --
442 -- This actually happened in the defn of errorIO in IOBase.hs:
443 -- errorIO (ST io) = case (errorIO# io) of
444 -- _ -> bottom
445 -- where
446 -- bottom = bottom -- Never evaluated
447
448 delBinderFV b s = (s `delDVarSet` b) `unionFVs` dVarTypeTyCoVars b
449 -- Include coercion variables too!
450
451 varTypeTyCoVars :: Var -> TyCoVarSet
452 -- Find the type/kind variables free in the type of the id/tyvar
453 varTypeTyCoVars var = runFVSet $ varTypeTyCoVarsAcc var
454
455 dVarTypeTyCoVars :: Var -> DTyCoVarSet
456 -- Find the type/kind/coercion variables free in the type of the id/tyvar
457 dVarTypeTyCoVars var = runFVDSet $ varTypeTyCoVarsAcc var
458
459 varTypeTyCoVarsAcc :: Var -> FV
460 varTypeTyCoVarsAcc var = tyCoVarsOfTypeAcc (varType var)
461
462 idFreeVars :: Id -> VarSet
463 idFreeVars id = ASSERT( isId id) runFVSet $ idFreeVarsAcc id
464
465 dIdFreeVars :: Id -> DVarSet
466 dIdFreeVars id = runFVDSet $ idFreeVarsAcc id
467
468 idFreeVarsAcc :: Id -> FV
469 -- Type variables, rule variables, and inline variables
470 idFreeVarsAcc id = ASSERT( isId id)
471 varTypeTyCoVarsAcc id `unionFV`
472 idRuleAndUnfoldingVarsAcc id
473
474 bndrRuleAndUnfoldingVarsAcc :: Var -> FV
475 bndrRuleAndUnfoldingVarsAcc v | isTyVar v = noVars
476 | otherwise = idRuleAndUnfoldingVarsAcc v
477
478 idRuleAndUnfoldingVars :: Id -> VarSet
479 idRuleAndUnfoldingVars id = runFVSet $ idRuleAndUnfoldingVarsAcc id
480
481 idRuleAndUnfoldingVarsDSet :: Id -> DVarSet
482 idRuleAndUnfoldingVarsDSet id = runFVDSet $ idRuleAndUnfoldingVarsAcc id
483
484 idRuleAndUnfoldingVarsAcc :: Id -> FV
485 idRuleAndUnfoldingVarsAcc id = ASSERT( isId id)
486 idRuleVarsAcc id `unionFV` idUnfoldingVarsAcc id
487
488
489 idRuleVars ::Id -> VarSet -- Does *not* include CoreUnfolding vars
490 idRuleVars id = runFVSet $ idRuleVarsAcc id
491
492 idRuleVarsAcc :: Id -> FV
493 idRuleVarsAcc id = ASSERT( isId id)
494 someVars (dVarSetElems $ ruleInfoFreeVars (idSpecialisation id))
495
496 idUnfoldingVars :: Id -> VarSet
497 -- Produce free vars for an unfolding, but NOT for an ordinary
498 -- (non-inline) unfolding, since it is a dup of the rhs
499 -- and we'll get exponential behaviour if we look at both unf and rhs!
500 -- But do look at the *real* unfolding, even for loop breakers, else
501 -- we might get out-of-scope variables
502 idUnfoldingVars id = runFVSet $ idUnfoldingVarsAcc id
503
504 idUnfoldingVarsAcc :: Id -> FV
505 idUnfoldingVarsAcc id = stableUnfoldingVarsAcc (realIdUnfolding id) `orElse` noVars
506
507 stableUnfoldingVars :: Unfolding -> Maybe VarSet
508 stableUnfoldingVars unf = runFVSet `fmap` stableUnfoldingVarsAcc unf
509
510 stableUnfoldingVarsAcc :: Unfolding -> Maybe FV
511 stableUnfoldingVarsAcc unf
512 = case unf of
513 CoreUnfolding { uf_tmpl = rhs, uf_src = src }
514 | isStableSource src
515 -> Just (filterFV isLocalVar $ expr_fvs rhs)
516 DFunUnfolding { df_bndrs = bndrs, df_args = args }
517 -> Just (filterFV isLocalVar $ FV.delFVs (mkVarSet bndrs) $ exprs_fvs args)
518 -- DFuns are top level, so no fvs from types of bndrs
519 _other -> Nothing
520
521
522 {-
523 ************************************************************************
524 * *
525 \subsection{Free variables (and types)}
526 * *
527 ************************************************************************
528 -}
529
530 freeVars :: CoreExpr -> CoreExprWithFVs
531 -- ^ Annotate a 'CoreExpr' with its (non-global) free type and value variables at every tree node
532 freeVars = go
533 where
534 go :: CoreExpr -> CoreExprWithFVs
535 go (Var v)
536 = (FVAnn fvs ty_fvs (idType v), AnnVar v)
537 where
538 -- ToDo: insert motivating example for why we *need*
539 -- to include the idSpecVars in the FV list.
540 -- Actually [June 98] I don't think it's necessary
541 -- fvs = fvs_v `unionVarSet` idSpecVars v
542
543 (fvs, ty_fvs)
544 | isLocalVar v = (aFreeVar v `unionFVs` ty_fvs, dVarTypeTyCoVars v)
545 | otherwise = (emptyDVarSet, emptyDVarSet)
546
547 go (Lit lit) = (FVAnn emptyDVarSet emptyDVarSet (literalType lit), AnnLit lit)
548 go (Lam b body)
549 = ( FVAnn { fva_fvs = b_fvs `unionFVs` (b `delBinderFV` body_fvs)
550 , fva_ty_fvs = b_fvs `unionFVs` (b `delBinderFV` body_ty_fvs)
551 , fva_ty = mkFunTy b_ty body_ty }
552 , AnnLam b body' )
553 where
554 body'@(FVAnn { fva_fvs = body_fvs, fva_ty_fvs = body_ty_fvs
555 , fva_ty = body_ty }, _) = go body
556 b_ty = idType b
557 b_fvs = tyCoVarsOfTypeDSet b_ty
558
559 go (App fun arg)
560 = ( FVAnn { fva_fvs = freeVarsOf fun' `unionFVs` freeVarsOf arg'
561 , fva_ty_fvs = tyCoVarsOfTypeDSet res_ty
562 , fva_ty = res_ty }
563 , AnnApp fun' arg' )
564 where
565 fun' = go fun
566 fun_ty = exprTypeFV fun'
567 arg' = go arg
568 res_ty = applyTypeToArg fun_ty arg
569
570 go (Case scrut bndr ty alts)
571 = ( FVAnn { fva_fvs = (bndr `delBinderFV` alts_fvs)
572 `unionFVs` freeVarsOf scrut2
573 `unionFVs` tyCoVarsOfTypeDSet ty
574 -- don't need to look at (idType bndr)
575 -- b/c that's redundant with scrut
576 , fva_ty_fvs = tyCoVarsOfTypeDSet ty
577 , fva_ty = ty }
578 , AnnCase scrut2 bndr ty alts2 )
579 where
580 scrut2 = go scrut
581
582 (alts_fvs_s, alts2) = mapAndUnzip fv_alt alts
583 alts_fvs = unionFVss alts_fvs_s
584
585 fv_alt (con,args,rhs) = (delBindersFV args (freeVarsOf rhs2),
586 (con, args, rhs2))
587 where
588 rhs2 = go rhs
589
590 go (Let (NonRec binder rhs) body)
591 = ( FVAnn { fva_fvs = freeVarsOf rhs2
592 `unionFVs` body_fvs
593 `unionFVs` runFVDSet
594 (bndrRuleAndUnfoldingVarsAcc binder)
595 -- Remember any rules; cf rhs_fvs above
596 , fva_ty_fvs = freeVarsOfType body2
597 , fva_ty = exprTypeFV body2 }
598 , AnnLet (AnnNonRec binder rhs2) body2 )
599 where
600 rhs2 = go rhs
601 body2 = go body
602 body_fvs = binder `delBinderFV` freeVarsOf body2
603
604 go (Let (Rec binds) body)
605 = ( FVAnn { fva_fvs = delBindersFV binders all_fvs
606 , fva_ty_fvs = freeVarsOfType body2
607 , fva_ty = exprTypeFV body2 }
608 , AnnLet (AnnRec (binders `zip` rhss2)) body2 )
609 where
610 (binders, rhss) = unzip binds
611
612 rhss2 = map go rhss
613 rhs_body_fvs = foldr (unionFVs . freeVarsOf) body_fvs rhss2
614 binders_fvs = runFVDSet $ mapUnionFV idRuleAndUnfoldingVarsAcc binders
615 all_fvs = rhs_body_fvs `unionFVs` binders_fvs
616 -- The "delBinderFV" happens after adding the idSpecVars,
617 -- since the latter may add some of the binders as fvs
618
619 body2 = go body
620 body_fvs = freeVarsOf body2
621
622 go (Cast expr co)
623 = ( FVAnn (freeVarsOf expr2 `unionFVs` cfvs) (tyCoVarsOfTypeDSet to_ty) to_ty
624 , AnnCast expr2 (c_ann, co) )
625 where
626 expr2 = go expr
627 cfvs = tyCoVarsOfCoDSet co
628 c_ann = FVAnn cfvs (tyCoVarsOfTypeDSet co_ki) co_ki
629 co_ki = coercionType co
630 Just (_, to_ty) = splitCoercionType_maybe co_ki
631
632
633 go (Tick tickish expr)
634 = ( FVAnn { fva_fvs = tickishFVs tickish `unionFVs` freeVarsOf expr2
635 , fva_ty_fvs = freeVarsOfType expr2
636 , fva_ty = exprTypeFV expr2 }
637 , AnnTick tickish expr2 )
638 where
639 expr2 = go expr
640 tickishFVs (Breakpoint _ ids) = mkDVarSet ids
641 tickishFVs _ = emptyDVarSet
642
643 go (Type ty) = ( FVAnn (tyCoVarsOfTypeDSet ty)
644 (tyCoVarsOfTypeDSet ki)
645 ki
646 , AnnType ty)
647 where
648 ki = typeKind ty
649
650 go (Coercion co) = ( FVAnn (tyCoVarsOfCoDSet co)
651 (tyCoVarsOfTypeDSet ki)
652 ki
653 , AnnCoercion co)
654 where
655 ki = coercionType co