CSE.hs 26.6 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3
{-
(c) The AQUA Project, Glasgow University, 1993-1998

4
\section{Common subexpression}
Austin Seipp's avatar
Austin Seipp committed
5
-}
6

7 8
{-# LANGUAGE CPP #-}

9 10 11
{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns   #-}

Joachim Breitner's avatar
Joachim Breitner committed
12
module CSE (cseProgram, cseOneExpr) where
13 14 15

#include "HsVersions.h"

16 17
import GhcPrelude

18
import CoreSubst
19
import Var              ( Var )
20
import VarEnv           ( elemInScopeSet, mkInScopeSet )
21
import Id               ( Id, idType, isDeadBinder, idHasRules
22
                        , idInlineActivation, setInlineActivation
23
                        , zapIdOccInfo, zapIdUsageInfo, idInlinePragma
24
                        , isJoinId, isJoinId_maybe )
25
import CoreUtils        ( mkAltExpr, eqExpr
26
                        , exprIsTickedString
27
                        , stripTicksE, stripTicksT, mkTicks )
28
import CoreFVs          ( exprFreeVars )
29
import Type             ( tyConAppArgs )
30 31
import CoreSyn
import Outputable
32
import BasicTypes
33
import CoreMap
David Feuer's avatar
David Feuer committed
34
import Util             ( filterOut )
35
import Data.List        ( mapAccumL )
36

Austin Seipp's avatar
Austin Seipp committed
37
{-
38 39
                        Simple common sub-expression
                        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
40
When we see
41 42
        x1 = C a b
        x2 = C x1 b
43
we build up a reverse mapping:   C a b  -> x1
44
                                 C x1 b -> x2
45 46 47
and apply that to the rest of the program.

When we then see
48 49
        y1 = C a b
        y2 = C y1 b
50
we replace the C a b with x1.  But then we *dont* want to
51
add   x1 -> y1  to the mapping.  Rather, we want the reverse, y1 -> x1
52
so that a subsequent binding
53 54
        y2 = C y1 b
will get transformed to C x1 b, and then to x2.
55

56
So we carry an extra var->var substitution which we apply *before* looking up in the
57 58 59
reverse mapping.


simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
60 61
Note [Shadowing]
~~~~~~~~~~~~~~~~
62
We have to be careful about shadowing.
63
For example, consider
64 65 66
        f = \x -> let y = x+x in
                      h = \x -> x+x
                  in ...
67

68 69 70 71
Here we must *not* do CSE on the inner x+x!  The simplifier used to guarantee no
shadowing, but it doesn't any more (it proved too hard), so we clone as we go.
We can simply add clones to the substitution already described.

72

73 74
Note [CSE for bindings]
~~~~~~~~~~~~~~~~~~~~~~~
75 76
Let-bindings have two cases, implemented by addBinding.

77
* SUBSTITUTE: applies when the RHS is a variable
78 79 80 81 82 83

     let x = y in ...(h x)....

  Here we want to extend the /substitution/ with x -> y, so that the
  (h x) in the body might CSE with an enclosing (let v = h y in ...).
  NB: the substitution maps InIds, so we extend the substitution with
84
      a binding for the original InId 'x'
85

86
  How can we have a variable on the RHS? Doesn't the simplifier inline them?
87 88 89

    - First, the original RHS might have been (g z) which has CSE'd
      with an enclosing (let y = g z in ...).  This is super-important.
90
      See #5996:
91 92 93 94 95 96 97
         x1 = C a b
         x2 = C x1 b
         y1 = C a b
         y2 = C y1 b
      Here we CSE y1's rhs to 'x1', and then we must add (y1->x1) to
      the substitution so that we can CSE the binding for y2.

98
    - Second, we use addBinding for case expression scrutinees too;
99
      see Note [CSE for case expressions]
100

101 102
* EXTEND THE REVERSE MAPPING: applies in all other cases

103
     let x = h y in ...(h y)...
104

105 106
  Here we want to extend the /reverse mapping (cs_map)/ so that
  we CSE the (h y) call to x.
107

108 109
  Note that we use EXTEND even for a trivial expression, provided it
  is not a variable or literal. In particular this /includes/ type
110
  applications. This can be important (#13156); e.g.
111 112 113 114 115 116 117 118
     case f @ Int of { r1 ->
     case f @ Int of { r2 -> ...
  Here we want to common-up the two uses of (f @ Int) so we can
  remove one of the case expressions.

  See also Note [Corner case for case expressions] for another
  reason not to use SUBSTITUTE for all trivial expressions.

119
Notice that
120 121
  - The SUBSTITUTE situation extends the substitution (cs_subst)
  - The EXTEND situation extends the reverse mapping (cs_map)
122

123
Notice also that in the SUBSTITUTE case we leave behind a binding
124 125 126 127 128 129 130 131
  x = y
even though we /also/ carry a substitution x -> y.  Can we just drop
the binding instead?  Well, not at top level! See SimplUtils
Note [Top level and postInlineUnconditionally]; and in any case CSE
applies only to the /bindings/ of the program, and we leave it to the
simplifier to propate effects to the RULES.  Finally, it doesn't seem
worth the effort to discard the nested bindings because the simplifier
will do it next.
132 133 134

Note [CSE for case expressions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
135
Consider
136 137 138
  case scrut_expr of x { ...alts... }
This is very like a strict let-binding
  let !x = scrut_expr in ...
139
So we use (addBinding x scrut_expr) to process scrut_expr and x, and as a
140 141 142 143 144 145 146 147
result all the stuff under Note [CSE for bindings] applies directly.

For example:

* Trivial scrutinee
     f = \x -> case x of wild {
                 (a:as) -> case a of wild1 {
                             (p,q) -> ...(wild1:as)...
148

149 150 151 152 153
  Here, (wild1:as) is morally the same as (a:as) and hence equal to
  wild. But that's not quite obvious.  In the rest of the compiler we
  want to keep it as (wild1:as), but for CSE purpose that's a bad
  idea.

154
  By using addBinding we add the binding (wild1 -> a) to the substitution,
155 156 157 158 159
  which does exactly the right thing.

  (Notice this is exactly backwards to what the simplifier does, which
  is to try to replaces uses of 'a' with uses of 'wild1'.)

160
  This is the main reason that addBinding is called with a trivial rhs.
161 162

* Non-trivial scrutinee
163
     case (f x) of y { pat -> ...let z = f x in ... }
164

165
  By using addBinding we'll add (f x :-> y) to the cs_map, and
166
  thereby CSE the inner (f x) to y.
167

168 169
Note [CSE for INLINE and NOINLINE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
170 171 172
There are some subtle interactions of CSE with functions that the user
has marked as INLINE or NOINLINE. (Examples from Roman Leshchinskiy.)
Consider
173

174
        yes :: Int  {-# NOINLINE yes #-}
175
        yes = undefined
176

177
        no :: Int   {-# NOINLINE no #-}
178
        no = undefined
179

180
        foo :: Int -> Int -> Int  {-# NOINLINE foo #-}
181
        foo m n = n
182

183
        {-# RULES "foo/no" foo no = id #-}
184

185 186
        bar :: Int -> Int
        bar = foo yes
187

188
We do not expect the rule to fire.  But if we do CSE, then we risk
Gabor Greif's avatar
Gabor Greif committed
189
getting yes=no, and the rule does fire.  Actually, it won't because
190
NOINLINE means that 'yes' will never be inlined, not even if we have
Gabor Greif's avatar
Gabor Greif committed
191
yes=no.  So that's fine (now; perhaps in the olden days, yes=no would
Simon Peyton Jones's avatar
Simon Peyton Jones committed
192
have substituted even if 'yes' was NOINLINE).
193

194
But we do need to take care.  Consider
195

196 197
        {-# NOINLINE bar #-}
        bar = <rhs>     -- Same rhs as foo
198

199 200
        foo = <rhs>

201
If CSE produces
202
        foo = bar
203 204 205 206 207 208 209 210 211 212
then foo will never be inlined to <rhs> (when it should be, if <rhs>
is small).  The conclusion here is this:

   We should not add
       <rhs> :-> bar
  to the CSEnv if 'bar' has any constraints on when it can inline;
  that is, if its 'activation' not always active.  Otherwise we
  might replace <rhs> by 'bar', and then later be unable to see that it
  really was <rhs>.

213 214 215 216
An except to the rule is when the INLINE pragma is not from the user, e.g. from
WorkWrap (see Note [Wrapper activation]). We can tell because noUserInlineSpec
is then true.

217
Note that we do not (currently) do CSE on the unfolding stored inside
Gabor Greif's avatar
Gabor Greif committed
218
an Id, even if it is a 'stable' unfolding.  That means that when an
219 220 221
unfolding happens, it is always faithful to what the stable unfolding
originally was.

222 223
Note [CSE for stable unfoldings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
224
Consider
225 226 227 228 229
   {-# Unf = Stable (\pq. build blah) #-}
   foo = x

Here 'foo' has a stable unfolding, but its (optimised) RHS is trivial.
(Turns out that this actually happens for the enumFromTo method of
rwbarton's avatar
rwbarton committed
230 231 232 233
the Integer instance of Enum in GHC.Enum.)  Suppose moreover that foo's
stable unfolding originates from an INLINE or INLINEABLE pragma on foo.
Then we obviously do NOT want to extend the substitution with (foo->x),
because we promised to inline foo as what the user wrote.  See similar
234 235 236 237 238 239 240 241 242 243 244 245 246
SimplUtils Note [Stable unfoldings and postInlineUnconditionally].

Nor do we want to change the reverse mapping. Suppose we have

   {-# Unf = Stable (\pq. build blah) #-}
   foo = <expr>
   bar = <expr>

There could conceivably be merit in rewriting the RHS of bar:
   bar = foo
but now bar's inlining behaviour will change, and importing
modules might see that.  So it seems dodgy and we don't do it.

rwbarton's avatar
rwbarton committed
247 248 249 250 251
Stable unfoldings are also created during worker/wrapper when we decide
that a function's definition is so small that it should always inline.
In this case we still want to do CSE (#13340). Hence the use of
isAnyInlinePragma rather than isStableUnfolding.

252 253
Note [Corner case for case expressions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
254 255
Here is another reason that we do not use SUBSTITUTE for
all trivial expressions. Consider
256 257
   case x |> co of (y::Array# Int) { ... }

258
We do not want to extend the substitution with (y -> x |> co); since y
259
is of unlifted type, this would destroy the let/app invariant if (x |>
260
co) was not ok-for-speculation.
261

262
But surely (x |> co) is ok-for-speculation, because it's a trivial
263 264 265
expression, and x's type is also unlifted, presumably.  Well, maybe
not if you are using unsafe casts.  I actually found a case where we
had
266 267
   (x :: HValue) |> (UnsafeCo :: HValue ~ Array# Int)

lukemaurer's avatar
lukemaurer committed
268 269 270 271 272 273 274 275 276 277 278
Note [CSE for join points?]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
We must not be naive about join points in CSE:
   join j = e in
   if b then jump j else 1 + e
The expression (1 + jump j) is not good (see Note [Invariants on join points] in
CoreSyn). This seems to come up quite seldom, but it happens (first seen
compiling ppHtml in Haddock.Backends.Xhtml).

We could try and be careful by tracking which join points are still valid at
each subexpression, but since join points aren't allocated or shared, there's
279 280
less to gain by trying to CSE them. (#13219)

Simon Peyton Jones's avatar
Simon Peyton Jones committed
281 282
Note [Look inside join-point binders]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
283 284 285 286 287 288 289 290 291 292 293 294 295 296 297
Another way how CSE for joint points is tricky is

  let join foo x = (x, 42)
      join bar x = (x, 42)
  in … jump foo 1 … jump bar 2 …

naively, CSE would turn this into

  let join foo x = (x, 42)
      join bar = foo
  in … jump foo 1 … jump bar 2 …

but now bar is a join point that claims arity one, but its right-hand side
is not a lambda, breaking the join-point invariant (this was #15002).

Simon Peyton Jones's avatar
Simon Peyton Jones committed
298 299 300
So `cse_bind` must zoom past the lambdas of a join point (using
`collectNBinders`) and resume searching for CSE opportunities only in
the body of the join point.
301

302 303 304 305 306 307 308 309 310 311 312 313 314 315
Note [CSE for recursive bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  f = \x ... f....
  g = \y ... g ...
where the "..." are identical.  Could we CSE them?  In full generality
with mutual recursion it's quite hard; but for self-recursive bindings
(which are very common) it's rather easy:

* Maintain a separate cs_rec_map, that maps
      (\f. (\x. ...f...) ) -> f
  Note the \f in the domain of the mapping!

* When we come across the binding for 'g', look up (\g. (\y. ...g...))
Ben Gamari's avatar
Ben Gamari committed
316
  Bingo we get a hit.  So we can replace the 'g' binding with
317 318 319 320 321 322
     g = f

We can't use cs_map for this, because the key isn't an expression of
the program; it's a kind of synthetic key for recursive bindings.


Austin Seipp's avatar
Austin Seipp committed
323 324
************************************************************************
*                                                                      *
325
\section{Common subexpression}
Austin Seipp's avatar
Austin Seipp committed
326 327 328
*                                                                      *
************************************************************************
-}
329

330
cseProgram :: CoreProgram -> CoreProgram
331
cseProgram binds = snd (mapAccumL (cseBind TopLevel) emptyCSEnv binds)
332

333
cseBind :: TopLevelFlag -> CSEnv -> CoreBind -> (CSEnv, CoreBind)
334
cseBind toplevel env (NonRec b e)
335
  = (env2, NonRec b2 e2)
336
  where
337 338
    (env1, b1)       = addBinder env b
    (env2, (b2, e2)) = cse_bind toplevel env1 (b,e) b1
339

340
cseBind toplevel env (Rec [(in_id, rhs)])
341 342 343 344 345 346
  | noCSE in_id
  = (env1, Rec [(out_id, rhs')])

  -- See Note [CSE for recursive bindings]
  | Just previous <- lookupCSRecEnv env out_id rhs''
  , let previous' = mkTicks ticks previous
347 348 349
        out_id'   = delayInlining toplevel out_id
  = -- We have a hit in the recursive-binding cache
    (extendCSSubst env1 in_id previous', NonRec out_id' previous')
350 351 352 353 354 355 356 357 358 359 360 361

  | otherwise
  = (extendCSRecEnv env1 out_id rhs'' id_expr', Rec [(zapped_id, rhs')])

  where
    (env1, [out_id]) = addRecBinders env [in_id]
    rhs'  = cseExpr env1 rhs
    rhs'' = stripTicksE tickishFloatable rhs'
    ticks = stripTicksT tickishFloatable rhs'
    id_expr'  = varToCoreExpr out_id
    zapped_id = zapIdUsageInfo out_id

362
cseBind toplevel env (Rec pairs)
363
  = (env2, Rec pairs')
364
  where
365 366 367 368 369
    (env1, bndrs1) = addRecBinders env (map fst pairs)
    (env2, pairs') = mapAccumL do_one env1 (zip pairs bndrs1)

    do_one env (pr, b1) = cse_bind toplevel env pr b1

370 371 372 373 374
-- | Given a binding of @in_id@ to @in_rhs@, and a fresh name to refer
-- to @in_id@ (@out_id@, created from addBinder or addRecBinders),
-- first try to CSE @in_rhs@, and then add the resulting (possibly CSE'd)
-- binding to the 'CSEnv', so that we attempt to CSE any expressions
-- which are equal to @out_rhs@.
375 376
cse_bind :: TopLevelFlag -> CSEnv -> (InId, InExpr) -> OutId -> (CSEnv, (OutId, OutExpr))
cse_bind toplevel env (in_id, in_rhs) out_id
377
  | isTopLevel toplevel, exprIsTickedString in_rhs
378
      -- See Note [Take care with literal strings]
Simon Peyton Jones's avatar
Simon Peyton Jones committed
379
  = (env', (out_id', in_rhs))
380

381
  | Just arity <- isJoinId_maybe in_id
Simon Peyton Jones's avatar
Simon Peyton Jones committed
382
      -- See Note [Look inside join-point binders]
383 384 385 386 387
  = let (params, in_body) = collectNBinders arity in_rhs
        (env', params') = addBinders env params
        out_body = tryForCSE env' in_body
    in (env, (out_id, mkLams params' out_body))

388
  | otherwise
389
  = (env', (out_id'', out_rhs))
390 391
  where
    (env', out_id') = addBinding env in_id out_id out_rhs
392 393 394 395 396 397
    (cse_done, out_rhs) = try_for_cse env in_rhs
    out_id'' | cse_done  = delayInlining toplevel out_id'
             | otherwise = out_id'

delayInlining :: TopLevelFlag -> Id -> Id
-- Add a NOINLINE[2] if the Id doesn't have an INLNE pragma already
398
-- See Note [Delay inlining after CSE]
399 400 401
delayInlining top_lvl bndr
  | isTopLevel top_lvl
  , isAlwaysActive (idInlineActivation bndr)
402 403 404 405 406
  , idHasRules bndr  -- Only if the Id has some RULES,
                     -- which might otherwise get lost
       -- These rules are probably auto-generated specialisations,
       -- since Ids with manual rules usually have manually-inserted
       -- delayed inlining anyway
407 408 409
  = bndr `setInlineActivation` activeAfterInitial
  | otherwise
  = bndr
410 411

addBinding :: CSEnv                      -- Includes InId->OutId cloning
Simon Peyton Jones's avatar
Simon Peyton Jones committed
412
           -> InVar                      -- Could be a let-bound type
413
           -> OutId -> OutExpr           -- Processed binding
414
           -> (CSEnv, OutId)             -- Final env, final bndr
415 416
-- Extend the CSE env with a mapping [rhs -> out-id]
-- unless we can instead just substitute [in-id -> rhs]
Simon Peyton Jones's avatar
Simon Peyton Jones committed
417 418 419
--
-- It's possible for the binder to be a type variable (see
-- Note [Type-let] in CoreSyn), in which case we can just substitute.
420
addBinding env in_id out_id rhs'
Simon Peyton Jones's avatar
Simon Peyton Jones committed
421 422 423 424
  | not (isId in_id) = (extendCSSubst env in_id rhs',     out_id)
  | noCSE in_id      = (env,                              out_id)
  | use_subst        = (extendCSSubst env in_id rhs',     out_id)
  | otherwise        = (extendCSEnv env rhs' id_expr', zapped_id)
425
  where
426 427 428
    id_expr'  = varToCoreExpr out_id
    zapped_id = zapIdUsageInfo out_id
       -- Putting the Id into the cs_map makes it possible that
429
       -- it'll become shared more than it is now, which would
430
       -- invalidate (the usage part of) its demand info.
431
       --    This caused #100218.
432 433 434 435 436 437
       -- Easiest thing is to zap the usage info; subsequently
       -- performing late demand-analysis will restore it.  Don't zap
       -- the strictness info; it's not necessary to do so, and losing
       -- it is bad for performance if you don't do late demand
       -- analysis

438
    -- Should we use SUBSTITUTE or EXTEND?
439
    -- See Note [CSE for bindings]
440 441 442
    use_subst = case rhs' of
                   Var {} -> True
                   _      -> False
443

444 445
-- | Given a binder `let x = e`, this function
-- determines whether we should add `e -> x` to the cs_map
446
noCSE :: InId -> Bool
447 448
noCSE id =  not (isAlwaysActive (idInlineActivation id)) &&
            not (noUserInlineSpec (inlinePragmaSpec (idInlinePragma id)))
449 450 451 452 453 454 455
             -- See Note [CSE for INLINE and NOINLINE]
         || isAnyInlinePragma (idInlinePragma id)
             -- See Note [CSE for stable unfoldings]
         || isJoinId id
             -- See Note [CSE for join points?]


456 457
{- Note [Take care with literal strings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479
Consider this example:

  x = "foo"#
  y = "foo"#
  ...x...y...x...y....

We would normally turn this into:

  x = "foo"#
  y = x
  ...x...x...x...x....

But this breaks an invariant of Core, namely that the RHS of a top-level binding
of type Addr# must be a string literal, not another variable. See Note
[CoreSyn top-level string literals] in CoreSyn.

For this reason, we special case top-level bindings to literal strings and leave
the original RHS unmodified. This produces:

  x = "foo"#
  y = "foo"#
  ...x...x...x...x....
480 481 482 483 484 485 486

Now 'y' will be discarded as dead code, and we are done.

The net effect is that for the y-binding we want to
  - Use SUBSTITUTE, by extending the substitution with  y :-> x
  - but leave the original binding for y undisturbed

487
This is done by cse_bind.  I got it wrong the first time (#13367).
488 489 490

Note [Delay inlining after CSE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
491
Suppose (#15445) we have
492 493 494 495 496 497 498 499 500 501 502 503 504 505
   f,g :: Num a => a -> a
   f x = ...f (x-1).....
   g y = ...g (y-1) ....

and we make some specialisations of 'g', either automatically, or via
a SPECIALISE pragma.  Then CSE kicks in and notices that the RHSs of
'f' and 'g' are identical, so we get
   f x = ...f (x-1)...
   g = f
   {-# RULES g @Int _ = $sg #-}

Now there is terrible danger that, in an importing module, we'll inline
'g' before we have a chance to run its specialisation!

506
Solution: during CSE, after a "hit" in the CSE cache
507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545
  * when adding a binding
        g = f
  * for a top-level function g
  * and g has specialisation RULES
add a NOINLINE[2] activation to it, to ensure it's not inlined
right away.

Notes:
* Why top level only?  Because for nested bindings we are already past
  phase 2 and will never return there.

* Why "only if g has RULES"?  Because there is no point in
  doing this if there are no RULES; and other things being
  equal it delays optimisation to delay inlining (#17409)


---- Historical note ---

This patch is simpler and more direct than an earlier
version:

  commit 2110738b280543698407924a16ac92b6d804dc36
  Author: Simon Peyton Jones <simonpj@microsoft.com>
  Date:   Mon Jul 30 13:43:56 2018 +0100

  Don't inline functions with RULES too early

We had to revert this patch because it made GHC itself slower.

Why? It delayed inlining of /all/ functions with RULES, and that was
very bad in TcFlatten.flatten_ty_con_app

* It delayed inlining of liftM
* That delayed the unravelling of the recursion in some dictionary
  bindings.
* That delayed some eta expansion, leaving
     flatten_ty_con_app = \x y. let <stuff> in \z. blah
* That allowed the float-out pass to put sguff between
  the \y and \z.
546
* And that permanently stopped eta expansion of the function,
547
  even once <stuff> was simplified.
548

549 550
-}

551
tryForCSE :: CSEnv -> InExpr -> OutExpr
552 553 554 555 556 557 558 559 560 561 562 563
tryForCSE env expr = snd (try_for_cse env expr)

try_for_cse :: CSEnv -> InExpr -> (Bool, OutExpr)
-- (False, e') => We did not CSE the entire expression,
--                but we might have CSE'd some sub-expressions,
--                yielding e'
--
-- (True, te') => We CSE'd the entire expression,
--                yielding the trivial expression te'
try_for_cse env expr
  | Just e <- lookupCSEnv env expr'' = (True,  mkTicks ticks e)
  | otherwise                        = (False, expr')
564 565 566 567
    -- The varToCoreExpr is needed if we have
    --   case e of xco { ...case e of yco { ... } ... }
    -- Then CSE will substitute yco -> xco;
    -- but these are /coercion/ variables
568
  where
569
    expr'  = cseExpr env expr
570
    expr'' = stripTicksE tickishFloatable expr'
571 572 573 574 575
    ticks  = stripTicksT tickishFloatable expr'
    -- We don't want to lose the source notes when a common sub
    -- expression gets eliminated. Hence we push all (!) of them on
    -- top of the replaced sub-expression. This is probably not too
    -- useful in practice, but upholds our semantics.
576

577 578 579 580
-- | Runs CSE on a single expression.
--
-- This entry point is not used in the compiler itself, but is provided
-- as a convenient entry point for users of the GHC API.
581
cseOneExpr :: InExpr -> OutExpr
582 583
cseOneExpr e = cseExpr env e
  where env = emptyCSEnv {cs_subst = mkEmptySubst (mkInScopeSet (exprFreeVars e)) }
Joachim Breitner's avatar
Joachim Breitner committed
584

585
cseExpr :: CSEnv -> InExpr -> OutExpr
586 587 588 589
cseExpr env (Type t)              = Type (substTy (csEnvSubst env) t)
cseExpr env (Coercion c)          = Coercion (substCo (csEnvSubst env) c)
cseExpr _   (Lit lit)             = Lit lit
cseExpr env (Var v)               = lookupSubst env v
590
cseExpr env (App f a)             = App (cseExpr env f) (tryForCSE env a)
591
cseExpr env (Tick t e)            = Tick t (cseExpr env e)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
592
cseExpr env (Cast e co)           = Cast (tryForCSE env e) (substCo (csEnvSubst env) co)
593 594
cseExpr env (Lam b e)             = let (env', b') = addBinder env b
                                    in Lam b' (cseExpr env' e)
595
cseExpr env (Let bind e)          = let (env', bind') = cseBind NotTopLevel env bind
596 597
                                    in Let bind' (cseExpr env' e)
cseExpr env (Case e bndr ty alts) = cseCase env e bndr ty alts
598 599 600

cseCase :: CSEnv -> InExpr -> InId -> InType -> [InAlt] -> OutExpr
cseCase env scrut bndr ty alts
601 602
  = Case scrut1 bndr3 ty' $
    combineAlts alt_env (map cse_alt alts)
603
  where
604
    ty' = substTy (csEnvSubst env) ty
605
    scrut1 = tryForCSE env scrut
606

607 608 609 610
    bndr1 = zapIdOccInfo bndr
      -- Zapping the OccInfo is needed because the extendCSEnv
      -- in cse_alt may mean that a dead case binder
      -- becomes alive, and Lint rejects that
611 612
    (env1, bndr2)    = addBinder env bndr1
    (alt_env, bndr3) = addBinding env1 bndr bndr2 scrut1
613
         -- addBinding: see Note [CSE for case expressions]
614

615 616
    con_target :: OutExpr
    con_target = lookupSubst alt_env bndr
617

618 619
    arg_tys :: [OutType]
    arg_tys = tyConAppArgs (idType bndr3)
620

621 622
    -- Given case x of { K y z -> ...K y z... }
    -- CSE K y z into x...
623
    cse_alt (DataAlt con, args, rhs)
624
        | not (null args)
625
                -- ... but don't try CSE if there are no args; it just increases the number
626 627 628 629
                -- of live vars.  E.g.
                --      case x of { True -> ....True.... }
                -- Don't replace True by x!
                -- Hence the 'null args', which also deal with literals and DEFAULT
630
        = (DataAlt con, args', tryForCSE new_env rhs)
631 632
        where
          (env', args') = addBinders alt_env args
633 634
          new_env       = extendCSEnv env' con_expr con_target
          con_expr      = mkAltExpr (DataAlt con) args' arg_tys
635 636

    cse_alt (con, args, rhs)
637
        = (con, args', tryForCSE env' rhs)
638 639
        where
          (env', args') = addBinders alt_env args
640

641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656
combineAlts :: CSEnv -> [InAlt] -> [InAlt]
-- See Note [Combine case alternatives]
combineAlts env ((_,bndrs1,rhs1) : rest_alts)
  | all isDeadBinder bndrs1
  = (DEFAULT, [], rhs1) : filtered_alts
  where
    in_scope = substInScope (csEnvSubst env)
    filtered_alts = filterOut identical rest_alts
    identical (_con, bndrs, rhs) = all ok bndrs && eqExpr in_scope rhs1 rhs
    ok bndr = isDeadBinder bndr || not (bndr `elemInScopeSet` in_scope)

combineAlts _ alts = alts  -- Default case

{- Note [Combine case alternatives]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
combineAlts is just a more heavyweight version of the use of
657
combineIdenticalAlts in SimplUtils.prepareAlts.  The basic idea is
658 659 660 661 662 663 664 665 666 667
to transform

    DEFAULT -> e1
    K x     -> e1
    W y z   -> e2
===>
   DEFAULT -> e1
   W y z   -> e2

In the simplifier we use cheapEqExpr, because it is called a lot.
Gabor Greif's avatar
Gabor Greif committed
668
But here in CSE we use the full eqExpr.  After all, two alternatives usually
669
differ near the root, so it probably isn't expensive to compare the full
670
alternative.  It seems like the same kind of thing that CSE is supposed
671 672
to be doing, which is why I put it here.

673
I actually saw some examples in the wild, where some inlining made e1 too
674 675 676
big for cheapEqExpr to catch it.


Austin Seipp's avatar
Austin Seipp committed
677 678
************************************************************************
*                                                                      *
679
\section{The CSE envt}
Austin Seipp's avatar
Austin Seipp committed
680 681 682
*                                                                      *
************************************************************************
-}
683

684 685 686 687
data CSEnv
  = CS { cs_subst :: Subst  -- Maps InBndrs to OutExprs
            -- The substitution variables to
            -- /trivial/ OutExprs, not arbitrary expressions
688

689 690 691
       , cs_map   :: CoreMap OutExpr   -- The reverse mapping
            -- Maps a OutExpr to a /trivial/ OutExpr
            -- The key of cs_map is stripped of all Ticks
692 693 694

       , cs_rec_map :: CoreMap OutExpr
            -- See Note [CSE for recursive bindings]
695
       }
696 697

emptyCSEnv :: CSEnv
698 699
emptyCSEnv = CS { cs_map = emptyCoreMap, cs_rec_map = emptyCoreMap
                , cs_subst = emptySubst }
700

701
lookupCSEnv :: CSEnv -> OutExpr -> Maybe OutExpr
702
lookupCSEnv (CS { cs_map = csmap }) expr
703 704 705 706 707 708 709
  = lookupCoreMap csmap expr

extendCSEnv :: CSEnv -> OutExpr -> OutExpr -> CSEnv
extendCSEnv cse expr triv_expr
  = cse { cs_map = extendCoreMap (cs_map cse) sexpr triv_expr }
  where
    sexpr = stripTicksE tickishFloatable expr
710

711 712 713
extendCSRecEnv :: CSEnv -> OutId -> OutExpr -> OutExpr -> CSEnv
-- See Note [CSE for recursive bindings]
extendCSRecEnv cse bndr expr triv_expr
Ben Gamari's avatar
Ben Gamari committed
714
  = cse { cs_rec_map = extendCoreMap (cs_rec_map cse) (Lam bndr expr) triv_expr }
715 716 717 718 719 720

lookupCSRecEnv :: CSEnv -> OutId -> OutExpr -> Maybe OutExpr
-- See Note [CSE for recursive bindings]
lookupCSRecEnv (CS { cs_rec_map = csmap }) bndr expr
  = lookupCoreMap csmap (Lam bndr expr)

721 722 723
csEnvSubst :: CSEnv -> Subst
csEnvSubst = cs_subst

724
lookupSubst :: CSEnv -> Id -> OutExpr
725
lookupSubst (CS { cs_subst = sub}) x = lookupIdSubst (text "CSE.lookupSubst") sub x
726

727 728
extendCSSubst :: CSEnv -> Id  -> CoreExpr -> CSEnv
extendCSSubst cse x rhs = cse { cs_subst = extendSubst (cs_subst cse) x rhs }
729

730 731 732
-- | Add clones to the substitution to deal with shadowing.  See
-- Note [Shadowing] for more details.  You should call this whenever
-- you go under a binder.
733
addBinder :: CSEnv -> Var -> (CSEnv, Var)
734
addBinder cse v = (cse { cs_subst = sub' }, v')
735 736
                where
                  (sub', v') = substBndr (cs_subst cse) v
737 738

addBinders :: CSEnv -> [Var] -> (CSEnv, [Var])
739
addBinders cse vs = (cse { cs_subst = sub' }, vs')
740 741
                where
                  (sub', vs') = substBndrs (cs_subst cse) vs
742

743
addRecBinders :: CSEnv -> [Id] -> (CSEnv, [Id])
744
addRecBinders cse vs = (cse { cs_subst = sub' }, vs')
745 746
                where
                  (sub', vs') = substRecBndrs (cs_subst cse) vs