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