CorePrep.hs 56.6 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3
{-
(c) The University of Glasgow, 1994-2006

Simon Marlow's avatar
Simon Marlow committed
4 5

Core pass to saturate constructors and PrimOps
Austin Seipp's avatar
Austin Seipp committed
6
-}
7

8
{-# LANGUAGE BangPatterns, CPP, MultiWayIf #-}
Ian Lynagh's avatar
Ian Lynagh committed
9

10
module CorePrep (
11
      corePrepPgm, corePrepExpr, cvtLitInteger,
12
      lookupMkIntegerName, lookupIntegerSDataConName
13 14 15 16
  ) where

#include "HsVersions.h"

17 18
import OccurAnal

19
import HscTypes
20
import PrelNames
21
import MkId             ( realWorldPrimId )
22 23
import CoreUtils
import CoreArity
Simon Marlow's avatar
Simon Marlow committed
24
import CoreFVs
25 26
import CoreMonad        ( CoreToDo(..) )
import CoreLint         ( endPassIO )
27
import CoreSyn
28
import CoreSubst
29
import MkCore hiding( FloatBind(..) )   -- We use our own FloatBind here
Simon Marlow's avatar
Simon Marlow committed
30
import Type
31
import Literal
Simon Marlow's avatar
Simon Marlow committed
32
import Coercion
33
import TcEnv
Simon Marlow's avatar
Simon Marlow committed
34
import TyCon
35
import Demand
Simon Marlow's avatar
Simon Marlow committed
36
import Var
37
import VarSet
38
import VarEnv
Simon Marlow's avatar
Simon Marlow committed
39
import Id
40
import IdInfo
41
import TysWiredIn
Simon Marlow's avatar
Simon Marlow committed
42 43 44
import DataCon
import PrimOp
import BasicTypes
45
import Module
46 47
import UniqSupply
import Maybes
48
import OrdList
49
import ErrUtils
50
import DynFlags
Simon Marlow's avatar
Simon Marlow committed
51
import Util
52
import Pair
53
import Outputable
54
import Platform
55
import FastString
56
import Config
Peter Wortmann's avatar
Peter Wortmann committed
57 58
import Name             ( NamedThing(..), nameSrcSpan )
import SrcLoc           ( SrcSpan(..), realSrcLocSpan, mkRealSrcLoc )
59
import Data.Bits
60
import MonadUtils       ( mapAccumLM )
61
import Data.List        ( mapAccumL )
62
import Control.Monad
63

Austin Seipp's avatar
Austin Seipp committed
64
{-
65 66 67
-- ---------------------------------------------------------------------------
-- Overview
-- ---------------------------------------------------------------------------
68

69
The goal of this pass is to prepare for code generation.
70

71
1.  Saturate constructor and primop applications.
72

73 74
2.  Convert to A-normal form; that is, function arguments
    are always variables.
75

76
    * Use case for strict arguments:
77 78
        f E ==> case E of x -> f x
        (where f is strict)
79

80
    * Use let for non-trivial lazy arguments
81 82
        f E ==> let x = E in f x
        (were f is lazy and x is non-trivial)
83

84
3.  Similarly, convert any unboxed lets into cases.
85
    [I'm experimenting with leaving 'ok-for-speculation'
86
     rhss in let-form right up to this point.]
87

88
4.  Ensure that *value* lambdas only occur as the RHS of a binding
89
    (The code generator can't deal with anything else.)
90
    Type lambdas are ok, however, because the code gen discards them.
91

92
5.  [Not any more; nuked Jun 2002] Do the seq/par munging.
93

94
6.  Clone all local Ids.
95
    This means that all such Ids are unique, rather than the
96 97 98
    weaker guarantee of no clashes which the simplifier provides.
    And that is what the code generator needs.

99
    We don't clone TyVars or CoVars. The code gen doesn't need that,
100
    and doing so would be tiresome because then we'd need
101
    to substitute in types and coercions.
102

103 104 105
7.  Give each dynamic CCall occurrence a fresh unique; this is
    rather like the cloning step above.

106
8.  Inject bindings for the "implicit" Ids:
107 108
        * Constructor wrappers
        * Constructor workers
109 110
    We want curried definitions for all of these in case they
    aren't inlined by some caller.
111

112
9.  Replace (lazy e) by e.  See Note [lazyId magic] in MkId.hs
113
    Also replace (noinline e) by e.
114

115 116
10. Convert (LitInteger i t) into the core representation
    for the Integer i. Normally this uses mkInteger, but if
117 118 119 120
    we are using the integer-gmp implementation then there is a
    special case where we use the S# constructor for Integers that
    are in the range of Int.

121 122 123 124
11. Uphold tick consistency while doing this: We move ticks out of
    (non-type) applications where we can, and make sure that we
    annotate according to scoping rules when floating.

125 126 127
This is all done modulo type applications and abstractions, so that
when type erasure is done for conversion to STG, we don't end up with
any trivial or useless bindings.
128

129

130 131
Note [CorePrep invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~
132
Here is the syntax of the Core produced by CorePrep:
133

134
    Trivial expressions
135 136 137
       arg ::= lit |  var
              | arg ty  |  /\a. arg
              | truv co  |  /\c. arg  |  arg |> co
138 139

    Applications
140
       app ::= lit  |  var  |  app arg  |  app ty  | app co | app |> co
141 142

    Expressions
143
       body ::= app
144 145
              | let(rec) x = rhs in body     -- Boxed only
              | case body of pat -> body
146
              | /\a. body | /\c. body
147 148
              | body |> co

149
    Right hand sides (only place where value lambdas can occur)
150 151 152 153
       rhs ::= /\a.rhs  |  \x.rhs  |  body

We define a synonym for each of these non-terminals.  Functions
with the corresponding name produce a result in that syntax.
Austin Seipp's avatar
Austin Seipp committed
154
-}
155

156
type CpeArg  = CoreExpr    -- Non-terminal 'arg'
157 158 159
type CpeApp  = CoreExpr    -- Non-terminal 'app'
type CpeBody = CoreExpr    -- Non-terminal 'body'
type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'
160

Austin Seipp's avatar
Austin Seipp committed
161 162 163
{-
************************************************************************
*                                                                      *
164
                Top level stuff
Austin Seipp's avatar
Austin Seipp committed
165 166 167
*                                                                      *
************************************************************************
-}
168

169 170 171 172 173 174
corePrepPgm :: HscEnv -> Module -> ModLocation -> CoreProgram -> [TyCon]
            -> IO CoreProgram
corePrepPgm hsc_env this_mod mod_loc binds data_tycons =
    withTiming (pure dflags)
               (text "CorePrep"<+>brackets (ppr this_mod))
               (const ()) $ do
175
    us <- mkSplitUniqSupply 's'
176
    initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
177

Peter Wortmann's avatar
Peter Wortmann committed
178
    let implicit_binds = mkDataConWorkers dflags mod_loc data_tycons
179 180 181 182
            -- NB: we must feed mkImplicitBinds through corePrep too
            -- so that they are suitably cloned and eta-expanded

        binds_out = initUs_ us $ do
183 184
                      floats1 <- corePrepTopBinds initialCorePrepEnv binds
                      floats2 <- corePrepTopBinds initialCorePrepEnv implicit_binds
185 186
                      return (deFloatTop (floats1 `appendFloats` floats2))

187
    endPassIO hsc_env alwaysQualify CorePrep binds_out []
188
    return binds_out
189 190
  where
    dflags = hsc_dflags hsc_env
191

192
corePrepExpr :: DynFlags -> HscEnv -> CoreExpr -> IO CoreExpr
193 194
corePrepExpr dflags hsc_env expr =
    withTiming (pure dflags) (text "CorePrep [expr]") (const ()) $ do
195
    us <- mkSplitUniqSupply 's'
196
    initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
197
    let new_expr = initUs_ us (cpeBodyNF initialCorePrepEnv expr)
198 199
    dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep" (ppr new_expr)
    return new_expr
200

201
corePrepTopBinds :: CorePrepEnv -> [CoreBind] -> UniqSM Floats
202
-- Note [Floating out of top level bindings]
203 204
corePrepTopBinds initialCorePrepEnv binds
  = go initialCorePrepEnv binds
205 206 207 208 209
  where
    go _   []             = return emptyFloats
    go env (bind : binds) = do (env', bind') <- cpeBind TopLevel env bind
                               binds' <- go env' binds
                               return (bind' `appendFloats` binds')
210

Peter Wortmann's avatar
Peter Wortmann committed
211
mkDataConWorkers :: DynFlags -> ModLocation -> [TyCon] -> [CoreBind]
212
-- See Note [Data constructor workers]
213
-- c.f. Note [Injecting implicit bindings] in TidyPgm
Peter Wortmann's avatar
Peter Wortmann committed
214 215 216
mkDataConWorkers dflags mod_loc data_tycons
  = [ NonRec id (tick_it (getName data_con) (Var id))
                                -- The ice is thin here, but it works
217
    | tycon <- data_tycons,     -- CorePrep will eta-expand it
218
      data_con <- tyConDataCons tycon,
Peter Wortmann's avatar
Peter Wortmann committed
219 220 221 222 223 224
      let id = dataConWorkId data_con
    ]
 where
   -- If we want to generate debug info, we put a source note on the
   -- worker. This is useful, especially for heap profiling.
   tick_it name
225
     | debugLevel dflags == 0                = id
Peter Wortmann's avatar
Peter Wortmann committed
226 227 228 229 230
     | RealSrcSpan span <- nameSrcSpan name  = tick span
     | Just file <- ml_hs_file mod_loc       = tick (span1 file)
     | otherwise                             = tick (span1 "???")
     where tick span  = Tick (SourceNote span $ showSDoc dflags (ppr name))
           span1 file = realSrcLocSpan $ mkRealSrcLoc (mkFastString file) 1 1
231

Austin Seipp's avatar
Austin Seipp committed
232
{-
233 234 235
Note [Floating out of top level bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
NB: we do need to float out of top-level bindings
236
Consider        x = length [True,False]
237
We want to get
238 239 240
                s1 = False : []
                s2 = True  : s1
                x  = length s2
241 242

We return a *list* of bindings, because we may start with
243
        x* = f (g y)
244
where x is demanded, in which case we want to finish with
245 246
        a = g y
        x* = f a
247 248 249 250
And then x will actually end up case-bound

Note [CafInfo and floating]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
251
What happens when we try to float bindings to the top level?  At this
252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269
point all the CafInfo is supposed to be correct, and we must make certain
that is true of the new top-level bindings.  There are two cases
to consider

a) The top-level binding is marked asCafRefs.  In that case we are
   basically fine.  The floated bindings had better all be lazy lets,
   so they can float to top level, but they'll all have HasCafRefs
   (the default) which is safe.

b) The top-level binding is marked NoCafRefs.  This really happens
   Example.  CoreTidy produces
      $fApplicativeSTM [NoCafRefs] = D:Alternative retry# ...blah...
   Now CorePrep has to eta-expand to
      $fApplicativeSTM = let sat = \xy. retry x y
                         in D:Alternative sat ...blah...
   So what we *want* is
      sat [NoCafRefs] = \xy. retry x y
      $fApplicativeSTM [NoCafRefs] = D:Alternative sat ...blah...
270

271
   So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
272
   *and* substutite the modified 'sat' into the old RHS.
273 274 275 276 277 278 279

   It should be the case that 'sat' is itself [NoCafRefs] (a value, no
   cafs) else the original top-level binding would not itself have been
   marked [NoCafRefs].  The DEBUG check in CoreToStg for
   consistentCafInfo will find this.

This is all very gruesome and horrible. It would be better to figure
280
out CafInfo later, after CorePrep.  We'll do that in due course.
281 282
Meanwhile this horrible hack works.

283 284 285

Note [Data constructor workers]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
286 287
Create any necessary "implicit" bindings for data con workers.  We
create the rather strange (non-recursive!) binding
288

289
        $wC = \x y -> $wC x y
290 291 292 293 294 295 296 297 298 299 300

i.e. a curried constructor that allocates.  This means that we can
treat the worker for a constructor like any other function in the rest
of the compiler.  The point here is that CoreToStg will generate a
StgConApp for the RHS, rather than a call to the worker (which would
give a loop).  As Lennart says: the ice is thin here, but it works.

Hmm.  Should we create bindings for dictionary constructors?  They are
always fully applied, and the bindings are just there to support
partial applications. But it's easier to let them through.

301

302 303
Note [Dead code in CorePrep]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Simon Peyton Jones's avatar
Simon Peyton Jones committed
304
Imagine that we got an input program like this (see Trac #4962):
305 306 307 308 309 310 311 312 313 314 315 316 317

  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
  f x = (g True (Just x) + g () (Just x), g)
    where
      g :: Show a => a -> Maybe Int -> Int
      g _ Nothing = x
      g y (Just z) = if z > 100 then g y (Just (z + length (show y))) else g y unknown

After specialisation and SpecConstr, we would get something like this:

  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
  f x = (g$Bool_True_Just x + g$Unit_Unit_Just x, g)
    where
318
      {-# RULES g $dBool = g$Bool
319 320 321 322 323 324 325 326 327
                g $dUnit = g$Unit #-}
      g = ...
      {-# RULES forall x. g$Bool True (Just x) = g$Bool_True_Just x #-}
      g$Bool = ...
      {-# RULES forall x. g$Unit () (Just x) = g$Unit_Unit_Just x #-}
      g$Unit = ...
      g$Bool_True_Just = ...
      g$Unit_Unit_Just = ...

328 329 330 331
Note that the g$Bool and g$Unit functions are actually dead code: they
are only kept alive by the occurrence analyser because they are
referred to by the rules of g, which is being kept alive by the fact
that it is used (unspecialised) in the returned pair.
332

333 334 335 336
However, at the CorePrep stage there is no way that the rules for g
will ever fire, and it really seems like a shame to produce an output
program that goes to the trouble of allocating a closure for the
unreachable g$Bool and g$Unit functions.
337 338 339

The way we fix this is to:
 * In cloneBndr, drop all unfoldings/rules
340 341 342 343 344 345

 * In deFloatTop, run a simple dead code analyser on each top-level
   RHS to drop the dead local bindings. For that call to OccAnal, we
   disable the binder swap, else the occurrence analyser sometimes
   introduces new let bindings for cased binders, which lead to the bug
   in #5433.
346 347 348 349 350 351

The reason we don't just OccAnal the whole output of CorePrep is that
the tidier ensures that all top-level binders are GlobalIds, so they
don't show up in the free variables any longer. So if you run the
occurrence analyser on the output of CoreTidy (or later) you e.g. turn
this program:
352 353 354 355 356 357 358 359 360 361 362 363

  Rec {
  f = ... f ...
  }

Into this one:

  f = ... f ...

(Since f is not considered to be free in its own RHS.)


Austin Seipp's avatar
Austin Seipp committed
364 365
************************************************************************
*                                                                      *
366
                The main code
Austin Seipp's avatar
Austin Seipp committed
367 368 369
*                                                                      *
************************************************************************
-}
370

371
cpeBind :: TopLevelFlag -> CorePrepEnv -> CoreBind
372
        -> UniqSM (CorePrepEnv, Floats)
373
cpeBind top_lvl env (NonRec bndr rhs)
374
  = do { (_, bndr1) <- cpCloneBndr env bndr
375
       ; let dmd         = idDemandInfo bndr
376
             is_unlifted = isUnliftedType (idType bndr)
377
       ; (floats, bndr2, rhs2) <- cpePair top_lvl NonRecursive
Austin Seipp's avatar
Austin Seipp committed
378
                                          dmd
379
                                          is_unlifted
380
                                          env bndr1 rhs
381
       -- See Note [Inlining in CorePrep]
382
       ; if exprIsTrivial rhs2 && isNotTopLevel top_lvl
383 384 385
            then return (extendCorePrepEnvExpr env bndr rhs2, floats)
            else do {

386
       ; let new_float = mkFloat dmd is_unlifted bndr2 rhs2
387

388 389
        -- We want bndr'' in the envt, because it records
        -- the evaluated-ness of the binder
390
       ; return (extendCorePrepEnv env bndr bndr2,
391
                 addFloat floats new_float) }}
392 393 394

cpeBind top_lvl env (Rec pairs)
  = do { let (bndrs,rhss) = unzip pairs
395
       ; (env', bndrs1) <- cpCloneBndrs env (map fst pairs)
396
       ; stuff <- zipWithM (cpePair top_lvl Recursive topDmd False env') bndrs1 rhss
397 398

       ; let (floats_s, bndrs2, rhss2) = unzip3 stuff
399
             all_pairs = foldrOL add_float (bndrs2 `zip` rhss2)
400
                                           (concatFloats floats_s)
401
       ; return (extendCorePrepEnvList env (bndrs `zip` bndrs2),
402
                 unitFloat (FloatLet (Rec all_pairs))) }
403
  where
404 405
        -- Flatten all the floats, and the currrent
        -- group into a single giant Rec
406 407 408 409 410
    add_float (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
    add_float (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
    add_float b                       _    = pprPanic "cpeBind" (ppr b)

---------------
411
cpePair :: TopLevelFlag -> RecFlag -> Demand -> Bool
412 413
        -> CorePrepEnv -> Id -> CoreExpr
        -> UniqSM (Floats, Id, CpeRhs)
414
-- Used for all bindings
415
cpePair top_lvl is_rec dmd is_unlifted env bndr rhs
416
  = do { (floats1, rhs1) <- cpeRhsE env rhs
417

418 419 420 421
       -- See if we are allowed to float this stuff out of the RHS
       ; (floats2, rhs2) <- float_from_rhs floats1 rhs1

       -- Make the arity match up
422
       ; (floats3, rhs3)
423 424 425 426 427
            <- if manifestArity rhs1 <= arity
               then return (floats2, cpeEtaExpand arity rhs2)
               else WARN(True, text "CorePrep: silly extra arguments:" <+> ppr bndr)
                               -- Note [Silly extra arguments]
                    (do { v <- newVar (idType bndr)
428
                        ; let float = mkFloat topDmd False v rhs2
429
                        ; return ( addFloat floats2 float
430
                                 , cpeEtaExpand arity (Var v)) })
431

432 433 434
        -- Wrap floating ticks
       ; let (floats4, rhs4) = wrapTicks floats3 rhs3

435 436 437 438 439
        -- Record if the binder is evaluated
        -- and otherwise trim off the unfolding altogether
        -- It's not used by the code generator; getting rid of it reduces
        -- heap usage and, since we may be changing uniques, we'd have
        -- to substitute to keep it right
440
       ; let bndr' | exprIsHNF rhs3 = bndr `setIdUnfolding` evaldUnfolding
441
                   | otherwise      = bndr `setIdUnfolding` noUnfolding
442

443
       ; return (floats4, bndr', rhs4) }
444
  where
445 446
    platform = targetPlatform (cpe_dynFlags env)

447
    arity = idArity bndr        -- We must match this arity
448 449

    ---------------------
450 451
    float_from_rhs floats rhs
      | isEmptyFloats floats = return (emptyFloats, rhs)
452 453
      | isTopLevel top_lvl   = float_top    floats rhs
      | otherwise            = float_nested floats rhs
454 455

    ---------------------
456
    float_nested floats rhs
457
      | wantFloatNested is_rec dmd is_unlifted floats rhs
458
                  = return (floats, rhs)
459
      | otherwise = dontFloat floats rhs
460 461

    ---------------------
462
    float_top floats rhs        -- Urhgh!  See Note [CafInfo and floating]
463
      | mayHaveCafRefs (idCafInfo bndr)
464 465 466 467
      , allLazyTop floats
      = return (floats, rhs)

      -- So the top-level binding is marked NoCafRefs
468
      | Just (floats', rhs') <- canFloatFromNoCaf platform floats rhs
469
      = return (floats', rhs')
470 471

      | otherwise
472 473 474 475 476 477 478 479 480 481 482
      = dontFloat floats rhs

dontFloat :: Floats -> CpeRhs -> UniqSM (Floats, CpeBody)
-- Non-empty floats, but do not want to float from rhs
-- So wrap the rhs in the floats
-- But: rhs1 might have lambdas, and we can't
--      put them inside a wrapBinds
dontFloat floats1 rhs
  = do { (floats2, body) <- rhsToBody rhs
        ; return (emptyFloats, wrapBinds floats1 $
                               wrapBinds floats2 body) }
483

484 485 486
{- Note [Silly extra arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we had this
487
        f{arity=1} = \x\y. e
488 489
We *must* match the arity on the Id, so we have to generate
        f' = \x\y. e
490
        f  = \x. f' x
491 492

It's a bizarre case: why is the arity on the Id wrong?  Reason
493
(in the days of __inline_me__):
494 495 496 497
        f{arity=0} = __inline_me__ (let v = expensive in \xy. e)
When InlineMe notes go away this won't happen any more.  But
it seems good for CorePrep to be robust.
-}
498

499
-- ---------------------------------------------------------------------------
500
--              CpeRhs: produces a result satisfying CpeRhs
501 502
-- ---------------------------------------------------------------------------

503
cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
504
-- If
505 506 507
--      e  ===>  (bs, e')
-- then
--      e = let bs in e'        (semantically, that is!)
508 509
--
-- For example
510
--      f (g x)   ===>   ([v = g x], f v)
511

512 513
cpeRhsE _env expr@(Type {})      = return (emptyFloats, expr)
cpeRhsE _env expr@(Coercion {})  = return (emptyFloats, expr)
514
cpeRhsE env (Lit (LitInteger i _))
515 516
    = cpeRhsE env (cvtLitInteger (cpe_dynFlags env) (getMkIntegerId env)
                   (cpe_integerSDataCon env) i)
517 518
cpeRhsE _env expr@(Lit {}) = return (emptyFloats, expr)
cpeRhsE env expr@(Var {})  = cpeApp env expr
519 520
cpeRhsE env expr@(App {}) = cpeApp env expr

521 522 523 524 525
cpeRhsE env (Let bind expr)
  = do { (env', new_binds) <- cpeBind NotTopLevel env bind
       ; (floats, body) <- cpeRhsE env' expr
       ; return (new_binds `appendFloats` floats, body) }

526
cpeRhsE env (Tick tickish expr)
527 528 529 530 531
  | tickishPlace tickish == PlaceNonLam && tickish `tickishScopesLike` SoftScope
  = do { (floats, body) <- cpeRhsE env expr
         -- See [Floating Ticks in CorePrep]
       ; return (unitFloat (FloatTick tickish) `appendFloats` floats, body) }
  | otherwise
532
  = do { body <- cpeBodyNF env expr
533
       ; return (emptyFloats, mkTick tickish' body) }
534 535
  where
    tickish' | Breakpoint n fvs <- tickish
536 537
             -- See also 'substTickish'
             = Breakpoint n (map (getIdFromTrivialExpr . lookupCorePrepEnv env) fvs)
538 539
             | otherwise
             = tickish
540 541 542 543 544 545 546

cpeRhsE env (Cast expr co)
   = do { (floats, expr') <- cpeRhsE env expr
        ; return (floats, Cast expr' co) }

cpeRhsE env expr@(Lam {})
   = do { let (bndrs,body) = collectBinders expr
547
        ; (env', bndrs') <- cpCloneBndrs env bndrs
548 549
        ; body' <- cpeBodyNF env' body
        ; return (emptyFloats, mkLams bndrs' body') }
550 551 552 553 554

cpeRhsE env (Case scrut bndr ty alts)
  = do { (floats, scrut') <- cpeBody env scrut
       ; let bndr1 = bndr `setIdUnfolding` evaldUnfolding
            -- Record that the case binder is evaluated in the alternatives
555
       ; (env', bndr2) <- cpCloneBndr env bndr1
556 557 558 559
       ; alts' <- mapM (sat_alt env') alts
       ; return (floats, Case scrut' bndr2 ty alts') }
  where
    sat_alt env (con, bs, rhs)
560
       = do { (env2, bs') <- cpCloneBndrs env bs
561 562
            ; rhs' <- cpeBodyNF env2 rhs
            ; return (con, bs', rhs') }
563

564
cvtLitInteger :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
565
-- Here we convert a literal Integer to the low-level
566
-- representation. Exactly how we do this depends on the
567 568
-- library that implements Integer.  If it's GMP we
-- use the S# data constructor for small literals.
569
-- See Note [Integer literals] in Literal
570 571 572
cvtLitInteger dflags _ (Just sdatacon) i
  | inIntRange dflags i -- Special case for small integers
    = mkConApp sdatacon [Lit (mkMachInt dflags i)]
573

574
cvtLitInteger dflags mk_integer _ i
575
    = mkApps (Var mk_integer) [isNonNegative, ints]
576 577 578 579 580 581
  where isNonNegative = if i < 0 then mkConApp falseDataCon []
                                 else mkConApp trueDataCon  []
        ints = mkListExpr intTy (f (abs i))
        f 0 = []
        f x = let low  = x .&. mask
                  high = x `shiftR` bits
582
              in mkConApp intDataCon [Lit (mkMachInt dflags low)] : f high
583 584
        bits = 31
        mask = 2 ^ bits - 1
585

586
-- ---------------------------------------------------------------------------
587
--              CpeBody: produces a result satisfying CpeBody
588
-- ---------------------------------------------------------------------------
589

590 591 592 593 594
-- | Convert a 'CoreExpr' so it satisfies 'CpeBody', without
-- producing any floats (any generated floats are immediately
-- let-bound using 'wrapBinds').  Generally you want this, esp.
-- when you've reached a binding form (e.g., a lambda) and
-- floating any further would be incorrect.
595
cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
596
cpeBodyNF env expr
597 598
  = do { (floats, body) <- cpeBody env expr
       ; return (wrapBinds floats body) }
599

600 601 602 603 604 605 606 607 608 609
-- | Convert a 'CoreExpr' so it satisfies 'CpeBody'; also produce
-- a list of 'Floats' which are being propagated upwards.  In
-- fact, this function is used in only two cases: to
-- implement 'cpeBodyNF' (which is what you usually want),
-- and in the case when a let-binding is in a case scrutinee--here,
-- we can always float out:
--
--      case (let x = y in z) of ...
--      ==> let x = y in case z of ...
--
610 611 612 613 614
cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
cpeBody env expr
  = do { (floats1, rhs) <- cpeRhsE env expr
       ; (floats2, body) <- rhsToBody rhs
       ; return (floats1 `appendFloats` floats2, body) }
615

616 617
--------
rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
618
-- Remove top level lambdas by let-binding
619

620
rhsToBody (Tick t expr)
Peter Wortmann's avatar
Peter Wortmann committed
621
  | tickishScoped t == NoScope  -- only float out of non-scoped annotations
622
  = do { (floats, expr') <- rhsToBody expr
623
       ; return (floats, mkTick t expr') }
624

625
rhsToBody (Cast e co)
626 627
        -- You can get things like
        --      case e of { p -> coerce t (\s -> ...) }
628 629
  = do { (floats, e') <- rhsToBody e
       ; return (floats, Cast e' co) }
630

631
rhsToBody expr@(Lam {})
632
  | Just no_lam_result <- tryEtaReducePrep bndrs body
633
  = return (emptyFloats, no_lam_result)
634
  | all isTyVar bndrs           -- Type lambdas are ok
635
  = return (emptyFloats, expr)
636
  | otherwise                   -- Some value lambdas
637 638
  = do { fn <- newVar (exprType expr)
       ; let rhs   = cpeEtaExpand (exprArity expr) expr
639
             float = FloatLet (NonRec fn rhs)
640
       ; return (unitFloat float, Var fn) }
641 642
  where
    (bndrs,body) = collectBinders expr
643

644 645
rhsToBody expr = return (emptyFloats, expr)

646

647 648

-- ---------------------------------------------------------------------------
649
--              CpeApp: produces a result satisfying CpeApp
650 651
-- ---------------------------------------------------------------------------

652 653 654
data ArgInfo = CpeApp  CoreArg
             | CpeCast Coercion
             | CpeTick (Tickish Id)
655 656 657 658 659 660 661 662 663 664 665 666

{- Note [runRW arg]
~~~~~~~~~~~~~~~~~~~
If we got, say
   runRW# (case bot of {})
which happened in Trac #11291, we do /not/ want to turn it into
   (case bot of {}) realWorldPrimId#
because that gives a panic in CoreToStg.myCollectArgs, which expects
only variables in function position.  But if we are sure to make
runRW# strict (which we do in MkId), this can't happen
-}

667 668
cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
-- May return a CpeRhs because of saturating primops
669 670
cpeApp top_env expr
  = do { let (terminal, args, depth) = collect_args expr
671 672
       ; cpe_app top_env terminal args depth
       }
673 674

  where
675 676
    -- We have a nested data structure of the form
    -- e `App` a1 `App` a2 ... `App` an, convert it into
677 678
    -- (e, [CpeApp a1, CpeApp a2, ..., CpeApp an], depth)
    -- We use 'ArgInfo' because we may also need to
679 680 681
    -- record casts and ticks.  Depth counts the number
    -- of arguments that would consume strictness information
    -- (so, no type or coercion arguments.)
682
    collect_args :: CoreExpr -> (CoreExpr, [ArgInfo], Int)
683 684 685
    collect_args e = go e [] 0
      where
        go (App fun arg)      as depth
686
            = go fun (CpeApp arg : as)
687 688 689 690 691 692 693 694 695 696 697
                (if isTyCoArg arg then depth else depth + 1)
        go (Cast fun co)      as depth
            = go fun (CpeCast co : as) depth
        go (Tick tickish fun) as depth
            | tickishPlace tickish == PlaceNonLam
            && tickish `tickishScopesLike` SoftScope
            = go fun (CpeTick tickish : as) depth
        go terminal as depth = (terminal, as, depth)

    cpe_app :: CorePrepEnv
            -> CoreExpr
698
            -> [ArgInfo]
699
            -> Int
700
            -> UniqSM (Floats, CpeRhs)
701
    cpe_app env (Var f) (CpeApp Type{} : CpeApp arg : args) depth
702 703
        | f `hasKey` lazyIdKey          -- Replace (lazy a) with a, and
       || f `hasKey` noinlineIdKey      -- Replace (noinline a) with a
704 705 706 707 708 709 710 711 712 713 714 715 716 717 718
        -- Consider the code:
        --
        --      lazy (f x) y
        --
        -- We need to make sure that we need to recursively collect arguments on
        -- "f x", otherwise we'll float "f x" out (it's not a variable) and
        -- end up with this awful -ddump-prep:
        --
        --      case f x of f_x {
        --        __DEFAULT -> f_x y
        --      }
        --
        -- rather than the far superior "f x y".  Test case is par01.
        = let (terminal, args', depth') = collect_args arg
          in cpe_app env terminal (args' ++ args) (depth + depth' - 1)
719
    cpe_app env (Var f) [CpeApp _runtimeRep@Type{}, CpeApp _type@Type{}, CpeApp arg] 1
720 721 722 723 724
        | f `hasKey` runRWKey
        -- Replace (runRW# f) by (f realWorld#), beta reducing if possible (this
        -- is why we return a CorePrepEnv as well)
        = case arg of
            Lam s body -> cpe_app (extendCorePrepEnv env s realWorldPrimId) body [] 0
725
            _          -> cpe_app env arg [CpeApp (Var realWorldPrimId)] 1
726
    cpe_app env (Var v) args depth
727
      = do { v1 <- fiddleCCall v
728
           ; let e2 = lookupCorePrepEnv env v1
729 730
                 hd = getIdFromTrivialExpr_maybe e2
           -- NB: depth from collect_args is right, because e2 is a trivial expression
731 732
           -- and thus its embedded Id *must* be at the same depth as any
           -- Apps it is under are type applications only (c.f.
733
           -- exprIsTrivial).  But note that we need the type of the
734
           -- expression, not the id.
735
           ; (app, floats) <- rebuild_app args e2 (exprType e2) emptyFloats stricts
736
           ; mb_saturate hd app floats depth }
737 738
        where
          stricts = case idStrictness v of
739 740
                            StrictSig (DmdType _ demands _)
                              | listLengthCmp demands depth /= GT -> demands
741
                                    -- length demands <= depth
742
                              | otherwise                         -> []
743 744 745 746 747
                -- If depth < length demands, then we have too few args to
                -- satisfy strictness  info so we have to  ignore all the
                -- strictness info, e.g. + (error "urk")
                -- Here, we can't evaluate the arg strictly, because this
                -- partial application might be seq'd
748

749 750 751 752
        -- We inlined into something that's not a var and has no args.
        -- Bounce it back up to cpeRhsE.
    cpe_app env fun [] _ = cpeRhsE env fun

753
        -- N-variable fun, better let-bind it
754
    cpe_app env fun args depth
755 756
      = do { (fun_floats, fun') <- cpeArg env evalDmd fun ty
                          -- The evalDmd says that it's sure to be evaluated,
757
                          -- so we'll end up case-binding it
758
           ; (app, floats) <- rebuild_app args fun' ty fun_floats []
759
           ; mb_saturate Nothing app floats depth }
760
        where
761
          ty = exprType fun
762

763 764 765 766 767 768 769
    -- Saturate if necessary
    mb_saturate head app floats depth =
       case head of
         Just fn_id -> do { sat_app <- maybeSaturate fn_id app depth
                          ; return (floats, sat_app) }
         _other              -> return (floats, app)

770 771 772 773 774 775
    -- Deconstruct and rebuild the application, floating any non-atomic
    -- arguments to the outside.  We collect the type of the expression,
    -- the head of the application, and the number of actual value arguments,
    -- all of which are used to possibly saturate this application if it
    -- has a constructor or primop at the head.
    rebuild_app
776
        :: [ArgInfo]                  -- The arguments (inner to outer)
777 778 779 780 781 782 783 784 785
        -> CpeApp
        -> Type
        -> Floats
        -> [Demand]
        -> UniqSM (CpeApp, Floats)
    rebuild_app [] app _ floats ss = do
      MASSERT(null ss) -- make sure we used all the strictness info
      return (app, floats)
    rebuild_app (a : as) fun' fun_ty floats ss = case a of
786
      CpeApp arg@(Type arg_ty) ->
787
        rebuild_app as (App fun' arg) (piResultTy fun_ty arg_ty) floats ss
788
      CpeApp arg@(Coercion {}) ->
789
        rebuild_app as (App fun' arg) (funResultTy fun_ty) floats ss
790
      CpeApp arg -> do
791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806
        let (ss1, ss_rest)  -- See Note [lazyId magic] in MkId
               = case (ss, isLazyExpr arg) of
                   (_   : ss_rest, True)  -> (topDmd, ss_rest)
                   (ss1 : ss_rest, False) -> (ss1,    ss_rest)
                   ([],            _)     -> (topDmd, [])
            (arg_ty, res_ty) = expectJust "cpeBody:collect_args" $
                               splitFunTy_maybe fun_ty
        (fs, arg') <- cpeArg top_env ss1 arg arg_ty
        rebuild_app as (App fun' arg') res_ty (fs `appendFloats` floats) ss_rest
      CpeCast co ->
        let Pair _ty1 ty2 = coercionKind co
        in rebuild_app as (Cast fun' co) ty2 floats ss
      CpeTick tickish ->
        -- See [Floating Ticks in CorePrep]
        rebuild_app as fun' fun_ty (addFloat floats (FloatTick tickish)) ss

807 808 809 810 811 812 813
isLazyExpr :: CoreExpr -> Bool
-- See Note [lazyId magic] in MkId
isLazyExpr (Cast e _)              = isLazyExpr e
isLazyExpr (Tick _ e)              = isLazyExpr e
isLazyExpr (Var f `App` _ `App` _) = f `hasKey` lazyIdKey
isLazyExpr _                       = False

814
-- ---------------------------------------------------------------------------
815
--      CpeArg: produces a result satisfying CpeArg
816 817
-- ---------------------------------------------------------------------------

818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851
{-
Note [ANF-ising literal string arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Consider a program like,

    data Foo = Foo Addr#

    foo = Foo "turtle"#

When we go to ANFise this we might think that we want to float the string
literal like we do any other non-trivial argument. This would look like,

    foo = u\ [] case "turtle"# of s { __DEFAULT__ -> Foo s }

However, this 1) isn't necessary since strings are in a sense "trivial"; and 2)
wreaks havoc on the CAF annotations that we produce here since we the result
above is caffy since it is updateable. Ideally at some point in the future we
would like to just float the literal to the top level as suggested in #11312,

    s = "turtle"#
    foo = Foo s

However, until then we simply add a special case excluding literals from the
floating done by cpeArg.
-}

-- | Is an argument okay to CPE?
okCpeArg :: CoreExpr -> Bool
-- Don't float literals. See Note [ANF-ising literal string arguments].
okCpeArg (Lit _) = False
-- Do not eta expand a trivial argument
okCpeArg expr    = not (exprIsTrivial expr)

852
-- This is where we arrange that a non-trivial argument is let-bound
Austin Seipp's avatar
Austin Seipp committed
853
cpeArg :: CorePrepEnv -> Demand
854
       -> CoreArg -> Type -> UniqSM (Floats, CpeArg)
855
cpeArg env dmd arg arg_ty
856
  = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
857 858
       ; (floats2, arg2) <- if want_float floats1 arg1
                            then return (floats1, arg1)
859
                            else dontFloat floats1 arg1
860 861
                -- Else case: arg1 might have lambdas, and we can't
                --            put them inside a wrapBinds
862

863 864 865 866 867 868 869
       ; if okCpeArg arg2
         then do { v <- newVar arg_ty
                 ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
                       arg_float = mkFloat dmd is_unlifted v arg3
                 ; return (addFloat floats2 arg_float, varToCoreExpr v) }
         else return (floats2, arg2)
       }
870
  where
871
    is_unlifted = isUnliftedType arg_ty
872
    want_float  = wantFloatNested NonRecursive dmd is_unlifted
873

Austin Seipp's avatar
Austin Seipp committed
874
{-
875 876 877 878 879 880 881 882 883 884 885 886 887
Note [Floating unlifted arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider    C (let v* = expensive in v)

where the "*" indicates "will be demanded".  Usually v will have been
inlined by now, but let's suppose it hasn't (see Trac #2756).  Then we
do *not* want to get

     let v* = expensive in C v

because that has different strictness.  Hence the use of 'allLazy'.
(NB: the let v* turns into a FloatCase, in mkLocalNonRec.)

888 889 890 891