CorePrep.hs 49.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 #-}
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

114 115
10. Convert (LitInteger i t) into the core representation
    for the Integer i. Normally this uses mkInteger, but if
116 117 118 119
    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.

120 121 122 123
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.

124 125 126
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.
127

128

129 130 131
Invariants
~~~~~~~~~~
Here is the syntax of the Core produced by CorePrep:
132

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

    Applications
139
       app ::= lit  |  var  |  app triv  |  app ty  | app co | app |> co
140 141

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

148
    Right hand sides (only place where value lambdas can occur)
149 150 151 152
       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
153
-}
154

155 156 157 158
type CpeTriv = CoreExpr    -- Non-terminal 'triv'
type CpeApp  = CoreExpr    -- Non-terminal 'app'
type CpeBody = CoreExpr    -- Non-terminal 'body'
type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'
159

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

Peter Wortmann's avatar
Peter Wortmann committed
168 169 170
corePrepPgm :: HscEnv -> ModLocation -> CoreProgram -> [TyCon] -> IO CoreProgram
corePrepPgm hsc_env mod_loc binds data_tycons = do
    let dflags = hsc_dflags hsc_env
171 172
    showPass dflags "CorePrep"
    us <- mkSplitUniqSupply 's'
173
    initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
174

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

        binds_out = initUs_ us $ do
180 181
                      floats1 <- corePrepTopBinds initialCorePrepEnv binds
                      floats2 <- corePrepTopBinds initialCorePrepEnv implicit_binds
182 183
                      return (deFloatTop (floats1 `appendFloats` floats2))

184
    endPassIO hsc_env alwaysQualify CorePrep binds_out []
185
    return binds_out
186

187 188
corePrepExpr :: DynFlags -> HscEnv -> CoreExpr -> IO CoreExpr
corePrepExpr dflags hsc_env expr = do
189 190
    showPass dflags "CorePrep"
    us <- mkSplitUniqSupply 's'
191
    initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
192
    let new_expr = initUs_ us (cpeBodyNF initialCorePrepEnv expr)
193 194
    dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep" (ppr new_expr)
    return new_expr
195

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

