7454d24a2cda4710eac5b1577f66a575f3809358
[ghc.git] / compiler / simplStg / StgCse.hs
1 {-# LANGUAGE TypeFamilies #-}
2
3 {-|
4 Note [CSE for Stg]
5 ~~~~~~~~~~~~~~~~~~
6 This module implements a simple common subexpression elimination pass for STG.
7 This is useful because there are expressions that we want to common up (because
8 they are operational equivalent), but that we cannot common up in Core, because
9 their types differ.
10 This was original reported as #9291.
11
12 There are two types of common code occurrences that we aim for, see
13 note [Case 1: CSEing allocated closures] and
14 note [Case 2: CSEing case binders] below.
15
16
17 Note [Case 1: CSEing allocated closures]
18 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
19 The fist kind of CSE opportunity we aim for is generated by this Haskell code:
20
21 bar :: a -> (Either Int a, Either Bool a)
22 bar x = (Right x, Right x)
23
24 which produces this Core:
25
26 bar :: forall a. a -> (Either Int a, Either Bool a)
27 bar @a x = (Right @Int @a x, Right @Bool @a x)
28
29 where the two components of the tuple are differnt terms, and cannot be
30 commoned up (easily). On the STG level we have
31
32 bar [x] = let c1 = Right [x]
33 c2 = Right [x]
34 in (c1,c2)
35
36 and now it is obvious that we can write
37
38 bar [x] = let c1 = Right [x]
39 in (c1,c1)
40
41 instead.
42
43
44 Note [Case 2: CSEing case binders]
45 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
46 The second kind of CSE opportunity we aim for is more interesting, and
47 came up in #9291 and #5344: The Haskell code
48
49 foo :: Either Int a -> Either Bool a
50 foo (Right x) = Right x
51 foo _ = Left False
52
53 produces this Core
54
55 foo :: forall a. Either Int a -> Either Bool a
56 foo @a e = case e of b { Left n -> …
57 , Right x -> Right @Bool @a x }
58
59 where we cannot CSE `Right @Bool @a x` with the case binder `b` as they have
60 different types. But in STG we have
61
62 foo [e] = case e of b { Left [n] -> …
63 , Right [x] -> Right [x] }
64
65 and nothing stops us from transforming that to
66
67 foo [e] = case e of b { Left [n] -> …
68 , Right [x] -> b}
69
70 -}
71 module StgCse (stgCse) where
72
73 import DataCon
74 import Id
75 import StgSyn
76 import Outputable
77 import VarEnv
78 import CoreSyn (AltCon(..))
79 import Data.List (mapAccumL)
80 import Data.Maybe (fromMaybe)
81 import TrieMap
82 import NameEnv
83 import Control.Monad( (>=>) )
84
85 --------------
86 -- The Trie --
87 --------------
88
89 -- A lookup trie for data constructor applications, i.e.
90 -- keys of type `(DataCon, [StgArg])`, following the patterns in TrieMap.
91
92 data StgArgMap a = SAM
93 { sam_var :: DVarEnv a
94 , sam_lit :: LiteralMap a
95 }
96
97 instance TrieMap StgArgMap where
98 type Key StgArgMap = StgArg
99 emptyTM = SAM { sam_var = emptyTM
100 , sam_lit = emptyTM }
101 lookupTM (StgVarArg var) = sam_var >.> lkDFreeVar var
102 lookupTM (StgLitArg lit) = sam_lit >.> lookupTM lit
103 alterTM (StgVarArg var) f m = m { sam_var = sam_var m |> xtDFreeVar var f }
104 alterTM (StgLitArg lit) f m = m { sam_lit = sam_lit m |> alterTM lit f }
105 foldTM k m = foldTM k (sam_var m) . foldTM k (sam_lit m)
106 mapTM f (SAM {sam_var = varm, sam_lit = litm}) =
107 SAM { sam_var = mapTM f varm, sam_lit = mapTM f litm }
108
109 newtype ConAppMap a = CAM { un_cam :: DNameEnv (ListMap StgArgMap a) }
110
111 instance TrieMap ConAppMap where
112 type Key ConAppMap = (DataCon, [StgArg])
113 emptyTM = CAM emptyTM
114 lookupTM (dataCon, args) = un_cam >.> lkDNamed dataCon >=> lookupTM args
115 alterTM (dataCon, args) f m =
116 m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
117 foldTM k = un_cam >.> foldTM (foldTM k)
118 mapTM f = un_cam >.> mapTM (mapTM f) >.> CAM
119
120 -----------------
121 -- The CSE Env --
122 -----------------
123
124 -- | The CSE environment. See note [CseEnv Example]
125 data CseEnv = CseEnv
126 { ce_conAppMap :: ConAppMap OutId
127 -- ^ The main component of the environment is the trie that maps
128 -- data constructor applications (with their `OutId` arguments)
129 -- to an in-scope name that can be used instead.
130 , ce_renaming :: IdEnv OutId
131 -- ^ CSE is simple to implement (and reason about) when there is no
132 -- shadowing. Unfortunately, we have to cope with shadowing
133 -- (see Note [Shadowing]). So we morally do a separate renaming pass
134 -- before CSE, and practically do both passes in one traversal of the tree.
135 -- It still causes less confusion to keep the renaming substitution
136 -- and the substitutions due to CSE separate.
137 , ce_subst :: IdEnv OutId
138 -- ^ This substitution contains CSE-specific entries. The domain are
139 -- OutIds, so ce_renaming has to be applied first.
140 -- It has an entry x ↦ y when a let-binding `let x = Con y` is
141 -- removed because `let y = Con z` is in scope.
142 --
143 -- Both substitutions are applied to data constructor arguments
144 -- before these are looked up in the conAppMap.
145 , ce_in_scope :: InScopeSet
146 -- ^ The third component is an in-scope set, to rename away any
147 -- shadowing binders
148 }
149
150 {-|
151 Note [CseEnv Example]
152 ~~~~~~~~~~~~~~~~~~~~~
153 The following tables shows how the CseEnvironment changes as code is traversed,
154 as well as the changes to that code.
155
156 InExpr OutExpr
157 conAppMap renaming subst in_scope
158 ──────────────────────────────────────────────────────────────────────
159 -- empty {} {} {}
160 case … as a of {Con x y -> case … as a of {Con x y ->
161 -- Con x y ↦ a {} {} {a,x,y}
162 let b = Con x y (removed)
163 -- Con x y ↦ a {} b↦a {a,x,y,b}
164 let c = Bar a let c = Bar a
165 -- Con x y ↦ a, Bar a ↦ c {} b↦a {a,x,y,b,c}
166 let c = some expression let c' = some expression
167 -- Con x y ↦ a, Bar a ↦ c c↦c' b↦a {a,x,y,b,c,c'}
168 let d = Bar b (removed)
169 -- Con x y ↦ a, Bar a ↦ c c↦c' b↦a, d↦c {a,x,y,b,c,c',d}
170 (a, b, c d) (a, a, c' c)
171 -}
172
173 initEnv :: InScopeSet -> CseEnv
174 initEnv in_scope = CseEnv
175 { ce_conAppMap = emptyTM
176 , ce_renaming = emptyVarEnv
177 , ce_subst = emptyVarEnv
178 , ce_in_scope = in_scope
179 }
180
181 envLookup :: DataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
182 envLookup dataCon args env = lookupTM (dataCon, args) (ce_conAppMap env)
183
184 addDataCon :: OutId -> DataCon -> [OutStgArg] -> CseEnv -> CseEnv
185 -- do not bother with nullary data constructors, they are static anyways
186 addDataCon _ _ [] env = env
187 addDataCon bndr dataCon args env = env { ce_conAppMap = new_env }
188 where
189 new_env = insertTM (dataCon, args) bndr (ce_conAppMap env)
190
191 forgetCse :: CseEnv -> CseEnv
192 forgetCse env = env { ce_conAppMap = emptyTM }
193 -- See note [Free variables of an StgClosure]
194
195 addSubst :: OutId -> OutId -> CseEnv -> CseEnv
196 addSubst from to env
197 = env { ce_subst = extendVarEnv (ce_subst env) from to }
198
199 substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
200 substArgs env = map (substArg env)
201
202 substArg :: CseEnv -> InStgArg -> OutStgArg
203 substArg env (StgVarArg from) = StgVarArg (substVar env from)
204 substArg _ (StgLitArg lit) = StgLitArg lit
205
206 substVars :: CseEnv -> [InId] -> [OutId]
207 substVars env = map (substVar env)
208
209 substVar :: CseEnv -> InId -> OutId
210 substVar env id0 = id2
211 where
212 id1 = fromMaybe id0 $ lookupVarEnv (ce_renaming env) id0
213 id2 = fromMaybe id1 $ lookupVarEnv (ce_subst env) id1
214
215 -- Functions to enter binders
216
217 -- This is much simpler than the requivalent code in CoreSubst:
218 -- * We do not substitute type variables, and
219 -- * There is nothing relevant in IdInfo at this stage
220 -- that needs substitutions.
221 -- Therefore, no special treatment for a recursive group is required.
222
223 substBndr :: CseEnv -> InId -> (CseEnv, OutId)
224 substBndr env old_id
225 = (new_env, new_id)
226 where
227 new_id = uniqAway (ce_in_scope env) old_id
228 no_change = new_id == old_id
229 env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
230 new_env | no_change = env' { ce_renaming = extendVarEnv (ce_subst env) old_id new_id }
231 | otherwise = env'
232
233 substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
234 substBndrs env bndrs = mapAccumL substBndr env bndrs
235
236 substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
237 substPairs env bndrs = mapAccumL go env bndrs
238 where go env (id, x) = let (env', id') = substBndr env id
239 in (env', (id', x))
240
241 -- Main entry point
242
243 stgCse :: [InStgBinding] -> [OutStgBinding]
244 stgCse binds = snd $ mapAccumL stgCseTopLvl emptyInScopeSet binds
245
246 -- Top level bindings.
247 --
248 -- We do not CSE these, as top-level closures are allocated statically anyways.
249 -- Also, they might be exported.
250 -- But we still have to collect the set of in-scope variables, otherwise
251 -- uniqAway might shadow a top-level closure.
252
253 stgCseTopLvl :: InScopeSet -> InStgBinding -> (InScopeSet, OutStgBinding)
254 stgCseTopLvl in_scope (StgNonRec bndr rhs)
255 = (in_scope'
256 , StgNonRec bndr (stgCseTopLvlRhs in_scope rhs))
257 where in_scope' = in_scope `extendInScopeSet` bndr
258
259 stgCseTopLvl in_scope (StgRec eqs)
260 = ( in_scope'
261 , StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ])
262 where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
263
264 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
265 stgCseTopLvlRhs in_scope (StgRhsClosure ccs info occs upd args body)
266 = let body' = stgCseExpr (initEnv in_scope) body
267 in StgRhsClosure ccs info occs upd args body'
268 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
269 = StgRhsCon ccs dataCon args
270
271 ------------------------------
272 -- The actual AST traversal --
273 ------------------------------
274
275 -- Trivial cases
276 stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
277 stgCseExpr env (StgApp fun args)
278 = StgApp fun' args'
279 where fun' = substVar env fun
280 args' = substArgs env args
281 stgCseExpr _ (StgLit lit)
282 = StgLit lit
283 stgCseExpr env (StgOpApp op args tys)
284 = StgOpApp op args' tys
285 where args' = substArgs env args
286 stgCseExpr _ (StgLam _ _)
287 = pprPanic "stgCseExp" (text "StgLam")
288 stgCseExpr env (StgTick tick body)
289 = let body' = stgCseExpr env body
290 in StgTick tick body'
291 stgCseExpr env (StgCase scrut bndr ty alts)
292 = StgCase scrut' bndr' ty alts'
293 where
294 scrut' = stgCseExpr env scrut
295 (env1, bndr') = substBndr env bndr
296 cse_bndr | StgApp trivial_scrut [] <- scrut' = trivial_scrut
297 -- See Note [Trivial case scrutinee]
298 | otherwise = bndr'
299 alts' = map (stgCseAlt env1 cse_bndr) alts
300
301
302 -- A constructor application.
303 -- To be removed by a variable use when found in the CSE environment
304 stgCseExpr env (StgConApp dataCon args tys)
305 | Just bndr' <- envLookup dataCon args' env
306 = StgApp bndr' []
307 | otherwise
308 = StgConApp dataCon args' tys
309 where args' = substArgs env args
310
311 -- Let bindings
312 -- The binding might be removed due to CSE (we do not want trivial bindings on
313 -- the STG level), so use the smart constructor `mkStgLet` to remove the binding
314 -- if empty.
315 stgCseExpr env (StgLet binds body)
316 = let (binds', env') = stgCseBind env binds
317 body' = stgCseExpr env' body
318 in mkStgLet StgLet binds' body'
319 stgCseExpr env (StgLetNoEscape binds body)
320 = let (binds', env') = stgCseBind env binds
321 body' = stgCseExpr env' body
322 in mkStgLet StgLetNoEscape binds' body'
323
324 -- Case alternatives
325 -- Extend the CSE environment
326 stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
327 stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
328 = let (env1, args') = substBndrs env args
329 env2 = addDataCon case_bndr dataCon (map StgVarArg args') env1
330 -- see note [Case 2: CSEing case binders]
331 rhs' = stgCseExpr env2 rhs
332 in (DataAlt dataCon, args', rhs')
333 stgCseAlt env _ (altCon, args, rhs)
334 = let (env1, args') = substBndrs env args
335 rhs' = stgCseExpr env1 rhs
336 in (altCon, args', rhs')
337
338 -- Bindings
339 stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
340 stgCseBind env (StgNonRec b e)
341 = let (env1, b') = substBndr env b
342 in case stgCseRhs env1 b' e of
343 (Nothing, env2) -> (Nothing, env2)
344 (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
345 stgCseBind env (StgRec pairs)
346 = let (env1, pairs1) = substPairs env pairs
347 in case stgCsePairs env1 pairs1 of
348 ([], env2) -> (Nothing, env2)
349 (pairs2, env2) -> (Just (StgRec pairs2), env2)
350
351 stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
352 stgCsePairs env [] = ([], env)
353 stgCsePairs env0 ((b,e):pairs)
354 = let (pairMB, env1) = stgCseRhs env0 b e
355 (pairs', env2) = stgCsePairs env1 pairs
356 in (pairMB `mbCons` pairs', env2)
357 where
358 mbCons = maybe id (:)
359
360 -- The RHS of a binding.
361 -- If it is an constructor application, either short-cut it or extend the environment
362 stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
363 stgCseRhs env bndr (StgRhsCon ccs dataCon args)
364 | Just other_bndr <- envLookup dataCon args' env
365 = let env' = addSubst bndr other_bndr env
366 in (Nothing, env')
367 | otherwise
368 = let env' = addDataCon bndr dataCon args' env
369 -- see note [Case 1: CSEing allocated closures]
370 pair = (bndr, StgRhsCon ccs dataCon args')
371 in (Just pair, env')
372 where args' = substArgs env args
373 stgCseRhs env bndr (StgRhsClosure ccs info occs upd args body)
374 = let (env1, args') = substBndrs env args
375 env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
376 body' = stgCseExpr env2 body
377 in (Just (substVar env bndr, StgRhsClosure ccs info occs' upd args' body'), env)
378 where occs' = substVars env occs
379
380 -- Utilities
381
382 -- | This function short-cuts let-bindings that are now obsolete
383 mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
384 mkStgLet _ Nothing body = body
385 mkStgLet stgLet (Just binds) body = stgLet binds body
386
387
388 {-
389 Note [Trivial case scrutinee]
390 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
391 If we find
392
393 case x as b of { Con a -> … }
394
395 we really want to replace uses of Con a in the body with x, and not just b, in
396 order to handle nested reconstruction of constructors as in
397
398 nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
399 nested (Right (Right x)) = Right (Right x)
400 nested _ = Left True
401
402 Therefore, we add
403 Con a ↦ x
404 to the ConAppMap respectively.
405 Compare Note [CSE for case expressions] in CSE.hs, which does the same for Core CSE.
406
407 If we find
408 case foo x as b of { Con a -> … }
409 we use
410 Con a ↦ b
411
412 Note [Free variables of an StgClosure]
413 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
414 StgClosures (function and thunks) have an explicit list of free variables:
415
416 foo [x] =
417 let not_a_free_var = Left [x]
418 let a_free_var = Right [x]
419 let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
420 in closure
421
422 If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
423 then the list of free variables would be wrong, so for now, we do not CSE
424 across such a closure, simply because I (Joachim) was not sure about possible
425 knock-on effects. If deemed safe and worth the slight code complication of
426 re-calculating this list during or after this pass, this can surely be done.
427 -}