80ce33ff7a8bb7aa3c83f91b566782722a4cb3be
[ghc.git] / compiler / stgSyn / StgFVs.hs
1 -- | Free variable analysis on STG terms.
2 module StgFVs (
3 annTopBindingsFreeVars
4 ) where
5
6 import GhcPrelude
7
8 import StgSyn
9 import Id
10 import VarSet
11 import CoreSyn ( Tickish(Breakpoint) )
12 import Outputable
13 import Util
14
15 import Data.Maybe ( mapMaybe )
16
17 newtype Env
18 = Env
19 { locals :: IdSet
20 }
21
22 emptyEnv :: Env
23 emptyEnv = Env emptyVarSet
24
25 addLocals :: [Id] -> Env -> Env
26 addLocals bndrs env
27 = env { locals = extendVarSetList (locals env) bndrs }
28
29 -- | Annotates a top-level STG binding with its free variables.
30 annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
31 annTopBindingsFreeVars = map go
32 where
33 go (StgTopStringLit id bs) = StgTopStringLit id bs
34 go (StgTopLifted bind)
35 = StgTopLifted (fst (binding emptyEnv emptyVarSet bind))
36
37 boundIds :: StgBinding -> [Id]
38 boundIds (StgNonRec b _) = [b]
39 boundIds (StgRec pairs) = map fst pairs
40
41 -- Note [Tracking local binders]
42 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
43 -- 'locals' contains non-toplevel, non-imported binders.
44 -- We maintain the set in 'expr', 'alt' and 'rhs', which are the only
45 -- places where new local binders are introduced.
46 -- Why do it there rather than in 'binding'? Two reasons:
47 --
48 -- 1. We call 'binding' from 'annTopBindingsFreeVars', which would
49 -- add top-level bindings to the 'locals' set.
50 -- 2. In the let(-no-escape) case, we need to extend the environment
51 -- prior to analysing the body, but we also need the fvs from the
52 -- body to analyse the RHSs. No way to do this without some
53 -- knot-tying.
54
55 -- | This makes sure that only local, non-global free vars make it into the set.
56 mkFreeVarSet :: Env -> [Id] -> IdSet
57 mkFreeVarSet env = mkVarSet . filter (`elemVarSet` locals env)
58
59 args :: Env -> [StgArg] -> IdSet
60 args env = mkFreeVarSet env . mapMaybe f
61 where
62 f (StgVarArg occ) = Just occ
63 f _ = Nothing
64
65 binding :: Env -> IdSet -> StgBinding -> (CgStgBinding, IdSet)
66 binding env body_fv (StgNonRec bndr r) = (StgNonRec bndr r', fvs)
67 where
68 -- See Note [Tacking local binders]
69 (r', rhs_fvs) = rhs env r
70 fvs = delVarSet body_fv bndr `unionVarSet` rhs_fvs
71 binding env body_fv (StgRec pairs) = (StgRec pairs', fvs)
72 where
73 -- See Note [Tacking local binders]
74 bndrs = map fst pairs
75 (rhss, rhs_fvss) = mapAndUnzip (rhs env . snd) pairs
76 pairs' = zip bndrs rhss
77 fvs = delVarSetList (unionVarSets (body_fv:rhs_fvss)) bndrs
78
79 expr :: Env -> StgExpr -> (CgStgExpr, IdSet)
80 expr env = go
81 where
82 go (StgApp occ as)
83 = (StgApp occ as, unionVarSet (args env as) (mkFreeVarSet env [occ]))
84 go (StgLit lit) = (StgLit lit, emptyVarSet)
85 go (StgConApp dc as tys) = (StgConApp dc as tys, args env as)
86 go (StgOpApp op as ty) = (StgOpApp op as ty, args env as)
87 go StgLam{} = pprPanic "StgFVs: StgLam" empty
88 go (StgCase scrut bndr ty alts) = (StgCase scrut' bndr ty alts', fvs)
89 where
90 (scrut', scrut_fvs) = go scrut
91 -- See Note [Tacking local binders]
92 (alts', alt_fvss) = mapAndUnzip (alt (addLocals [bndr] env)) alts
93 alt_fvs = unionVarSets alt_fvss
94 fvs = delVarSet (unionVarSet scrut_fvs alt_fvs) bndr
95 go (StgLet bind body) = go_bind StgLet bind body
96 go (StgLetNoEscape bind body) = go_bind StgLetNoEscape bind body
97 go (StgTick tick e) = (StgTick tick e', fvs')
98 where
99 (e', fvs) = go e
100 fvs' = unionVarSet (tickish tick) fvs
101 tickish (Breakpoint _ ids) = mkVarSet ids
102 tickish _ = emptyVarSet
103
104 go_bind dc bind body = (dc bind' body', fvs)
105 where
106 -- See Note [Tacking local binders]
107 env' = addLocals (boundIds bind) env
108 (body', body_fvs) = expr env' body
109 (bind', fvs) = binding env' body_fvs bind
110
111 rhs :: Env -> StgRhs -> (CgStgRhs, IdSet)
112 rhs env (StgRhsClosure _ ccs uf bndrs body)
113 = (StgRhsClosure fvs ccs uf bndrs body', fvs)
114 where
115 -- See Note [Tacking local binders]
116 (body', body_fvs) = expr (addLocals bndrs env) body
117 fvs = delVarSetList body_fvs bndrs
118 rhs env (StgRhsCon ccs dc as) = (StgRhsCon ccs dc as, args env as)
119
120 alt :: Env -> StgAlt -> (CgStgAlt, IdSet)
121 alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
122 where
123 -- See Note [Tacking local binders]
124 (e', rhs_fvs) = expr (addLocals bndrs env) e
125 fvs = delVarSetList rhs_fvs bndrs