never executed always true always false
1
2 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
3 {-# LANGUAGE TypeFamilies #-}
4
5 -- | Hides away distracting bookkeeping while lambda lifting into a 'LiftM'
6 -- monad.
7 module GHC.Stg.Lift.Monad (
8 decomposeStgBinding, mkStgBinding,
9 Env (..),
10 -- * #floats# Handling floats
11 -- $floats
12 FloatLang (..), collectFloats, -- Exported just for the docs
13 -- * Transformation monad
14 LiftM, runLiftM,
15 -- ** Adding bindings
16 startBindingGroup, endBindingGroup, addTopStringLit, addLiftedBinding,
17 -- ** Substitution and binders
18 withSubstBndr, withSubstBndrs, withLiftedBndr, withLiftedBndrs,
19 -- ** Occurrences
20 substOcc, isLifted, formerFreeVars, liftedIdsExpander
21 ) where
22
23 import GHC.Prelude
24
25 import GHC.Types.Basic
26 import GHC.Types.CostCentre ( isCurrentCCS, dontCareCCS )
27 import GHC.Driver.Session
28 import GHC.Data.FastString
29 import GHC.Types.Id
30 import GHC.Types.Name
31 import GHC.Utils.Outputable
32 import GHC.Data.OrdList
33 import GHC.Stg.Subst
34 import GHC.Stg.Syntax
35 import GHC.Core.Utils
36 import GHC.Types.Unique.Supply
37 import GHC.Utils.Panic
38 import GHC.Utils.Panic.Plain
39 import GHC.Types.Var.Env
40 import GHC.Types.Var.Set
41 import GHC.Core.Multiplicity
42
43 import Control.Arrow ( second )
44 import Control.Monad.Trans.Class
45 import Control.Monad.Trans.RWS.Strict ( RWST, runRWST )
46 import qualified Control.Monad.Trans.RWS.Strict as RWS
47 import Control.Monad.Trans.Cont ( ContT (..) )
48 import Data.ByteString ( ByteString )
49
50 -- | @uncurry 'mkStgBinding' . 'decomposeStgBinding' = id@
51 decomposeStgBinding :: GenStgBinding pass -> (RecFlag, [(BinderP pass, GenStgRhs pass)])
52 decomposeStgBinding (StgRec pairs) = (Recursive, pairs)
53 decomposeStgBinding (StgNonRec bndr rhs) = (NonRecursive, [(bndr, rhs)])
54
55 mkStgBinding :: RecFlag -> [(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
56 mkStgBinding Recursive = StgRec
57 mkStgBinding NonRecursive = uncurry StgNonRec . head
58
59 -- | Environment threaded around in a scoped, @Reader@-like fashion.
60 data Env
61 = Env
62 { e_dflags :: !DynFlags
63 -- ^ Read-only.
64 , e_subst :: !Subst
65 -- ^ We need to track the renamings of local 'InId's to their lifted 'OutId',
66 -- because shadowing might make a closure's free variables unavailable at its
67 -- call sites. Consider:
68 -- @
69 -- let f y = x + y in let x = 4 in f x
70 -- @
71 -- Here, @f@ can't be lifted to top-level, because its free variable @x@ isn't
72 -- available at its call site.
73 , e_expansions :: !(IdEnv DIdSet)
74 -- ^ Lifted 'Id's don't occur as free variables in any closure anymore, because
75 -- they are bound at the top-level. Every occurrence must supply the formerly
76 -- free variables of the lifted 'Id', so they in turn become free variables of
77 -- the call sites. This environment tracks this expansion from lifted 'Id's to
78 -- their free variables.
79 --
80 -- 'InId's to 'OutId's.
81 --
82 -- Invariant: 'Id's not present in this map won't be substituted.
83 }
84
85 emptyEnv :: DynFlags -> Env
86 emptyEnv dflags = Env dflags emptySubst emptyVarEnv
87
88
89 -- Note [Handling floats]
90 -- ~~~~~~~~~~~~~~~~~~~~~~
91 -- $floats
92 -- Consider the following expression:
93 --
94 -- @
95 -- f x =
96 -- let g y = ... f y ...
97 -- in g x
98 -- @
99 --
100 -- What happens when we want to lift @g@? Normally, we'd put the lifted @l_g@
101 -- binding above the binding for @f@:
102 --
103 -- @
104 -- g f y = ... f y ...
105 -- f x = g f x
106 -- @
107 --
108 -- But this very unnecessarily turns a known call to @f@ into an unknown one, in
109 -- addition to complicating matters for the analysis.
110 -- Instead, we'd really like to put both functions in the same recursive group,
111 -- thereby preserving the known call:
112 --
113 -- @
114 -- Rec {
115 -- g y = ... f y ...
116 -- f x = g x
117 -- }
118 -- @
119 --
120 -- But we don't want this to happen for just /any/ binding. That would create
121 -- possibly huge recursive groups in the process, calling for an occurrence
122 -- analyser on STG.
123 -- So, we need to track when we lift a binding out of a recursive RHS and add
124 -- the binding to the same recursive group as the enclosing recursive binding
125 -- (which must have either already been at the top-level or decided to be
126 -- lifted itself in order to preserve the known call).
127 --
128 -- This is done by expressing this kind of nesting structure as a 'Writer' over
129 -- @['FloatLang']@ and flattening this expression in 'runLiftM' by a call to
130 -- 'collectFloats'.
131 -- API-wise, the analysis will not need to know about the whole 'FloatLang'
132 -- business and will just manipulate it indirectly through actions in 'LiftM'.
133
134 -- | We need to detect when we are lifting something out of the RHS of a
135 -- recursive binding (c.f. "GHC.Stg.Lift.Monad#floats"), in which case that
136 -- binding needs to be added to the same top-level recursive group. This
137 -- requires we detect a certain nesting structure, which is encoded by
138 -- 'StartBindingGroup' and 'EndBindingGroup'.
139 --
140 -- Although 'collectFloats' will only ever care if the current binding to be
141 -- lifted (through 'LiftedBinding') will occur inside such a binding group or
142 -- not, e.g. doesn't care about the nesting level as long as its greater than 0.
143 data FloatLang
144 = StartBindingGroup
145 | EndBindingGroup
146 | PlainTopBinding OutStgTopBinding
147 | LiftedBinding OutStgBinding
148
149 instance Outputable FloatLang where
150 ppr StartBindingGroup = char '('
151 ppr EndBindingGroup = char ')'
152 ppr (PlainTopBinding StgTopStringLit{}) = text "<str>"
153 ppr (PlainTopBinding (StgTopLifted b)) = ppr (LiftedBinding b)
154 ppr (LiftedBinding bind) = (if isRec rec then char 'r' else char 'n') <+> ppr (map fst pairs)
155 where
156 (rec, pairs) = decomposeStgBinding bind
157
158 -- | Flattens an expression in @['FloatLang']@ into an STG program, see "GHC.Stg.Lift.Monad#floats".
159 -- Important pre-conditions: The nesting of opening 'StartBindinGroup's and
160 -- closing 'EndBindinGroup's is balanced. Also, it is crucial that every binding
161 -- group has at least one recursive binding inside. Otherwise there's no point
162 -- in announcing the binding group in the first place and an @ASSERT@ will
163 -- trigger.
164 collectFloats :: [FloatLang] -> [OutStgTopBinding]
165 collectFloats = go (0 :: Int) []
166 where
167 go 0 [] [] = []
168 go _ _ [] = pprPanic "collectFloats" (text "unterminated group")
169 go n binds (f:rest) = case f of
170 StartBindingGroup -> go (n+1) binds rest
171 EndBindingGroup
172 | n == 0 -> pprPanic "collectFloats" (text "no group to end")
173 | n == 1 -> StgTopLifted (merge_binds binds) : go 0 [] rest
174 | otherwise -> go (n-1) binds rest
175 PlainTopBinding top_bind
176 | n == 0 -> top_bind : go n binds rest
177 | otherwise -> pprPanic "collectFloats" (text "plain top binding inside group")
178 LiftedBinding bind
179 | n == 0 -> StgTopLifted (rm_cccs bind) : go n binds rest
180 | otherwise -> go n (bind:binds) rest
181
182 map_rhss f = uncurry mkStgBinding . second (map (second f)) . decomposeStgBinding
183 rm_cccs = map_rhss removeRhsCCCS
184 merge_binds binds = assert (any is_rec binds) $
185 StgRec (concatMap (snd . decomposeStgBinding . rm_cccs) binds)
186 is_rec StgRec{} = True
187 is_rec _ = False
188
189 -- | Omitting this makes for strange closure allocation schemes that crash the
190 -- GC.
191 removeRhsCCCS :: GenStgRhs pass -> GenStgRhs pass
192 removeRhsCCCS (StgRhsClosure ext ccs upd bndrs body)
193 | isCurrentCCS ccs
194 = StgRhsClosure ext dontCareCCS upd bndrs body
195 removeRhsCCCS (StgRhsCon ccs con mu ts args)
196 | isCurrentCCS ccs
197 = StgRhsCon dontCareCCS con mu ts args
198 removeRhsCCCS rhs = rhs
199
200 -- | The analysis monad consists of the following 'RWST' components:
201 --
202 -- * 'Env': Reader-like context. Contains a substitution, info about how
203 -- how lifted identifiers are to be expanded into applications and details
204 -- such as 'DynFlags'.
205 --
206 -- * @'OrdList' 'FloatLang'@: Writer output for the resulting STG program.
207 --
208 -- * No pure state component
209 --
210 -- * But wrapping around 'UniqSM' for generating fresh lifted binders.
211 -- (The @uniqAway@ approach could give the same name to two different
212 -- lifted binders, so this is necessary.)
213 newtype LiftM a
214 = LiftM { unwrapLiftM :: RWST Env (OrdList FloatLang) () UniqSM a }
215 deriving (Functor, Applicative, Monad)
216
217 instance HasDynFlags LiftM where
218 getDynFlags = LiftM (RWS.asks e_dflags)
219
220 instance MonadUnique LiftM where
221 getUniqueSupplyM = LiftM (lift getUniqueSupplyM)
222 getUniqueM = LiftM (lift getUniqueM)
223 getUniquesM = LiftM (lift getUniquesM)
224
225 runLiftM :: DynFlags -> UniqSupply -> LiftM () -> [OutStgTopBinding]
226 runLiftM dflags us (LiftM m) = collectFloats (fromOL floats)
227 where
228 (_, _, floats) = initUs_ us (runRWST m (emptyEnv dflags) ())
229
230 -- | Writes a plain 'StgTopStringLit' to the output.
231 addTopStringLit :: OutId -> ByteString -> LiftM ()
232 addTopStringLit id = LiftM . RWS.tell . unitOL . PlainTopBinding . StgTopStringLit id
233
234 -- | Starts a recursive binding group. See "GHC.Stg.Lift.Monad#floats" and 'collectFloats'.
235 startBindingGroup :: LiftM ()
236 startBindingGroup = LiftM $ RWS.tell $ unitOL $ StartBindingGroup
237
238 -- | Ends a recursive binding group. See "GHC.Stg.Lift.Monad#floats" and 'collectFloats'.
239 endBindingGroup :: LiftM ()
240 endBindingGroup = LiftM $ RWS.tell $ unitOL $ EndBindingGroup
241
242 -- | Lifts a binding to top-level. Depending on whether it's declared inside
243 -- a recursive RHS (see "GHC.Stg.Lift.Monad#floats" and 'collectFloats'), this might be added to
244 -- an existing recursive top-level binding group.
245 addLiftedBinding :: OutStgBinding -> LiftM ()
246 addLiftedBinding = LiftM . RWS.tell . unitOL . LiftedBinding
247
248 -- | Takes a binder and a continuation which is called with the substituted
249 -- binder. The continuation will be evaluated in a 'LiftM' context in which that
250 -- binder is deemed in scope. Think of it as a 'RWS.local' computation: After
251 -- the continuation finishes, the new binding won't be in scope anymore.
252 withSubstBndr :: Id -> (Id -> LiftM a) -> LiftM a
253 withSubstBndr bndr inner = LiftM $ do
254 subst <- RWS.asks e_subst
255 let (bndr', subst') = substBndr bndr subst
256 RWS.local (\e -> e { e_subst = subst' }) (unwrapLiftM (inner bndr'))
257
258 -- | See 'withSubstBndr'.
259 withSubstBndrs :: Traversable f => f Id -> (f Id -> LiftM a) -> LiftM a
260 withSubstBndrs = runContT . traverse (ContT . withSubstBndr)
261
262 -- | Similarly to 'withSubstBndr', this function takes a set of variables to
263 -- abstract over, the binder to lift (and generate a fresh, substituted name
264 -- for) and a continuation in which that fresh, lifted binder is in scope.
265 --
266 -- It takes care of all the details involved with copying and adjusting the
267 -- binder and fresh name generation.
268 withLiftedBndr :: DIdSet -> Id -> (Id -> LiftM a) -> LiftM a
269 withLiftedBndr abs_ids bndr inner = do
270 uniq <- getUniqueM
271 let str = "$l" ++ occNameString (getOccName bndr)
272 let ty = mkLamTypes (dVarSetElems abs_ids) (idType bndr)
273 let bndr'
274 -- See Note [transferPolyIdInfo] in GHC.Types.Id. We need to do this at least
275 -- for arity information.
276 = transferPolyIdInfo bndr (dVarSetElems abs_ids)
277 . mkSysLocal (mkFastString str) uniq Many
278 $ ty
279 LiftM $ RWS.local
280 (\e -> e
281 { e_subst = extendSubst bndr bndr' $ extendInScope bndr' $ e_subst e
282 , e_expansions = extendVarEnv (e_expansions e) bndr abs_ids
283 })
284 (unwrapLiftM (inner bndr'))
285
286 -- | See 'withLiftedBndr'.
287 withLiftedBndrs :: Traversable f => DIdSet -> f Id -> (f Id -> LiftM a) -> LiftM a
288 withLiftedBndrs abs_ids = runContT . traverse (ContT . withLiftedBndr abs_ids)
289
290 -- | Substitutes a binder /occurrence/, which was brought in scope earlier by
291 -- 'withSubstBndr' \/ 'withLiftedBndr'.
292 substOcc :: Id -> LiftM Id
293 substOcc id = LiftM (RWS.asks (lookupIdSubst id . e_subst))
294
295 -- | Whether the given binding was decided to be lambda lifted.
296 isLifted :: InId -> LiftM Bool
297 isLifted bndr = LiftM (RWS.asks (elemVarEnv bndr . e_expansions))
298
299 -- | Returns an empty list for a binding that was not lifted and the list of all
300 -- local variables the binding abstracts over (so, exactly the additional
301 -- arguments at adjusted call sites) otherwise.
302 formerFreeVars :: InId -> LiftM [OutId]
303 formerFreeVars f = LiftM $ do
304 expansions <- RWS.asks e_expansions
305 pure $ case lookupVarEnv expansions f of
306 Nothing -> []
307 Just fvs -> dVarSetElems fvs
308
309 -- | Creates an /expander function/ for the current set of lifted binders.
310 -- This expander function will replace any 'InId' by their corresponding 'OutId'
311 -- and, in addition, will expand any lifted binders by the former free variables
312 -- it abstracts over.
313 liftedIdsExpander :: LiftM (DIdSet -> DIdSet)
314 liftedIdsExpander = LiftM $ do
315 expansions <- RWS.asks e_expansions
316 subst <- RWS.asks e_subst
317 -- We use @noWarnLookupIdSubst@ here in order to suppress "not in scope"
318 -- warnings generated by 'lookupIdSubst' due to local bindings within RHS.
319 -- These are not in the InScopeSet of @subst@ and extending the InScopeSet in
320 -- @goodToLift@/@closureGrowth@ before passing it on to @expander@ is too much
321 -- trouble.
322 let go set fv = case lookupVarEnv expansions fv of
323 Nothing -> extendDVarSet set (noWarnLookupIdSubst fv subst) -- Not lifted
324 Just fvs' -> unionDVarSet set fvs'
325 let expander fvs = foldl' go emptyDVarSet (dVarSetElems fvs)
326 pure expander