never executed always true always false
1 {-# LANGUAGE TypeFamilies #-}
2
3 {- |
4 Non-global free variable analysis on STG terms. This pass annotates
5 non-top-level closure bindings with captured variables. Global variables are not
6 captured. For example, in a top-level binding like (pseudo-STG)
7
8 f = \[x,y] .
9 let g = \[p] . reverse (x ++ p)
10 in g y
11
12 In g, `reverse` and `(++)` are global variables so they're not considered free.
13 `p` is an argument, so `x` is the only actual free variable here. The annotated
14 version is thus:
15
16 f = \[x,y] .
17 let g = [x] \[p] . reverse (x ++ p)
18 in g y
19
20 Note that non-top-level recursive bindings are also considered free within the
21 group:
22
23 map = {} \r [f xs0]
24 let {
25 Rec {
26 go = {f, go} \r [xs1]
27 case xs1 of {
28 [] -> [] [];
29 : x xs2 ->
30 let { xs' = {go, xs2} \u [] go xs2; } in
31 let { x' = {f, x} \u [] f x; } in
32 : [x' xs'];
33 };
34 end Rec }
35 } in go xs0;
36
37 Here go is free in its RHS.
38
39 Top-level closure bindings never capture variables as all of their free
40 variables are global.
41 -}
42 module GHC.Stg.FVs (
43 annTopBindingsFreeVars,
44 annBindingFreeVars
45 ) where
46
47 import GHC.Prelude
48
49 import GHC.Stg.Syntax
50 import GHC.Types.Id
51 import GHC.Types.Var.Set
52 import GHC.Types.Tickish ( GenTickish(Breakpoint) )
53 import GHC.Utils.Misc
54
55 import Data.Maybe ( mapMaybe )
56
57 newtype Env
58 = Env
59 { locals :: IdSet
60 }
61
62 emptyEnv :: Env
63 emptyEnv = Env emptyVarSet
64
65 addLocals :: [Id] -> Env -> Env
66 addLocals bndrs env
67 = env { locals = extendVarSetList (locals env) bndrs }
68
69 -- | Annotates a top-level STG binding group with its free variables.
70 annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
71 annTopBindingsFreeVars = map go
72 where
73 go (StgTopStringLit id bs) = StgTopStringLit id bs
74 go (StgTopLifted bind)
75 = StgTopLifted (annBindingFreeVars bind)
76
77 -- | Annotates an STG binding with its free variables.
78 annBindingFreeVars :: StgBinding -> CgStgBinding
79 annBindingFreeVars = fst . binding emptyEnv emptyDVarSet
80
81 boundIds :: StgBinding -> [Id]
82 boundIds (StgNonRec b _) = [b]
83 boundIds (StgRec pairs) = map fst pairs
84
85 -- Note [Tracking local binders]
86 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
87 -- 'locals' contains non-toplevel, non-imported binders.
88 -- We maintain the set in 'expr', 'alt' and 'rhs', which are the only
89 -- places where new local binders are introduced.
90 -- Why do it there rather than in 'binding'? Two reasons:
91 --
92 -- 1. We call 'binding' from 'annTopBindingsFreeVars', which would
93 -- add top-level bindings to the 'locals' set.
94 -- 2. In the let(-no-escape) case, we need to extend the environment
95 -- prior to analysing the body, but we also need the fvs from the
96 -- body to analyse the RHSs. No way to do this without some
97 -- knot-tying.
98
99 -- | This makes sure that only local, non-global free vars make it into the set.
100 mkFreeVarSet :: Env -> [Id] -> DIdSet
101 mkFreeVarSet env = mkDVarSet . filter (`elemVarSet` locals env)
102
103 args :: Env -> [StgArg] -> DIdSet
104 args env = mkFreeVarSet env . mapMaybe f
105 where
106 f (StgVarArg occ) = Just occ
107 f _ = Nothing
108
109 binding :: Env -> DIdSet -> StgBinding -> (CgStgBinding, DIdSet)
110 binding env body_fv (StgNonRec bndr r) = (StgNonRec bndr r', fvs)
111 where
112 -- See Note [Tracking local binders]
113 (r', rhs_fvs) = rhs env r
114 fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_fvs
115 binding env body_fv (StgRec pairs) = (StgRec pairs', fvs)
116 where
117 -- See Note [Tracking local binders]
118 bndrs = map fst pairs
119 (rhss, rhs_fvss) = mapAndUnzip (rhs env . snd) pairs
120 pairs' = zip bndrs rhss
121 fvs = delDVarSetList (unionDVarSets (body_fv:rhs_fvss)) bndrs
122
123 expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
124 expr env = go
125 where
126 go (StgApp occ as)
127 = (StgApp occ as, unionDVarSet (args env as) (mkFreeVarSet env [occ]))
128 go (StgLit lit) = (StgLit lit, emptyDVarSet)
129 go (StgConApp dc n as tys) = (StgConApp dc n as tys, args env as)
130 go (StgOpApp op as ty) = (StgOpApp op as ty, args env as)
131 go (StgCase scrut bndr ty alts) = (StgCase scrut' bndr ty alts', fvs)
132 where
133 (scrut', scrut_fvs) = go scrut
134 -- See Note [Tracking local binders]
135 (alts', alt_fvss) = mapAndUnzip (alt (addLocals [bndr] env)) alts
136 alt_fvs = unionDVarSets alt_fvss
137 fvs = delDVarSet (unionDVarSet scrut_fvs alt_fvs) bndr
138 go (StgLet ext bind body) = go_bind (StgLet ext) bind body
139 go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
140 go (StgTick tick e) = (StgTick tick e', fvs')
141 where
142 (e', fvs) = go e
143 fvs' = unionDVarSet (tickish tick) fvs
144 tickish (Breakpoint _ _ ids) = mkDVarSet ids
145 tickish _ = emptyDVarSet
146
147 go_bind dc bind body = (dc bind' body', fvs)
148 where
149 -- See Note [Tracking local binders]
150 env' = addLocals (boundIds bind) env
151 (body', body_fvs) = expr env' body
152 (bind', fvs) = binding env' body_fvs bind
153
154 rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
155 rhs env (StgRhsClosure _ ccs uf bndrs body)
156 = (StgRhsClosure fvs ccs uf bndrs body', fvs)
157 where
158 -- See Note [Tracking local binders]
159 (body', body_fvs) = expr (addLocals bndrs env) body
160 fvs = delDVarSetList body_fvs bndrs
161 rhs env (StgRhsCon ccs dc mu ts as) = (StgRhsCon ccs dc mu ts as, args env as)
162
163 alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
164 alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
165 where
166 -- See Note [Tracking local binders]
167 (e', rhs_fvs) = expr (addLocals bndrs env) e
168 fvs = delDVarSetList rhs_fvs bndrs