fbccf80b64aff8a93a0fdb8f1bef05a3718878be
[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 operationally equivalent), but that we cannot common up in Core, because
9 their types differ.
10 This was originally 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 first 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 different 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 Note [StgCse after unarisation]
72 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
73 Consider two unboxed sum terms:
74
75 (# 1 | #) :: (# Int | Int# #)
76 (# 1 | #) :: (# Int | Int #)
77
78 These two terms are not equal as they unarise to different unboxed
79 tuples. However if we run StgCse before Unarise, it'll think the two
80 terms (# 1 | #) are equal, and replace one of these with a binder to
81 the other. That's bad -- Trac #15300.
82
83 Solution: do unarise first.
84
85 -}
86
87 module StgCse (stgCse) where
88
89 import GhcPrelude
90
91 import DataCon
92 import Id
93 import StgSyn
94 import Outputable
95 import VarEnv
96 import CoreSyn (AltCon(..))
97 import Data.List (mapAccumL)
98 import Data.Maybe (fromMaybe)
99 import CoreMap
100 import NameEnv
101 import Control.Monad( (>=>) )
102
103 --------------
104 -- The Trie --
105 --------------
106
107 -- A lookup trie for data constructor applications, i.e.
108 -- keys of type `(DataCon, [StgArg])`, following the patterns in TrieMap.
109
110 data StgArgMap a = SAM
111 { sam_var :: DVarEnv a
112 , sam_lit :: LiteralMap a
113 }
114
115 instance TrieMap StgArgMap where
116 type Key StgArgMap = StgArg
117 emptyTM = SAM { sam_var = emptyTM
118 , sam_lit = emptyTM }
119 lookupTM (StgVarArg var) = sam_var >.> lkDFreeVar var
120 lookupTM (StgLitArg lit) = sam_lit >.> lookupTM lit
121 alterTM (StgVarArg var) f m = m { sam_var = sam_var m |> xtDFreeVar var f }
122 alterTM (StgLitArg lit) f m = m { sam_lit = sam_lit m |> alterTM lit f }
123 foldTM k m = foldTM k (sam_var m) . foldTM k (sam_lit m)
124 mapTM f (SAM {sam_var = varm, sam_lit = litm}) =
125 SAM { sam_var = mapTM f varm, sam_lit = mapTM f litm }
126
127 newtype ConAppMap a = CAM { un_cam :: DNameEnv (ListMap StgArgMap a) }
128
129 instance TrieMap ConAppMap where
130 type Key ConAppMap = (DataCon, [StgArg])
131 emptyTM = CAM emptyTM
132 lookupTM (dataCon, args) = un_cam >.> lkDNamed dataCon >=> lookupTM args
133 alterTM (dataCon, args) f m =
134 m { un_cam = un_cam m |> xtDNamed dataCon |>> alterTM args f }
135 foldTM k = un_cam >.> foldTM (foldTM k)
136 mapTM f = un_cam >.> mapTM (mapTM f) >.> CAM
137
138 -----------------
139 -- The CSE Env --
140 -----------------
141
142 -- | The CSE environment. See note [CseEnv Example]
143 data CseEnv = CseEnv
144 { ce_conAppMap :: ConAppMap OutId
145 -- ^ The main component of the environment is the trie that maps
146 -- data constructor applications (with their `OutId` arguments)
147 -- to an in-scope name that can be used instead.
148 -- This name is always either a let-bound variable or a case binder.
149 , ce_subst :: IdEnv OutId
150 -- ^ This substitution is applied to the code as we traverse it.
151 -- Entries have one of two reasons:
152 --
153 -- * The input might have shadowing (see Note [Shadowing]), so we have
154 -- to rename some binders as we traverse the tree.
155 -- * If we remove `let x = Con z` because `let y = Con z` is in scope,
156 -- we note this here as x ↦ y.
157 , ce_bndrMap :: IdEnv OutId
158 -- ^ If we come across a case expression case x as b of … with a trivial
159 -- binder, we add b ↦ x to this.
160 -- This map is *only* used when looking something up in the ce_conAppMap.
161 -- See Note [Trivial case scrutinee]
162 , ce_in_scope :: InScopeSet
163 -- ^ The third component is an in-scope set, to rename away any
164 -- shadowing binders
165 }
166
167 {-|
168 Note [CseEnv Example]
169 ~~~~~~~~~~~~~~~~~~~~~
170 The following tables shows how the CseEnvironment changes as code is traversed,
171 as well as the changes to that code.
172
173 InExpr OutExpr
174 conAppMap subst in_scope
175 ───────────────────────────────────────────────────────────
176 -- empty {} {}
177 case … as a of {Con x y -> case … as a of {Con x y ->
178 -- Con x y ↦ a {} {a,x,y}
179 let b = Con x y (removed)
180 -- Con x y ↦ a b↦a {a,x,y,b}
181 let c = Bar a let c = Bar a
182 -- Con x y ↦ a, Bar a ↦ c b↦a {a,x,y,b,c}
183 let c = some expression let c' = some expression
184 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', {a,x,y,b,c,c'}
185 let d = Bar b (removed)
186 -- Con x y ↦ a, Bar a ↦ c b↦a, c↦c', d↦c {a,x,y,b,c,c',d}
187 (a, b, c d) (a, a, c' c)
188 -}
189
190 initEnv :: InScopeSet -> CseEnv
191 initEnv in_scope = CseEnv
192 { ce_conAppMap = emptyTM
193 , ce_subst = emptyVarEnv
194 , ce_bndrMap = emptyVarEnv
195 , ce_in_scope = in_scope
196 }
197
198 envLookup :: DataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
199 envLookup dataCon args env = lookupTM (dataCon, args') (ce_conAppMap env)
200 where args' = map go args -- See Note [Trivial case scrutinee]
201 go (StgVarArg v ) = StgVarArg (fromMaybe v $ lookupVarEnv (ce_bndrMap env) v)
202 go (StgLitArg lit) = StgLitArg lit
203
204 addDataCon :: OutId -> DataCon -> [OutStgArg] -> CseEnv -> CseEnv
205 -- do not bother with nullary data constructors, they are static anyways
206 addDataCon _ _ [] env = env
207 addDataCon bndr dataCon args env = env { ce_conAppMap = new_env }
208 where
209 new_env = insertTM (dataCon, args) bndr (ce_conAppMap env)
210
211 forgetCse :: CseEnv -> CseEnv
212 forgetCse env = env { ce_conAppMap = emptyTM }
213 -- See note [Free variables of an StgClosure]
214
215 addSubst :: OutId -> OutId -> CseEnv -> CseEnv
216 addSubst from to env
217 = env { ce_subst = extendVarEnv (ce_subst env) from to }
218
219 addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
220 addTrivCaseBndr from to env
221 = env { ce_bndrMap = extendVarEnv (ce_bndrMap env) from to }
222
223 substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
224 substArgs env = map (substArg env)
225
226 substArg :: CseEnv -> InStgArg -> OutStgArg
227 substArg env (StgVarArg from) = StgVarArg (substVar env from)
228 substArg _ (StgLitArg lit) = StgLitArg lit
229
230 substVar :: CseEnv -> InId -> OutId
231 substVar env id = fromMaybe id $ lookupVarEnv (ce_subst env) id
232
233 -- Functions to enter binders
234
235 -- This is much simpler than the equivalent code in CoreSubst:
236 -- * We do not substitute type variables, and
237 -- * There is nothing relevant in IdInfo at this stage
238 -- that needs substitutions.
239 -- Therefore, no special treatment for a recursive group is required.
240
241 substBndr :: CseEnv -> InId -> (CseEnv, OutId)
242 substBndr env old_id
243 = (new_env, new_id)
244 where
245 new_id = uniqAway (ce_in_scope env) old_id
246 no_change = new_id == old_id
247 env' = env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
248 new_env | no_change = env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }
249 | otherwise = env'
250
251 substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
252 substBndrs env bndrs = mapAccumL substBndr env bndrs
253
254 substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
255 substPairs env bndrs = mapAccumL go env bndrs
256 where go env (id, x) = let (env', id') = substBndr env id
257 in (env', (id', x))
258
259 -- Main entry point
260
261 stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
262 stgCse binds = snd $ mapAccumL stgCseTopLvl emptyInScopeSet binds
263
264 -- Top level bindings.
265 --
266 -- We do not CSE these, as top-level closures are allocated statically anyways.
267 -- Also, they might be exported.
268 -- But we still have to collect the set of in-scope variables, otherwise
269 -- uniqAway might shadow a top-level closure.
270
271 stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
272 stgCseTopLvl in_scope t@(StgTopStringLit _ _) = (in_scope, t)
273 stgCseTopLvl in_scope (StgTopLifted (StgNonRec bndr rhs))
274 = (in_scope'
275 , StgTopLifted (StgNonRec bndr (stgCseTopLvlRhs in_scope rhs)))
276 where in_scope' = in_scope `extendInScopeSet` bndr
277
278 stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
279 = ( in_scope'
280 , StgTopLifted (StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ]))
281 where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
282
283 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
284 stgCseTopLvlRhs in_scope (StgRhsClosure ext ccs upd args body)
285 = let body' = stgCseExpr (initEnv in_scope) body
286 in StgRhsClosure ext ccs upd args body'
287 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
288 = StgRhsCon ccs dataCon args
289
290 ------------------------------
291 -- The actual AST traversal --
292 ------------------------------
293
294 -- Trivial cases
295 stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
296 stgCseExpr env (StgApp fun args)
297 = StgApp fun' args'
298 where fun' = substVar env fun
299 args' = substArgs env args
300 stgCseExpr _ (StgLit lit)
301 = StgLit lit
302 stgCseExpr env (StgOpApp op args tys)
303 = StgOpApp op args' tys
304 where args' = substArgs env args
305 stgCseExpr _ (StgLam _ _)
306 = pprPanic "stgCseExp" (text "StgLam")
307 stgCseExpr env (StgTick tick body)
308 = let body' = stgCseExpr env body
309 in StgTick tick body'
310 stgCseExpr env (StgCase scrut bndr ty alts)
311 = mkStgCase scrut' bndr' ty alts'
312 where
313 scrut' = stgCseExpr env scrut
314 (env1, bndr') = substBndr env bndr
315 env2 | StgApp trivial_scrut [] <- scrut' = addTrivCaseBndr bndr trivial_scrut env1
316 -- See Note [Trivial case scrutinee]
317 | otherwise = env1
318 alts' = map (stgCseAlt env2 ty bndr') alts
319
320
321 -- A constructor application.
322 -- To be removed by a variable use when found in the CSE environment
323 stgCseExpr env (StgConApp dataCon args tys)
324 | Just bndr' <- envLookup dataCon args' env
325 = StgApp bndr' []
326 | otherwise
327 = StgConApp dataCon args' tys
328 where args' = substArgs env args
329
330 -- Let bindings
331 -- The binding might be removed due to CSE (we do not want trivial bindings on
332 -- the STG level), so use the smart constructor `mkStgLet` to remove the binding
333 -- if empty.
334 stgCseExpr env (StgLet ext binds body)
335 = let (binds', env') = stgCseBind env binds
336 body' = stgCseExpr env' body
337 in mkStgLet (StgLet ext) binds' body'
338 stgCseExpr env (StgLetNoEscape ext binds body)
339 = let (binds', env') = stgCseBind env binds
340 body' = stgCseExpr env' body
341 in mkStgLet (StgLetNoEscape ext) binds' body'
342
343 -- Case alternatives
344 -- Extend the CSE environment
345 stgCseAlt :: CseEnv -> AltType -> OutId -> InStgAlt -> OutStgAlt
346 stgCseAlt env ty case_bndr (DataAlt dataCon, args, rhs)
347 = let (env1, args') = substBndrs env args
348 env2
349 -- To avoid dealing with unboxed sums StgCse runs after unarise and
350 -- should maintain invariants listed in Note [Post-unarisation
351 -- invariants]. One of the invariants is that some binders are not
352 -- used (unboxed tuple case binders) which is what we check with
353 -- `stgCaseBndrInScope` here. If the case binder is not in scope we
354 -- don't add it to the CSE env. See also #15300.
355 | stgCaseBndrInScope ty True -- CSE runs after unarise
356 = addDataCon case_bndr dataCon (map StgVarArg args') env1
357 | otherwise
358 = env1
359 -- see note [Case 2: CSEing case binders]
360 rhs' = stgCseExpr env2 rhs
361 in (DataAlt dataCon, args', rhs')
362 stgCseAlt env _ _ (altCon, args, rhs)
363 = let (env1, args') = substBndrs env args
364 rhs' = stgCseExpr env1 rhs
365 in (altCon, args', rhs')
366
367 -- Bindings
368 stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
369 stgCseBind env (StgNonRec b e)
370 = let (env1, b') = substBndr env b
371 in case stgCseRhs env1 b' e of
372 (Nothing, env2) -> (Nothing, env2)
373 (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
374 stgCseBind env (StgRec pairs)
375 = let (env1, pairs1) = substPairs env pairs
376 in case stgCsePairs env1 pairs1 of
377 ([], env2) -> (Nothing, env2)
378 (pairs2, env2) -> (Just (StgRec pairs2), env2)
379
380 stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
381 stgCsePairs env [] = ([], env)
382 stgCsePairs env0 ((b,e):pairs)
383 = let (pairMB, env1) = stgCseRhs env0 b e
384 (pairs', env2) = stgCsePairs env1 pairs
385 in (pairMB `mbCons` pairs', env2)
386 where
387 mbCons = maybe id (:)
388
389 -- The RHS of a binding.
390 -- If it is a constructor application, either short-cut it or extend the environment
391 stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
392 stgCseRhs env bndr (StgRhsCon ccs dataCon args)
393 | Just other_bndr <- envLookup dataCon args' env
394 = let env' = addSubst bndr other_bndr env
395 in (Nothing, env')
396 | otherwise
397 = let env' = addDataCon bndr dataCon args' env
398 -- see note [Case 1: CSEing allocated closures]
399 pair = (bndr, StgRhsCon ccs dataCon args')
400 in (Just pair, env')
401 where args' = substArgs env args
402 stgCseRhs env bndr (StgRhsClosure ext ccs upd args body)
403 = let (env1, args') = substBndrs env args
404 env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
405 body' = stgCseExpr env2 body
406 in (Just (substVar env bndr, StgRhsClosure ext ccs upd args' body'), env)
407
408
409 mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
410 mkStgCase scrut bndr ty alts | all isBndr alts = scrut
411 | otherwise = StgCase scrut bndr ty alts
412
413 where
414 -- see Note [All alternatives are the binder]
415 isBndr (_, _, StgApp f []) = f == bndr
416 isBndr _ = False
417
418
419 -- Utilities
420
421 -- | This function short-cuts let-bindings that are now obsolete
422 mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
423 mkStgLet _ Nothing body = body
424 mkStgLet stgLet (Just binds) body = stgLet binds body
425
426
427 {-
428 Note [All alternatives are the binder]
429 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
430
431 When all alternatives simply refer to the case binder, then we do not have
432 to bother with the case expression at all (#13588). CoreSTG does this as well,
433 but sometimes, types get into the way:
434
435 newtype T = MkT Int
436 f :: (Int, Int) -> (T, Int)
437 f (x, y) = (MkT x, y)
438
439 Core cannot just turn this into
440
441 f p = p
442
443 as this would not be well-typed. But to STG, where MkT is no longer in the way,
444 we can.
445
446 Note [Trivial case scrutinee]
447 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
448 We want to be able to handle nested reconstruction of constructors as in
449
450 nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
451 nested (Right (Right v)) = Right (Right v)
452 nested _ = Left True
453
454 So if we come across
455
456 case x of r1
457 Right a -> case a of r2
458 Right b -> let v = Right b
459 in Right v
460
461 we first replace v with r2. Next we want to replace Right r2 with r1. But the
462 ce_conAppMap contains Right a!
463
464 Therefore, we add r1 ↦ x to ce_bndrMap when analysing the outer case, and use
465 this substitution before looking Right r2 up in ce_conAppMap, and everything
466 works out.
467
468 Note [Free variables of an StgClosure]
469 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
470 StgClosures (function and thunks) have an explicit list of free variables:
471
472 foo [x] =
473 let not_a_free_var = Left [x]
474 let a_free_var = Right [x]
475 let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
476 in closure
477
478 If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
479 then the list of free variables would be wrong, so for now, we do not CSE
480 across such a closure, simply because I (Joachim) was not sure about possible
481 knock-on effects. If deemed safe and worth the slight code complication of
482 re-calculating this list during or after this pass, this can surely be done.
483 -}