Peter Wortmann's avatar
Peter Wortmann committed
206
mkDataConWorkers :: DynFlags -> ModLocation -> [TyCon] -> [CoreBind]
207
-- See Note [Data constructor workers]
208
-- c.f. Note [Injecting implicit bindings] in TidyPgm
Peter Wortmann's avatar
Peter Wortmann committed
209 210 211
mkDataConWorkers dflags mod_loc data_tycons
  = [ NonRec id (tick_it (getName data_con) (Var id))
                                -- The ice is thin here, but it works
212
    | tycon <- data_tycons,     -- CorePrep will eta-expand it
213
      data_con <- tyConDataCons tycon,
Peter Wortmann's avatar
Peter Wortmann committed
214 215 216 217 218 219
      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
220
     | debugLevel dflags == 0                = id
Peter Wortmann's avatar
Peter Wortmann committed
221 222 223 224 225
     | 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
226

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

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

Note [CafInfo and floating]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
246
What happens when we try to float bindings to the top level?  At this
247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
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...
265

266
   So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
267
   *and* substutite the modified 'sat' into the old RHS.
268 269 270 271 272 273 274

   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
275
out CafInfo later, after CorePrep.  We'll do that in due course.
276 277
Meanwhile this horrible hack works.

278 279 280

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

284
        $wC = \x y -> $wC x y
285 286 287 288 289 290 291 292 293 294 295

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.

296

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

  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
313
      {-# RULES g $dBool = g$Bool
314 315 316 317 318 319 320 321 322
                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 = ...

323 324 325 326
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.
327

328 329 330 331
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.
332 333 334

The way we fix this is to:
 * In cloneBndr, drop all unfoldings/rules
335 336 337 338 339 340

 * 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.
341 342 343 344 345 346

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:
347 348 349 350 351 352 353 354 355 356 357 358

  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
359 360
************************************************************************
*                                                                      *
361
                The main code
Austin Seipp's avatar
Austin Seipp committed
362 363 364
*                                                                      *
************************************************************************
-}
365

366
cpeBind :: TopLevelFlag -> CorePrepEnv -> CoreBind
367
        -> UniqSM (CorePrepEnv, Floats)
368
cpeBind top_lvl env (NonRec bndr rhs)
369
  = do { (_, bndr1) <- cpCloneBndr env bndr
370
       ; let dmd         = idDemandInfo bndr
371
             is_unlifted = isUnLiftedType (idType bndr)
372
       ; (floats, bndr2, rhs2) <- cpePair top_lvl NonRecursive
Austin Seipp's avatar
Austin Seipp committed
373
                                          dmd
374
                                          is_unlifted
375
                                          env bndr1 rhs
376
       ; let new_float = mkFloat dmd is_unlifted bndr2 rhs2
377

378 379
        -- We want bndr'' in the envt, because it records
        -- the evaluated-ness of the binder
380 381
       ; return (extendCorePrepEnv env bndr bndr2,
                 addFloat floats new_float) }
382 383 384

cpeBind top_lvl env (Rec pairs)
  = do { let (bndrs,rhss) = unzip pairs
385
       ; (env', bndrs1) <- cpCloneBndrs env (map fst pairs)
386
       ; stuff <- zipWithM (cpePair top_lvl Recursive topDmd False env') bndrs1 rhss
387 388

       ; let (floats_s, bndrs2, rhss2) = unzip3 stuff
389
             all_pairs = foldrOL add_float (bndrs2 `zip` rhss2)
390
                                           (concatFloats floats_s)
391
       ; return (extendCorePrepEnvList env (bndrs `zip` bndrs2),
392
                 unitFloat (FloatLet (Rec all_pairs))) }
393
  where
394 395
        -- Flatten all the floats, and the currrent
        -- group into a single giant Rec
396 397 398 399 400
    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)

---------------
401
cpePair :: TopLevelFlag -> RecFlag -> Demand -> Bool
402 403
        -> CorePrepEnv -> Id -> CoreExpr
        -> UniqSM (Floats, Id, CpeRhs)
404
-- Used for all bindings
405
cpePair top_lvl is_rec dmd is_unlifted env bndr rhs
406
  = do { (floats1, rhs1) <- cpeRhsE env rhs
407

408 409 410 411
       -- 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
412
       ; (floats3, rhs3)
413 414 415 416 417
            <- 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)
418
                        ; let float = mkFloat topDmd False v rhs2
419
                        ; return ( addFloat floats2 float
420
                                 , cpeEtaExpand arity (Var v)) })
421

422 423 424
        -- Wrap floating ticks
       ; let (floats4, rhs4) = wrapTicks floats3 rhs3

425 426 427 428 429
        -- 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
430
       ; let bndr' | exprIsHNF rhs3 = bndr `setIdUnfolding` evaldUnfolding
431
                   | otherwise      = bndr `setIdUnfolding` noUnfolding
432

