CorePrep.hs 50.3 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

168 169 170 171 172 173
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
174
    us <- mkSplitUniqSupply 's'
175
    initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
176

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

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

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

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

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

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

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

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

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

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

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

282 283 284

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

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

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.

300

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

  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
317
      {-# RULES g $dBool = g$Bool
318 319 320 321 322 323 324 325 326
                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 = ...

327 328 329 330
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.
331

332 333 334 335
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.
336 337 338

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

 * 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.
345 346 347 348 349 350

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:
351 352 353 354 355 356 357 358 359 360 361 362

  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
363 364
************************************************************************
*                                                                      *
365
                The main code
Austin Seipp's avatar
Austin Seipp committed
366 367 368
*                                                                      *
************************************************************************
-}
369

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

382 383
        -- We want bndr'' in the envt, because it records
        -- the evaluated-ness of the binder
384 385
       ; return (extendCorePrepEnv env bndr bndr2,
                 addFloat floats new_float) }
386 387 388

cpeBind top_lvl env (Rec pairs)
  = do { let (bndrs,rhss) = unzip pairs
389
       ; (env', bndrs1) <- cpCloneBndrs env (map fst pairs)
390
       ; stuff <- zipWithM (cpePair top_lvl Recursive topDmd False env') bndrs1 rhss
391 392

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

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

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

426 427 428
        -- Wrap floating ticks
       ; let (floats4, rhs4) = wrapTicks floats3 rhs3

429 430 431 432 433
        -- 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
434
       ; let bndr' | exprIsHNF rhs3 = bndr `setIdUnfolding` evaldUnfolding
435
                   | otherwise      = bndr `setIdUnfolding` noUnfolding
436

437
       ; return (floats4, bndr', rhs4) }
438
  where
439 440
    is_strict_or_unlifted = (isStrictDmd dmd) || is_unlifted

441 442
    platform = targetPlatform (cpe_dynFlags env)

443
    arity = idArity bndr        -- We must match this arity
444 445

    ---------------------
446 447 448 449
    float_from_rhs floats rhs
      | isEmptyFloats floats = return (emptyFloats, rhs)
      | isTopLevel top_lvl    = float_top    floats rhs
      | otherwise             = float_nested floats rhs
450 451

    ---------------------
452 453 454 455
    float_nested floats rhs
      | wantFloatNested is_rec is_strict_or_unlifted floats rhs
                  = return (floats, rhs)
      | otherwise = dont_float floats rhs
456 457

    ---------------------
458
    float_top floats rhs        -- Urhgh!  See Note [CafInfo and floating]
459
      | mayHaveCafRefs (idCafInfo bndr)
460 461 462 463
      , allLazyTop floats
      = return (floats, rhs)

      -- So the top-level binding is marked NoCafRefs
464
      | Just (floats', rhs') <- canFloatFromNoCaf platform floats rhs
465
      = return (floats', rhs')
466 467

      | otherwise
468
      = dont_float floats rhs
469 470

    ---------------------
471
    dont_float floats rhs
472 473 474 475
      -- 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
476
      = do { body <- rhsToBodyNF rhs
477
           ; return (emptyFloats, wrapBinds floats body) }
478

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

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

494
-- ---------------------------------------------------------------------------
495
--              CpeRhs: produces a result satisfying CpeRhs
496 497
-- ---------------------------------------------------------------------------

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

507 508
cpeRhsE _env expr@(Type {})      = return (emptyFloats, expr)
cpeRhsE _env expr@(Coercion {})  = return (emptyFloats, expr)
509
cpeRhsE env (Lit (LitInteger i _))
510 511
    = cpeRhsE env (cvtLitInteger (cpe_dynFlags env) (getMkIntegerId env)
                   (cpe_integerSDataCon env) i)
512 513
cpeRhsE _env expr@(Lit {}) = return (emptyFloats, expr)
cpeRhsE env expr@(Var {})  = cpeApp env expr
514

515
cpeRhsE env (Var f `App` _{-type-} `App` arg)
516 517
  | f `hasKey` lazyIdKey          -- Replace (lazy a) by a
  = cpeRhsE env arg               -- See Note [lazyId magic] in MkId
518

519
cpeRhsE env (Var f `App` _runtimeRep `App` _type `App` arg)
520 521 522
    -- See Note [runRW magic] in MkId
  | f `hasKey` runRWKey           -- Replace (runRW# f) by (f realWorld#),
  = case arg of                   -- beta reducing if possible
523
      Lam s body -> cpeRhsE (extendCorePrepEnv env s realWorldPrimId) body
524
      _          -> cpeRhsE env (arg `App` Var realWorldPrimId)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
525 526 527 528 529 530 531 532 533 534 535 536
                    -- 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
-}
537

538 539
cpeRhsE env expr@(App {}) = cpeApp env expr

540 541 542 543 544
cpeRhsE env (Let bind expr)
  = do { (env', new_binds) <- cpeBind NotTopLevel env bind
       ; (floats, body) <- cpeRhsE env' expr
       ; return (new_binds `appendFloats` floats, body) }

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

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
565
        ; (env', bndrs') <- cpCloneBndrs env bndrs
566 567
        ; body' <- cpeBodyNF env' body
        ; return (emptyFloats, mkLams bndrs' body') }
568 569 570 571 572

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

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

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

604
-- ---------------------------------------------------------------------------
605
--              CpeBody: produces a result satisfying CpeBody
606
-- ---------------------------------------------------------------------------
607

608
cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
609
cpeBodyNF env expr
610 611
  = do { (floats, body) <- cpeBody env expr
       ; return (wrapBinds floats body) }
612

613 614 615 616 617 618
--------
cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
cpeBody env expr
  = do { (floats1, rhs) <- cpeRhsE env expr
       ; (floats2, body) <- rhsToBody rhs
       ; return (floats1 `appendFloats` floats2, body) }
619

620 621 622
--------
rhsToBodyNF :: CpeRhs -> UniqSM CpeBody
rhsToBodyNF rhs = do { (floats,body) <- rhsToBody rhs
623
                     ; return (wrapBinds floats body) }
624

625 626
--------
rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
627
-- Remove top level lambdas by let-binding
628

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

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

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

653 654
rhsToBody expr = return (emptyFloats, expr)

655

656 657

-- ---------------------------------------------------------------------------
658
--              CpeApp: produces a result satisfying CpeApp
659 660 661 662
-- ---------------------------------------------------------------------------

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

667
        -- Now deal with the function
668
       ; case head of
669 670 671
           Just (fn_id, depth) -> do { sat_app <- maybeSaturate fn_id app depth
                                     ; return (floats, sat_app) }
           _other              -> return (floats, app) }
672 673

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

680
    collect_args
681
        :: CoreExpr
682 683 684 685 686 687 688
        -> Int                       -- Current app depth
        -> UniqSM (CpeApp,           -- The rebuilt expression
                   Maybe (Id, 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
689 690 691

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

694
    collect_args (App fun arg@(Coercion {})) depth
695
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
696
           ; return (App fun' arg, hd, funResultTy fun_ty, floats, ss) }
697

698 699
    collect_args (App fun arg) depth
      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun (depth+1)
700 701 702 703 704 705 706
           ; 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
707

708
           ; (fs, arg') <- cpeArg env ss1 arg arg_ty
709
           ; return (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest) }
710

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

727
    collect_args (Cast fun co) depth
728
      = do { let Pair _ty1 ty2 = coercionKind co
729 730
           ; (fun', hd, _, floats, ss) <- collect_args fun depth
           ; return (Cast fun' co, hd, ty2, floats, ss) }
731

732
    collect_args (Tick tickish fun) depth
733 734 735 736 737
      | 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) }
738

739
        -- N-variable fun, better let-bind it
740
    collect_args fun _
741 742
      = do { (fun_floats, fun') <- cpeArg env evalDmd fun ty
                          -- The evalDmd says that it's sure to be evaluated,
743
                          -- so we'll end up case-binding it
744
           ; return (fun', Nothing, ty, fun_floats, []) }
745
        where
746
          ty = exprType fun
747

748 749 750 751 752 753 754
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

755
-- ---------------------------------------------------------------------------
756
--      CpeArg: produces a result satisfying CpeArg
757 758 759
-- ---------------------------------------------------------------------------

-- This is where we arrange that a non-trivial argument is let-bound
Austin Seipp's avatar
Austin Seipp committed
760
cpeArg :: CorePrepEnv -> Demand
761
       -> CoreArg -> Type -> UniqSM (Floats, CpeTriv)
762
cpeArg env dmd arg arg_ty
763
  = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
764 765 766 767 768 769
       ; (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
770

771 772 773 774
       ; if cpe_ExprIsTrivial arg2    -- Do not eta expand a trivial argument
         then return (floats2, arg2)
         else do
       { v <- newVar arg_ty
775
       ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
776
             arg_float = mkFloat dmd is_unlifted v arg3
777
       ; return (addFloat floats2 arg_float, varToCoreExpr v) } }
778
  where
779
    is_unlifted = isUnliftedType arg_ty
780 781
    is_strict   = isStrictDmd dmd
    want_float  = wantFloatNested NonRecursive (is_strict || is_unlifted)
782

Austin Seipp's avatar
Austin Seipp committed
783
{-
784 785 786 787 788 789 790 791 792 793 794 795 796
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.)

797 798 799 800 801

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

802 803
maybeSaturate deals with saturating primops and constructors
The type is the type of the entire application
Austin Seipp's avatar
Austin Seipp committed
804
-}
805 806 807

maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
maybeSaturate fn expr n_args