Allow top-level string literals in Core (#8472)
[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 :: [InStgTopBinding] -> [OutStgTopBinding]
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 -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
254 stgCseTopLvl in_scope t@(StgTopStringLit _ _) = (in_scope, t)
255 stgCseTopLvl in_scope (StgTopLifted (StgNonRec bndr rhs))
256 = (in_scope'
257 , StgTopLifted (StgNonRec bndr (stgCseTopLvlRhs in_scope rhs)))
258 where in_scope' = in_scope `extendInScopeSet` bndr
259
260 stgCseTopLvl in_scope (StgTopLifted (StgRec eqs))
261 = ( in_scope'
262 , StgTopLifted (StgRec [ (bndr, stgCseTopLvlRhs in_scope' rhs) | (bndr, rhs) <- eqs ]))
263 where in_scope' = in_scope `extendInScopeSetList` [ bndr | (bndr, _) <- eqs ]
264
265 stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
266 stgCseTopLvlRhs in_scope (StgRhsClosure ccs info occs upd args body)
267 = let body' = stgCseExpr (initEnv in_scope) body
268 in StgRhsClosure ccs info occs upd args body'
269 stgCseTopLvlRhs _ (StgRhsCon ccs dataCon args)
270 = StgRhsCon ccs dataCon args
271
272 ------------------------------
273 -- The actual AST traversal --
274 ------------------------------
275
276 -- Trivial cases
277 stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
278 stgCseExpr env (StgApp fun args)
279 = StgApp fun' args'
280 where fun' = substVar env fun
281 args' = substArgs env args
282 stgCseExpr _ (StgLit lit)
283 = StgLit lit
284 stgCseExpr env (StgOpApp op args tys)
285 = StgOpApp op args' tys
286 where args' = substArgs env args
287 stgCseExpr _ (StgLam _ _)
288 = pprPanic "stgCseExp" (text "StgLam")
289 stgCseExpr env (StgTick tick body)
290 = let body' = stgCseExpr env body
291 in StgTick tick body'
292 stgCseExpr env (StgCase scrut bndr ty alts)
293 = StgCase scrut' bndr' ty alts'
294 where
295 scrut' = stgCseExpr env scrut
296 (env1, bndr') = substBndr env bndr
297 cse_bndr | StgApp trivial_scrut [] <- scrut' = trivial_scrut
298 -- See Note [Trivial case scrutinee]
299 | otherwise = bndr'
300 alts' = map (stgCseAlt env1 cse_bndr) alts
301
302
303 -- A constructor application.
304 -- To be removed by a variable use when found in the CSE environment
305 stgCseExpr env (StgConApp dataCon args tys)
306 | Just bndr' <- envLookup dataCon args' env
307 = StgApp bndr' []
308 | otherwise
309 = StgConApp dataCon args' tys
310 where args' = substArgs env args
311
312 -- Let bindings
313 -- The binding might be removed due to CSE (we do not want trivial bindings on
314 -- the STG level), so use the smart constructor `mkStgLet` to remove the binding
315 -- if empty.
316 stgCseExpr env (StgLet binds body)
317 = let (binds', env') = stgCseBind env binds
318 body' = stgCseExpr env' body
319 in mkStgLet StgLet binds' body'
320 stgCseExpr env (StgLetNoEscape binds body)
321 = let (binds', env') = stgCseBind env binds
322 body' = stgCseExpr env' body
323 in mkStgLet StgLetNoEscape binds' body'
324
325 -- Case alternatives
326 -- Extend the CSE environment
327 stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
328 stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
329 = let (env1, args') = substBndrs env args
330 env2 = addDataCon case_bndr dataCon (map StgVarArg args') env1
331 -- see note [Case 2: CSEing case binders]
332 rhs' = stgCseExpr env2 rhs
333 in (DataAlt dataCon, args', rhs')
334 stgCseAlt env _ (altCon, args, rhs)
335 = let (env1, args') = substBndrs env args
336 rhs' = stgCseExpr env1 rhs
337 in (altCon, args', rhs')
338
339 -- Bindings
340 stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
341 stgCseBind env (StgNonRec b e)
342 = let (env1, b') = substBndr env b
343 in case stgCseRhs env1 b' e of
344 (Nothing, env2) -> (Nothing, env2)
345 (Just (b2,e'), env2) -> (Just (StgNonRec b2 e'), env2)
346 stgCseBind env (StgRec pairs)
347 = let (env1, pairs1) = substPairs env pairs
348 in case stgCsePairs env1 pairs1 of
349 ([], env2) -> (Nothing, env2)
350 (pairs2, env2) -> (Just (StgRec pairs2), env2)
351
352 stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
353 stgCsePairs env [] = ([], env)
354 stgCsePairs env0 ((b,e):pairs)
355 = let (pairMB, env1) = stgCseRhs env0 b e
356 (pairs', env2) = stgCsePairs env1 pairs
357 in (pairMB `mbCons` pairs', env2)
358 where
359 mbCons = maybe id (:)
360
361 -- The RHS of a binding.
362 -- If it is an constructor application, either short-cut it or extend the environment
363 stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
364 stgCseRhs env bndr (StgRhsCon ccs dataCon args)
365 | Just other_bndr <- envLookup dataCon args' env
366 = let env' = addSubst bndr other_bndr env
367 in (Nothing, env')
368 | otherwise
369 = let env' = addDataCon bndr dataCon args' env
370 -- see note [Case 1: CSEing allocated closures]
371 pair = (bndr, StgRhsCon ccs dataCon args')
372 in (Just pair, env')
373 where args' = substArgs env args
374 stgCseRhs env bndr (StgRhsClosure ccs info occs upd args body)
375 = let (env1, args') = substBndrs env args
376 env2 = forgetCse env1 -- See note [Free variables of an StgClosure]
377 body' = stgCseExpr env2 body
378 in (Just (substVar env bndr, StgRhsClosure ccs info occs' upd args' body'), env)
379 where occs' = substVars env occs
380
381 -- Utilities
382
383 -- | This function short-cuts let-bindings that are now obsolete
384 mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
385 mkStgLet _ Nothing body = body
386 mkStgLet stgLet (Just binds) body = stgLet binds body
387
388
389 {-
390 Note [Trivial case scrutinee]
391 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
392 If we find
393
394 case x as b of { Con a -> … }
395
396 we really want to replace uses of Con a in the body with x, and not just b, in
397 order to handle nested reconstruction of constructors as in
398
399 nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
400 nested (Right (Right x)) = Right (Right x)
401 nested _ = Left True
402
403 Therefore, we add
404 Con a ↦ x
405 to the ConAppMap respectively.
406 Compare Note [CSE for case expressions] in CSE.hs, which does the same for Core CSE.
407
408 If we find
409 case foo x as b of { Con a -> … }
410 we use
411 Con a ↦ b
412
413 Note [Free variables of an StgClosure]
414 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
415 StgClosures (function and thunks) have an explicit list of free variables:
416
417 foo [x] =
418 let not_a_free_var = Left [x]
419 let a_free_var = Right [x]
420 let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
421 in closure
422
423 If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
424 then the list of free variables would be wrong, so for now, we do not CSE
425 across such a closure, simply because I (Joachim) was not sure about possible
426 knock-on effects. If deemed safe and worth the slight code complication of
427 re-calculating this list during or after this pass, this can surely be done.
428 -}