433
       ; return (floats4, bndr', rhs4) }
434
  where
435 436
    is_strict_or_unlifted = (isStrictDmd dmd) || is_unlifted

437 438
    platform = targetPlatform (cpe_dynFlags env)

439
    arity = idArity bndr        -- We must match this arity
440 441

    ---------------------
442 443 444 445
    float_from_rhs floats rhs
      | isEmptyFloats floats = return (emptyFloats, rhs)
      | isTopLevel top_lvl    = float_top    floats rhs
      | otherwise             = float_nested floats rhs
446 447

    ---------------------
448 449 450 451
    float_nested floats rhs
      | wantFloatNested is_rec is_strict_or_unlifted floats rhs
                  = return (floats, rhs)
      | otherwise = dont_float floats rhs
452 453

    ---------------------
454
    float_top floats rhs        -- Urhgh!  See Note [CafInfo and floating]
455
      | mayHaveCafRefs (idCafInfo bndr)
456 457 458 459
      , allLazyTop floats
      = return (floats, rhs)

      -- So the top-level binding is marked NoCafRefs
460
      | Just (floats', rhs') <- canFloatFromNoCaf platform floats rhs
461
      = return (floats', rhs')
462 463

      | otherwise
464
      = dont_float floats rhs
465 466

    ---------------------
467
    dont_float floats rhs
468 469 470 471
      -- 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
472
      = do { body <- rhsToBodyNF rhs
473
           ; return (emptyFloats, wrapBinds floats body) }
474

475 476 477
{- Note [Silly extra arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we had this
478
        f{arity=1} = \x\y. e
479 480
We *must* match the arity on the Id, so we have to generate
        f' = \x\y. e
481
        f  = \x. f' x
482 483

It's a bizarre case: why is the arity on the Id wrong?  Reason
484
(in the days of __inline_me__):
485 486 487 488
        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.
-}
489

490
-- ---------------------------------------------------------------------------
491
--              CpeRhs: produces a result satisfying CpeRhs
492 493
-- ---------------------------------------------------------------------------

494
cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
495
-- If
496 497 498
--      e  ===>  (bs, e')
-- then
--      e = let bs in e'        (semantically, that is!)
499 500
--
-- For example
501
--      f (g x)   ===>   ([v = g x], f v)
502

503 504
cpeRhsE _env expr@(Type {})      = return (emptyFloats, expr)
cpeRhsE _env expr@(Coercion {})  = return (emptyFloats, expr)
505
cpeRhsE env (Lit (LitInteger i _))
506 507
    = cpeRhsE env (cvtLitInteger (cpe_dynFlags env) (getMkIntegerId env)
                   (cpe_integerSDataCon env) i)
508 509
cpeRhsE _env expr@(Lit {}) = return (emptyFloats, expr)
cpeRhsE env expr@(Var {})  = cpeApp env expr
510

511
cpeRhsE env (Var f `App` _{-type-} `App` arg)
512 513
  | f `hasKey` lazyIdKey          -- Replace (lazy a) by a
  = cpeRhsE env arg               -- See Note [lazyId magic] in MkId
514

515
cpeRhsE env (Var f `App` _levity `App` _type `App` arg)
516 517 518
    -- See Note [runRW magic] in MkId
  | f `hasKey` runRWKey           -- Replace (runRW# f) by (f realWorld#),
  = case arg of                   -- beta reducing if possible
519
      Lam s body -> cpeRhsE (extendCorePrepEnv env s realWorldPrimId) body
520
      _          -> cpeRhsE env (arg `App` Var realWorldPrimId)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
521 522 523 524 525 526 527 528 529 530 531 532
                    -- See Note [runRW arg]

{- 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
-}
533

534 535
cpeRhsE env expr@(App {}) = cpeApp env expr

536 537 538 539 540
cpeRhsE env (Let bind expr)
  = do { (env', new_binds) <- cpeBind NotTopLevel env bind
       ; (floats, body) <- cpeRhsE env' expr
       ; return (new_binds `appendFloats` floats, body) }

541
cpeRhsE env (Tick tickish expr)
542 543 544 545 546
  | 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
547
  = do { body <- cpeBodyNF env expr
548
       ; return (emptyFloats, mkTick tickish' body) }
549 550 551 552 553
  where
    tickish' | Breakpoint n fvs <- tickish
             = Breakpoint n (map (lookupCorePrepEnv env) fvs)
             | otherwise
             = tickish
554 555 556 557 558 559 560

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
561
        ; (env', bndrs') <- cpCloneBndrs env bndrs
562 563
        ; body' <- cpeBodyNF env' body
        ; return (emptyFloats, mkLams bndrs' body') }
564 565 566 567 568

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
569
       ; (env', bndr2) <- cpCloneBndr env bndr1
570 571 572 573
       ; alts' <- mapM (sat_alt env') alts
       ; return (floats, Case scrut' bndr2 ty alts') }
  where
    sat_alt env (con, bs, rhs)
574
       = do { (env2, bs') <- cpCloneBndrs env bs
575 576
            ; rhs' <- cpeBodyNF env2 rhs
            ; return (con, bs', rhs') }
577

578
cvtLitInteger :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
579 580
-- Here we convert a literal Integer to the low-level
-- represenation. Exactly how we do this depends on the
581 582
-- library that implements Integer.  If it's GMP we
-- use the S# data constructor for small literals.
583
-- See Note [Integer literals] in Literal
584 585 586
cvtLitInteger dflags _ (Just sdatacon) i
  | inIntRange dflags i -- Special case for small integers
    = mkConApp sdatacon [Lit (mkMachInt dflags i)]
587

588
cvtLitInteger dflags mk_integer _ i
589
    = mkApps (Var mk_integer) [isNonNegative, ints]
590 591 592 593 594 595
  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
596
              in mkConApp intDataCon [Lit (mkMachInt dflags low)] : f high
597 598
        bits = 31
        mask = 2 ^ bits - 1
599

600
-- ---------------------------------------------------------------------------
601
--              CpeBody: produces a result satisfying CpeBody
602
-- ---------------------------------------------------------------------------
603

604
cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
605
cpeBodyNF env expr
606 607
  = do { (floats, body) <- cpeBody env expr
       ; return (wrapBinds floats body) }
608

609 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 618
--------
rhsToBodyNF :: CpeRhs -> UniqSM CpeBody
rhsToBodyNF rhs = do { (floats,body) <- rhsToBody rhs
619
                     ; return (wrapBinds floats body) }
620

621 622
--------
rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
623
-- Remove top level lambdas by let-binding
624

625
rhsToBody (Tick t expr)
Peter Wortmann's avatar
Peter Wortmann committed
626
  | tickishScoped t == NoScope  -- only float out of non-scoped annotations
627
  = do { (floats, expr') <- rhsToBody expr
628
       ; return (floats, mkTick t expr') }
629

630
rhsToBody (Cast e co)
631 632
        -- You can get things like
        --      case e of { p -> coerce t (\s -> ...) }
633 634
  = do { (floats, e') <- rhsToBody e
       ; return (floats, Cast e' co) }
635

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

649 650
rhsToBody expr = return (emptyFloats, expr)

651

652 653

-- ---------------------------------------------------------------------------
654
--              CpeApp: produces a result satisfying CpeApp
655 656 657 658
-- ---------------------------------------------------------------------------

cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
-- May return a CpeRhs because of saturating primops
659
cpeApp env expr
660
  = do { (app, (head,depth), _, floats, ss) <- collect_args expr 0
661
       ; MASSERT(null ss)       -- make sure we used all the strictness info
662

663
        -- Now deal with the function
664 665
       ; case head of
           Var fn_id -> do { sat_app <- maybeSaturate fn_id app depth
666
                           ; return (floats, sat_app) }
667
           _other    -> return (floats, app) }
668 669

  where
670 671
    -- Deconstruct and rebuild the application, floating any non-atomic
    -- arguments to the outside.  We collect the type of the expression,
672
    -- the head of the application, and the number of actual value arguments,
673 674 675
    -- all of which are used to possibly saturate this application if it
    -- has a constructor or primop at the head.

676
    collect_args
677 678 679 680 681 682 683 684
        :: CoreExpr
        -> Int                     -- Current app depth
        -> UniqSM (CpeApp,         -- The rebuilt expression
                   (CoreExpr,Int), -- The head of the application,
                                   -- and no. of args it was applied to
                   Type,           -- Type of the whole expr
                   Floats,         -- Any floats we pulled out
                   [Demand])       -- Remaining argument demands
685 686 687

    collect_args (App fun arg@(Type arg_ty)) depth
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
688
           ; return (App fun' arg, hd, piResultTy fun_ty arg_ty, floats, ss) }
689

690
    collect_args (App fun arg@(Coercion {})) depth
691
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
692
           ; return (App fun' arg, hd, funResultTy fun_ty, floats, ss) }
693

694 695
    collect_args (App fun arg) depth
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun (depth+1)
696
           ; let
697
              (ss1, ss_rest)   = case ss of
698 699
                                   (ss1:ss_rest)             -> (ss1,     ss_rest)
                                   []                        -> (topDmd, [])
700
              (arg_ty, res_ty) = expectJust "cpeBody:collect_args" $
701
                                 splitFunTy_maybe fun_ty
702

703
           ; (fs, arg') <- cpeArg env ss1 arg arg_ty
704
           ; return (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest) }
705

706
    collect_args (Var v) depth
707 708 709
      = do { v1 <- fiddleCCall v
           ; let v2 = lookupCorePrepEnv env v1
           ; return (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts) }
710 711
        where
          stricts = case idStrictness v of
712 713
                            StrictSig (DmdType _ demands _)
                              | listLengthCmp demands depth /= GT -> demands
714
                                    -- length demands <= depth
715
                              | otherwise                         -> []
716 717 718 719 720
                -- 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
721

722
    collect_args (Cast fun co) depth
723
      = do { let Pair _ty1 ty2 = coercionKind co
724 725
           ; (fun', hd, _, floats, ss) <- collect_args fun depth
           ; return (Cast fun' co, hd, ty2, floats, ss) }
726

727
    collect_args (Tick tickish fun) depth
728 729 730 731 732
      | tickishPlace tickish == PlaceNonLam
        && tickish `tickishScopesLike` SoftScope
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
             -- See [Floating Ticks in CorePrep]
           ; return (fun',hd,fun_ty,addFloat floats (FloatTick tickish),ss) }
733

734
        -- N-variable fun, better let-bind it
735
    collect_args fun depth
736 737
      = do { (fun_floats, fun') <- cpeArg env evalDmd fun ty
                          -- The evalDmd says that it's sure to be evaluated,
738
                          -- so we'll end up case-binding it
739
           ; return (fun', (fun', depth), ty, fun_floats, []) }
740
        where
741
          ty = exprType fun
742

743
-- ---------------------------------------------------------------------------
744
--      CpeArg: produces a result satisfying CpeArg
745 746 747
-- ---------------------------------------------------------------------------

-- This is where we arrange that a non-trivial argument is let-bound
Austin Seipp's avatar
Austin Seipp committed
748
cpeArg :: CorePrepEnv -> Demand
749
       -> CoreArg -> Type -> UniqSM (Floats, CpeTriv)
750
cpeArg env dmd arg arg_ty
751
  = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
752 753 754 755 756 757
       ; (floats2, arg2) <- if want_float floats1 arg1
                            then return (floats1, arg1)
                            else do { body1 <- rhsToBodyNF arg1
                                    ; return (emptyFloats, wrapBinds floats1 body1) }
                -- Else case: arg1 might have lambdas, and we can't
                --            put them inside a wrapBinds
758

759 760 761 762
       ; if cpe_ExprIsTrivial arg2    -- Do not eta expand a trivial argument
         then return (floats2, arg2)
         else do
       { v <- newVar arg_ty
763
       ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
764
             arg_float = mkFloat dmd is_unlifted v arg3
765
       ; return (addFloat floats2 arg_float, varToCoreExpr v) } }
766 767
  where
    is_unlifted = isUnLiftedType arg_ty
768 769
    is_strict   = isStrictDmd dmd
    want_float  = wantFloatNested NonRecursive (is_strict || is_unlifted)
770

Austin Seipp's avatar
Austin Seipp committed
771
{-
772 773 774 775 776 777 778 779 780 781 782 783 784
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.)

785 786 787 788 789

------------------------------------------------------------------------------
-- Building the saturated syntax
-- ---------------------------------------------------------------------------

790 791
maybeSaturate deals with saturating primops and constructors
The type is the type of the entire application
Austin Seipp's avatar
Austin Seipp committed
792
-}
793 794 795

maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
maybeSaturate fn expr n_args
796 797
  | Just DataToTagOp <- isPrimOpId_maybe fn     -- DataToTag must have an evaluated arg
                                                -- A gruesome special case
798
  = saturateDataToTag sat_expr
799

800
  | hasNoBinding fn        -- There's no binding