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, exprsFreeVars,
   9         exprSomeFreeVars, exprsSomeFreeVars,
  10         idRuleVars, idFreeVars, ruleSomeFreeVars, ruleSomeLhsFreeVars,
  11
  12         CoreExprWithFVs, CoreBindWithFVs, freeVars, freeVarsOf,
  13     ) where
  14
  15 #include "HsVersions.h"
  16
  17 import CoreSyn
  18 import Id               ( Id, idFreeTyVars, getIdSpecialisation )
  19 import VarSet
  20 import Var              ( IdOrTyVar, isId )
  21 import Name             ( isLocallyDefined )
  22 import Type             ( tyVarsOfType, Type )
  23 import Util             ( mapAndUnzip )
  24 \end{code}
  25
  26 %************************************************************************
  27 %*                                                                      *
  28 \section{Finding the free variables of an expression}
  29 %*                                                                      *
  30 %************************************************************************
  31
  32 This function simply finds the free variables of an expression.
  33 So far as type variables are concerned, it only finds tyvars that are
  34
  35         * free in type arguments,
  36         * free in the type of a binder,
  37
  38 but not those that are free in the type of variable occurrence.
  39
  40 \begin{code}
  41 exprFreeVars :: CoreExpr -> IdOrTyVarSet        -- Find all locally-defined free Ids or tyvars
  42 exprFreeVars = exprSomeFreeVars isLocallyDefined
  43
  44 exprsFreeVars :: [CoreExpr] -> IdOrTyVarSet
  45 exprsFreeVars = foldr (unionVarSet . exprFreeVars) emptyVarSet
  46
  47 exprSomeFreeVars :: InterestingVarFun   -- Says which Vars are interesting
  48                  -> CoreExpr
  49                  -> IdOrTyVarSet
  50 exprSomeFreeVars fv_cand e = expr_fvs e fv_cand emptyVarSet
  51
  52 exprsSomeFreeVars :: InterestingVarFun  -- Says which Vars are interesting
  53                   -> [CoreExpr]
  54                   -> IdOrTyVarSet
  55 exprsSomeFreeVars fv_cand = foldr (unionVarSet . exprSomeFreeVars fv_cand) emptyVarSet
  56
  57 type InterestingVarFun = IdOrTyVar -> Bool      -- True <=> interesting
  58 \end{code}
  59
  60
  61 \begin{code}
  62 type FV = InterestingVarFun
  63           -> IdOrTyVarSet       -- In scope
  64           -> IdOrTyVarSet       -- Free vars
  65
  66 union :: FV -> FV -> FV
  67 union fv1 fv2 fv_cand in_scope = fv1 fv_cand in_scope `unionVarSet` fv2 fv_cand in_scope
  68
  69 noVars :: FV
  70 noVars fv_cand in_scope = emptyVarSet
  71
  72 -- At a variable occurrence, add in any free variables of its rule rhss
  73 -- Curiously, we gather the Id's free *type* variables from its binding
  74 -- site, but its free *rule-rhs* variables from its usage sites.  This
  75 -- is a little weird.  The reason is that the former is more efficient,
  76 -- but the latter is more fine grained, and a makes a difference when
  77 -- a variable mentions itself one of its own rule RHSs
  78 oneVar :: IdOrTyVar -> FV
  79 oneVar var fv_cand in_scope
  80   = foldVarSet add_rule_var var_itself_set (idRuleVars var)
  81   where
  82     var_itself_set | keep_it fv_cand in_scope var = unitVarSet var
  83                    | otherwise                = emptyVarSet
  84     add_rule_var var set | keep_it fv_cand in_scope var = extendVarSet set var
  85                          | otherwise                    = set
  86
  87 someVars :: IdOrTyVarSet -> FV
  88 someVars vars fv_cand in_scope
  89   = filterVarSet (keep_it fv_cand in_scope) vars
  90
  91 keep_it fv_cand in_scope var
  92   | var `elemVarSet` in_scope = False
  93   | fv_cand var               = True
  94   | otherwise                 = False
  95
  96
  97 addBndr :: CoreBndr -> FV -> FV
  98 addBndr bndr fv fv_cand in_scope
  99   | isId bndr = inside_fvs `unionVarSet` someVars (idFreeTyVars bndr) fv_cand in_scope
 100   | otherwise = inside_fvs
 101   where
 102     inside_fvs = fv fv_cand (in_scope `extendVarSet` bndr)
 103
 104 addBndrs :: [CoreBndr] -> FV -> FV
 105 addBndrs bndrs fv = foldr addBndr fv bndrs
 106 \end{code}
 107
 108
 109 \begin{code}
 110 expr_fvs :: CoreExpr -> FV
 111
 112 expr_fvs (Type ty)       = someVars (tyVarsOfType ty)
 113 expr_fvs (Var var)       = oneVar var
 114 expr_fvs (Con con args)  = foldr (union . expr_fvs) noVars args
 115 expr_fvs (Note _ expr)   = expr_fvs expr
 116 expr_fvs (App fun arg)   = expr_fvs fun `union` expr_fvs arg
 117 expr_fvs (Lam bndr body) = addBndr bndr (expr_fvs body)
 118
 119 expr_fvs (Case scrut bndr alts)
 120   = expr_fvs scrut `union` addBndr bndr (foldr (union . alt_fvs) noVars alts)
 121   where
 122     alt_fvs (con, bndrs, rhs) = addBndrs bndrs (expr_fvs rhs)
 123
 124 expr_fvs (Let (NonRec bndr rhs) body)
 125   = expr_fvs rhs `union` addBndr bndr (expr_fvs body)
 126
 127 expr_fvs (Let (Rec pairs) body)
 128   = addBndrs bndrs (foldr (union . expr_fvs) (expr_fvs body) rhss)
 129   where
 130     (bndrs,rhss) = unzip pairs
 131 \end{code}
 132
 133
 134
 135 \begin{code}
 136 idRuleVars ::Id -> IdOrTyVarSet
 137 idRuleVars id = rulesRhsFreeVars (getIdSpecialisation id)
 138
 139 idFreeVars :: Id -> IdOrTyVarSet
 140 idFreeVars id = idRuleVars id `unionVarSet` idFreeTyVars id
 141
 142 rulesSomeFreeVars :: InterestingVarFun -> CoreRules -> IdOrTyVarSet
 143 rulesSomeFreeVars interesting (Rules rules _)
 144   = foldr (unionVarSet . ruleSomeFreeVars interesting) emptyVarSet rules
 145
 146 ruleSomeFreeVars :: InterestingVarFun -> CoreRule -> IdOrTyVarSet
 147 ruleSomeFreeVars interesting (BuiltinRule _) = noFVs
 148 ruleSomeFreeVars interesting (Rule _ tpl_vars tpl_args rhs)
 149   = rule_fvs interesting emptyVarSet
 150   where
 151     rule_fvs = addBndrs tpl_vars $
 152                foldr (union . expr_fvs) (expr_fvs rhs) tpl_args
 153
 154 ruleSomeLhsFreeVars :: InterestingVarFun -> CoreRule -> IdOrTyVarSet
 155 ruleSomeLhsFreeVars fn (BuiltinRule _) = noFVs
 156 ruleSomeLhsFreeVars fn (Rule _ tpl_vars tpl_args rhs)
 157   = foldl delVarSet (exprsSomeFreeVars fn tpl_args) tpl_vars
 158 \end{code}
 159
 160
 161 %************************************************************************
 162 %*                                                                      *
 163 \section[freevars-everywhere]{Attaching free variables to every sub-expression}
 164 %*                                                                      *
 165 %************************************************************************
 166
 167 The free variable pass annotates every node in the expression with its
 168 NON-GLOBAL free variables and type variables.
 169
 170 \begin{code}
 171 type CoreBindWithFVs = AnnBind Id IdOrTyVarSet
 172 type CoreExprWithFVs = AnnExpr Id IdOrTyVarSet
 173         -- Every node annotated with its free variables,
 174         -- both Ids and TyVars
 175
 176 freeVarsOf :: CoreExprWithFVs -> IdSet
 177 freeVarsOf (free_vars, _) = free_vars
 178
 179 noFVs    = emptyVarSet
 180 aFreeVar = unitVarSet
 181 unionFVs = unionVarSet
 182
 183 filters :: IdOrTyVar -> IdOrTyVarSet -> IdOrTyVarSet
 184
 185 -- (b `filters` s) removes the binder b from the free variable set s,
 186 -- but *adds* to s
 187 --      (a) the free variables of b's type
 188 --      (b) the idSpecVars of b
 189 --
 190 -- This is really important for some lambdas:
 191 --      In (\x::a -> x) the only mention of "a" is in the binder.
 192 --
 193 -- Also in
 194 --      let x::a = b in ...
 195 -- we should really note that "a" is free in this expression.
 196 -- It'll be pinned inside the /\a by the binding for b, but
 197 -- it seems cleaner to make sure that a is in the free-var set
 198 -- when it is mentioned.
 199 --
 200 -- This also shows up in recursive bindings.  Consider:
 201 --      /\a -> letrec x::a = x in E
 202 -- Now, there are no explicit free type variables in the RHS of x,
 203 -- but nevertheless "a" is free in its definition.  So we add in
 204 -- the free tyvars of the types of the binders, and include these in the
 205 -- free vars of the group, attached to the top level of each RHS.
 206 --
 207 -- This actually happened in the defn of errorIO in IOBase.lhs:
 208 --      errorIO (ST io) = case (errorIO# io) of
 209 --                          _ -> bottom
 210 --                        where
 211 --                          bottom = bottom -- Never evaluated
 212
 213 filters b s | isId b    = (s `delVarSet` b) `unionFVs` idFreeVars b
 214             | otherwise = s `delVarSet` b
 215 \end{code}
 216
 217
 218 %************************************************************************
 219 %*                                                                      *
 220 \subsection{Free variables (and types)}
 221 %*                                                                      *
 222 %************************************************************************
 223
 224 \begin{code}
 225 freeVars :: CoreExpr -> CoreExprWithFVs
 226
 227 freeVars (Var v)
 228   = (fvs, AnnVar v)
 229   where
 230         -- ToDo: insert motivating example for why we *need*
 231         -- to include the idSpecVars in the FV list.
 232         --      Actually [June 98] I don't think it's necessary
 233         -- fvs = fvs_v `unionVarSet` idSpecVars v
 234
 235     fvs | isLocallyDefined v = aFreeVar v
 236         | otherwise          = noFVs
 237
 238 freeVars (Con con args)
 239   = (foldr (unionFVs . freeVarsOf) noFVs args2, AnnCon con args2)
 240   where
 241     args2 = map freeVars args
 242
 243 freeVars (Lam b body)
 244   = (b `filters` freeVarsOf body', AnnLam b body')
 245   where
 246     body' = freeVars body
 247
 248 freeVars (App fun arg)
 249   = (freeVarsOf fun2 `unionFVs` freeVarsOf arg2, AnnApp fun2 arg2)
 250   where
 251     fun2 = freeVars fun
 252     arg2 = freeVars arg
 253
 254 freeVars (Case scrut bndr alts)
 255   = ((bndr `filters` alts_fvs) `unionFVs` freeVarsOf scrut2,
 256      AnnCase scrut2 bndr alts2)
 257   where
 258     scrut2 = freeVars scrut
 259
 260     (alts_fvs_s, alts2) = mapAndUnzip fv_alt alts
 261     alts_fvs            = foldr1 unionFVs alts_fvs_s
 262
 263     fv_alt (con,args,rhs) = (foldr filters (freeVarsOf rhs2) args,
 264                              (con, args, rhs2))
 265                           where
 266                              rhs2 = freeVars rhs
 267
 268 freeVars (Let (NonRec binder rhs) body)
 269   = (freeVarsOf rhs2 `unionFVs` body_fvs,
 270      AnnLet (AnnNonRec binder rhs2) body2)
 271   where
 272     rhs2     = freeVars rhs
 273     body2    = freeVars body
 274     body_fvs = binder `filters` freeVarsOf body2
 275
 276 freeVars (Let (Rec binds) body)
 277   = (foldl delVarSet group_fvs binders,
 278         -- The "filters" part may have added one of the binders
 279         -- via the idSpecVars part, so we must delete it again
 280      AnnLet (AnnRec (binders `zip` rhss2)) body2)
 281   where
 282     (binders, rhss) = unzip binds
 283
 284     rhss2     = map freeVars rhss
 285     all_fvs   = foldr (unionFVs . fst) body_fvs rhss2
 286     group_fvs = foldr filters all_fvs binders
 287
 288     body2     = freeVars body
 289     body_fvs  = freeVarsOf body2
 290
 291 freeVars (Note (Coerce to_ty from_ty) expr)
 292   = (freeVarsOf expr2 `unionFVs` tfvs1 `unionFVs` tfvs2,
 293      AnnNote (Coerce to_ty from_ty) expr2)
 294   where
 295     expr2  = freeVars expr
 296     tfvs1  = tyVarsOfType from_ty
 297     tfvs2  = tyVarsOfType to_ty
 298
 299 freeVars (Note other_note expr)
 300   = (freeVarsOf expr2, AnnNote other_note expr2)
 301   where
 302     expr2 = freeVars expr
 303
 304 freeVars (Type ty) = (tyVarsOfType ty, AnnType ty)
 305 \end{code}
 306