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