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