compiler/coreSyn/CorePrep.lhs

   1 %
   2 % (c) The University of Glasgow, 1994-2006
   3 %
   4
   5 Core pass to saturate constructors and PrimOps
   6
   7 \begin{code}
   8 module CorePrep (
   9       corePrepPgm, corePrepExpr
  10   ) where
  11
  12 #include "HsVersions.h"
  13
  14 import CoreUtils hiding (exprIsTrivial)
  15 import CoreFVs
  16 import CoreLint
  17 import CoreSyn
  18 import Type
  19 import Coercion
  20 import TyCon
  21 import NewDemand
  22 import Var
  23 import VarSet
  24 import VarEnv
  25 import Id
  26 import IdInfo
  27 import DataCon
  28 import PrimOp
  29 import BasicTypes
  30 import UniqSupply
  31 import Maybes
  32 import OrdList
  33 import ErrUtils
  34 import DynFlags
  35 import Util
  36 import Outputable
  37 import TysWiredIn
  38 import MkId
  39 import TysPrim
  40 \end{code}
  41
  42 -- ---------------------------------------------------------------------------
  43 -- Overview
  44 -- ---------------------------------------------------------------------------
  45
  46 The goal of this pass is to prepare for code generation.
  47
  48 1.  Saturate constructor and primop applications.
  49
  50 2.  Convert to A-normal form; that is, function arguments
  51     are always variables.
  52
  53     * Use case for strict arguments:
  54         f E ==> case E of x -> f x
  55         (where f is strict)
  56
  57     * Use let for non-trivial lazy arguments
  58         f E ==> let x = E in f x
  59         (were f is lazy and x is non-trivial)
  60
  61 3.  Similarly, convert any unboxed lets into cases.
  62     [I'm experimenting with leaving 'ok-for-speculation'
  63      rhss in let-form right up to this point.]
  64
  65 4.  Ensure that lambdas only occur as the RHS of a binding
  66     (The code generator can't deal with anything else.)
  67
  68 5.  [Not any more; nuked Jun 2002] Do the seq/par munging.
  69
  70 6.  Clone all local Ids.
  71     This means that all such Ids are unique, rather than the
  72     weaker guarantee of no clashes which the simplifier provides.
  73     And that is what the code generator needs.
  74
  75     We don't clone TyVars. The code gen doesn't need that,
  76     and doing so would be tiresome because then we'd need
  77     to substitute in types.
  78
  79
  80 7.  Give each dynamic CCall occurrence a fresh unique; this is
  81     rather like the cloning step above.
  82
  83 8.  Inject bindings for the "implicit" Ids:
  84         * Constructor wrappers
  85         * Constructor workers
  86         * Record selectors
  87     We want curried definitions for all of these in case they
  88     aren't inlined by some caller.
  89
  90 This is all done modulo type applications and abstractions, so that
  91 when type erasure is done for conversion to STG, we don't end up with
  92 any trivial or useless bindings.
  93
  94
  95
  96 -- -----------------------------------------------------------------------------
  97 -- Top level stuff
  98 -- -----------------------------------------------------------------------------
  99
 100 \begin{code}
 101 corePrepPgm :: DynFlags -> [CoreBind] -> [TyCon] -> IO [CoreBind]
 102 corePrepPgm dflags binds data_tycons
 103   = do  showPass dflags "CorePrep"
 104         us <- mkSplitUniqSupply 's'
 105
 106         let implicit_binds = mkDataConWorkers data_tycons
 107                 -- NB: we must feed mkImplicitBinds through corePrep too
 108                 -- so that they are suitably cloned and eta-expanded
 109
 110             binds_out = initUs_ us (
 111                           corePrepTopBinds binds        `thenUs` \ floats1 ->
 112                           corePrepTopBinds implicit_binds       `thenUs` \ floats2 ->
 113                           returnUs (deFloatTop (floats1 `appendFloats` floats2))
 114                         )
 115
 116         endPass dflags "CorePrep" Opt_D_dump_prep binds_out
 117         return binds_out
 118
 119 corePrepExpr :: DynFlags -> CoreExpr -> IO CoreExpr
 120 corePrepExpr dflags expr
 121   = do showPass dflags "CorePrep"
 122        us <- mkSplitUniqSupply 's'
 123        let new_expr = initUs_ us (corePrepAnExpr emptyCorePrepEnv expr)
 124        dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep"
 125                      (ppr new_expr)
 126        return new_expr
 127 \end{code}
 128
 129 -- -----------------------------------------------------------------------------
 130 -- Implicit bindings
 131 -- -----------------------------------------------------------------------------
 132
 133 Create any necessary "implicit" bindings for data con workers.  We
 134 create the rather strange (non-recursive!) binding
 135
 136         $wC = \x y -> $wC x y
 137
 138 i.e. a curried constructor that allocates.  This means that we can
 139 treat the worker for a constructor like any other function in the rest
 140 of the compiler.  The point here is that CoreToStg will generate a
 141 StgConApp for the RHS, rather than a call to the worker (which would
 142 give a loop).  As Lennart says: the ice is thin here, but it works.
 143
 144 Hmm.  Should we create bindings for dictionary constructors?  They are
 145 always fully applied, and the bindings are just there to support
 146 partial applications. But it's easier to let them through.
 147
 148 \begin{code}
 149 mkDataConWorkers data_tycons
 150   = [ NonRec id (Var id)        -- The ice is thin here, but it works
 151     | tycon <- data_tycons,     -- CorePrep will eta-expand it
 152       data_con <- tyConDataCons tycon,
 153       let id = dataConWorkId data_con ]
 154 \end{code}
 155
 156
 157 \begin{code}
 158 -- ---------------------------------------------------------------------------
 159 -- Dealing with bindings
 160 -- ---------------------------------------------------------------------------
 161
 162 data FloatingBind = FloatLet CoreBind
 163                   | FloatCase Id CoreExpr Bool
 164                         -- The bool indicates "ok-for-speculation"
 165
 166 data Floats = Floats OkToSpec (OrdList FloatingBind)
 167
 168 -- Can we float these binds out of the rhs of a let?  We cache this decision
 169 -- to avoid having to recompute it in a non-linear way when there are
 170 -- deeply nested lets.
 171 data OkToSpec
 172    = NotOkToSpec        -- definitely not
 173    | OkToSpec           -- yes
 174    | IfUnboxedOk        -- only if floating an unboxed binding is ok
 175
 176 emptyFloats :: Floats
 177 emptyFloats = Floats OkToSpec nilOL
 178
 179 addFloat :: Floats -> FloatingBind -> Floats
 180 addFloat (Floats ok_to_spec floats) new_float
 181   = Floats (combine ok_to_spec (check new_float)) (floats `snocOL` new_float)
 182   where
 183     check (FloatLet _)                = OkToSpec
 184     check (FloatCase _ _ ok_for_spec)
 185         | ok_for_spec  =  IfUnboxedOk
 186         | otherwise    =  NotOkToSpec
 187         -- The ok-for-speculation flag says that it's safe to
 188         -- float this Case out of a let, and thereby do it more eagerly
 189         -- We need the top-level flag because it's never ok to float
 190         -- an unboxed binding to the top level
 191
 192 unitFloat :: FloatingBind -> Floats
 193 unitFloat = addFloat emptyFloats
 194
 195 appendFloats :: Floats -> Floats -> Floats
 196 appendFloats (Floats spec1 floats1) (Floats spec2 floats2)
 197   = Floats (combine spec1 spec2) (floats1 `appOL` floats2)
 198
 199 concatFloats :: [Floats] -> Floats
 200 concatFloats = foldr appendFloats emptyFloats
 201
 202 combine NotOkToSpec _ = NotOkToSpec
 203 combine _ NotOkToSpec = NotOkToSpec
 204 combine IfUnboxedOk _ = IfUnboxedOk
 205 combine _ IfUnboxedOk = IfUnboxedOk
 206 combine _ _           = OkToSpec
 207
 208 instance Outputable FloatingBind where
 209   ppr (FloatLet bind)        = text "FloatLet" <+> ppr bind
 210   ppr (FloatCase b rhs spec) = text "FloatCase" <+> ppr b <+> ppr spec <+> equals <+> ppr rhs
 211
 212 deFloatTop :: Floats -> [CoreBind]
 213 -- For top level only; we don't expect any FloatCases
 214 deFloatTop (Floats _ floats)
 215   = foldrOL get [] floats
 216   where
 217     get (FloatLet b) bs = b:bs
 218     get b            bs = pprPanic "corePrepPgm" (ppr b)
 219
 220 allLazy :: TopLevelFlag -> RecFlag -> Floats -> Bool
 221 allLazy top_lvl is_rec (Floats ok_to_spec _)
 222   = case ok_to_spec of
 223         OkToSpec    -> True
 224         NotOkToSpec -> False
 225         IfUnboxedOk -> isNotTopLevel top_lvl && isNonRec is_rec
 226
 227 -- ---------------------------------------------------------------------------
 228 --                      Bindings
 229 -- ---------------------------------------------------------------------------
 230
 231 corePrepTopBinds :: [CoreBind] -> UniqSM Floats
 232 corePrepTopBinds binds
 233   = go emptyCorePrepEnv binds
 234   where
 235     go env []             = returnUs emptyFloats
 236     go env (bind : binds) = corePrepTopBind env bind    `thenUs` \ (env', bind') ->
 237                             go env' binds               `thenUs` \ binds' ->
 238                             returnUs (bind' `appendFloats` binds')
 239
 240 -- NB: we do need to float out of top-level bindings
 241 -- Consider     x = length [True,False]
 242 -- We want to get
 243 --              s1 = False : []
 244 --              s2 = True  : s1
 245 --              x  = length s2
 246
 247 -- We return a *list* of bindings, because we may start with
 248 --      x* = f (g y)
 249 -- where x is demanded, in which case we want to finish with
 250 --      a = g y
 251 --      x* = f a
 252 -- And then x will actually end up case-bound
 253 --
 254 -- What happens to the CafInfo on the floated bindings?  By
 255 -- default, all the CafInfos will be set to MayHaveCafRefs,
 256 -- which is safe.
 257 --
 258 -- This might be pessimistic, because eg. s1 & s2
 259 -- might not refer to any CAFs and the GC will end up doing
 260 -- more traversal than is necessary, but it's still better
 261 -- than not floating the bindings at all, because then
 262 -- the GC would have to traverse the structure in the heap
 263 -- instead.  Given this, we decided not to try to get
 264 -- the CafInfo on the floated bindings correct, because
 265 -- it looks difficult.
 266
 267 --------------------------------
 268 corePrepTopBind :: CorePrepEnv -> CoreBind -> UniqSM (CorePrepEnv, Floats)
 269 corePrepTopBind env (NonRec bndr rhs)
 270   = cloneBndr env bndr                                  `thenUs` \ (env', bndr') ->
 271     corePrepRhs TopLevel NonRecursive env (bndr, rhs)   `thenUs` \ (floats, rhs') ->
 272     returnUs (env', addFloat floats (FloatLet (NonRec bndr' rhs')))
 273
 274 corePrepTopBind env (Rec pairs) = corePrepRecPairs TopLevel env pairs
 275
 276 --------------------------------
 277 corePrepBind ::  CorePrepEnv -> CoreBind -> UniqSM (CorePrepEnv, Floats)
 278         -- This one is used for *local* bindings
 279 corePrepBind env (NonRec bndr rhs)
 280   = etaExpandRhs bndr rhs                               `thenUs` \ rhs1 ->
 281     corePrepExprFloat env rhs1                          `thenUs` \ (floats, rhs2) ->
 282     cloneBndr env bndr                                  `thenUs` \ (_, bndr') ->
 283     mkLocalNonRec bndr' (bdrDem bndr) floats rhs2       `thenUs` \ (floats', bndr'') ->
 284         -- We want bndr'' in the envt, because it records
 285         -- the evaluated-ness of the binder
 286     returnUs (extendCorePrepEnv env bndr bndr'', floats')
 287
 288 corePrepBind env (Rec pairs) = corePrepRecPairs NotTopLevel env pairs
 289
 290 --------------------------------
 291 corePrepRecPairs :: TopLevelFlag -> CorePrepEnv
 292                  -> [(Id,CoreExpr)]     -- Recursive bindings
 293                  -> UniqSM (CorePrepEnv, Floats)
 294 -- Used for all recursive bindings, top level and otherwise
 295 corePrepRecPairs lvl env pairs
 296   = cloneBndrs env (map fst pairs)                              `thenUs` \ (env', bndrs') ->
 297     mapAndUnzipUs (corePrepRhs lvl Recursive env') pairs        `thenUs` \ (floats_s, rhss') ->
 298     returnUs (env', unitFloat (FloatLet (Rec (flatten (concatFloats floats_s) bndrs' rhss'))))
 299   where
 300         -- Flatten all the floats, and the currrent
 301         -- group into a single giant Rec
 302     flatten (Floats _ floats) bndrs rhss = foldrOL get (bndrs `zip` rhss) floats
 303
 304     get (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
 305     get (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
 306     get b                       prs2 = pprPanic "corePrepRecPairs" (ppr b)
 307
 308 --------------------------------
 309 corePrepRhs :: TopLevelFlag -> RecFlag
 310             -> CorePrepEnv -> (Id, CoreExpr)
 311             -> UniqSM (Floats, CoreExpr)
 312 -- Used for top-level bindings, and local recursive bindings
 313 corePrepRhs top_lvl is_rec env (bndr, rhs)
 314   = etaExpandRhs bndr rhs       `thenUs` \ rhs' ->
 315     corePrepExprFloat env rhs'  `thenUs` \ floats_w_rhs ->
 316     floatRhs top_lvl is_rec bndr floats_w_rhs
 317
 318
 319 -- ---------------------------------------------------------------------------
 320 -- Making arguments atomic (function args & constructor args)
 321 -- ---------------------------------------------------------------------------
 322
 323 -- This is where we arrange that a non-trivial argument is let-bound
 324 corePrepArg :: CorePrepEnv -> CoreArg -> RhsDemand
 325            -> UniqSM (Floats, CoreArg)
 326 corePrepArg env arg dem
 327   = corePrepExprFloat env arg           `thenUs` \ (floats, arg') ->
 328     if exprIsTrivial arg'
 329     then returnUs (floats, arg')
 330     else newVar (exprType arg')                 `thenUs` \ v ->
 331          mkLocalNonRec v dem floats arg'        `thenUs` \ (floats', v') ->
 332          returnUs (floats', Var v')
 333
 334 -- version that doesn't consider an scc annotation to be trivial.
 335 exprIsTrivial (Var v)                  = True
 336 exprIsTrivial (Type _)                 = True
 337 exprIsTrivial (Lit lit)                = True
 338 exprIsTrivial (App e arg)              = isTypeArg arg && exprIsTrivial e
 339 exprIsTrivial (Note (SCC _) e)         = False
 340 exprIsTrivial (Note _ e)               = exprIsTrivial e
 341 exprIsTrivial (Cast e co)              = exprIsTrivial e
 342 exprIsTrivial (Lam b body) | isTyVar b = exprIsTrivial body
 343 exprIsTrivial other                    = False
 344
 345 -- ---------------------------------------------------------------------------
 346 -- Dealing with expressions
 347 -- ---------------------------------------------------------------------------
 348
 349 corePrepAnExpr :: CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
 350 corePrepAnExpr env expr
 351   = corePrepExprFloat env expr          `thenUs` \ (floats, expr) ->
 352     mkBinds floats expr
 353
 354
 355 corePrepExprFloat :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
 356 -- If
 357 --      e  ===>  (bs, e')
 358 -- then
 359 --      e = let bs in e'        (semantically, that is!)
 360 --
 361 -- For example
 362 --      f (g x)   ===>   ([v = g x], f v)
 363
 364 corePrepExprFloat env (Var v)
 365   = fiddleCCall v                               `thenUs` \ v1 ->
 366     let
 367         v2 = lookupCorePrepEnv env v1
 368     in
 369     maybeSaturate v2 (Var v2) 0 emptyFloats (idType v2)
 370
 371 corePrepExprFloat env expr@(Type _)
 372   = returnUs (emptyFloats, expr)
 373
 374 corePrepExprFloat env expr@(Lit lit)
 375   = returnUs (emptyFloats, expr)
 376
 377 corePrepExprFloat env (Let bind body)
 378   = corePrepBind env bind               `thenUs` \ (env', new_binds) ->
 379     corePrepExprFloat env' body         `thenUs` \ (floats, new_body) ->
 380     returnUs (new_binds `appendFloats` floats, new_body)
 381
 382 corePrepExprFloat env (Note n@(SCC _) expr)
 383   = corePrepAnExpr env expr             `thenUs` \ expr1 ->
 384     deLamFloat expr1                    `thenUs` \ (floats, expr2) ->
 385     returnUs (floats, Note n expr2)
 386
 387 corePrepExprFloat env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
 388   | Just (TickBox {}) <- isTickBoxOp_maybe id
 389   = corePrepAnExpr env expr            `thenUs` \ expr1 ->
 390     deLamFloat expr1                   `thenUs` \ (floats, expr2) ->
 391     return (floats, Case (Var id) bndr ty [(DEFAULT,[],expr2)])
 392
 393 -- Translate Binary tickBox into standard tickBox
 394 corePrepExprFloat env (App (Var id) expr)
 395   | Just (BinaryTickBox m t e) <- isTickBoxOp_maybe id
 396   = corePrepAnExpr env expr            `thenUs` \ expr1 ->
 397     deLamFloat expr1                   `thenUs` \ (floats, expr2) ->
 398     getUniqueUs                        `thenUs` \ u1 ->
 399     getUniqueUs                        `thenUs` \ u2 ->
 400     getUniqueUs                        `thenUs` \ u3 ->
 401     getUniqueUs                        `thenUs` \ u4 ->
 402     getUniqueUs                        `thenUs` \ u5 ->
 403     let bndr1 = mkSysLocal FSLIT("t1") u1 boolTy in
 404     let bndr2 = mkSysLocal FSLIT("t2") u2 realWorldStatePrimTy in
 405     let bndr3 = mkSysLocal FSLIT("t3") u3 realWorldStatePrimTy in
 406     let tick_e  = mkTickBoxOpId u4 m e in
 407     let tick_t  = mkTickBoxOpId u5 m t in
 408     return (floats, Case expr2
 409                        bndr1
 410                        boolTy
 411                        [ (DataAlt falseDataCon, [],
 412                           Case (Var tick_e) bndr2 boolTy [(DEFAULT,[],Var falseDataConId)])
 413                        , (DataAlt trueDataCon,  [],
 414                           Case (Var tick_t) bndr3 boolTy [(DEFAULT,[],Var trueDataConId)])
 415                        ])
 416
 417 corePrepExprFloat env (Note other_note expr)
 418   = corePrepExprFloat env expr          `thenUs` \ (floats, expr') ->
 419     returnUs (floats, Note other_note expr')
 420
 421 corePrepExprFloat env (Cast expr co)
 422   = corePrepExprFloat env expr          `thenUs` \ (floats, expr') ->
 423     returnUs (floats, Cast expr' co)
 424
 425 corePrepExprFloat env expr@(Lam _ _)
 426   = cloneBndrs env bndrs                `thenUs` \ (env', bndrs') ->
 427     corePrepAnExpr env' body            `thenUs` \ body' ->
 428     returnUs (emptyFloats, mkLams bndrs' body')
 429   where
 430     (bndrs,body) = collectBinders expr
 431
 432 -- This is an (important) optimization.
 433 -- case <btick,A,B>  e of { T -> e1 ; F -> e2 }
 434 --  ==> case e of { T -> <tick,A> e1 ; F -> <tick,B> e2 }
 435 -- This could move into the simplifier.
 436
 437 corePrepExprFloat env (Case (App (Var id) expr) bndr ty alts)
 438   | Just (BinaryTickBox m t e) <- isTickBoxOp_maybe id
 439   = getUniqueUs                        `thenUs` \ u1 ->
 440     getUniqueUs                        `thenUs` \ u2 ->
 441     getUniqueUs                        `thenUs` \ u3 ->
 442     getUniqueUs                        `thenUs` \ u4 ->
 443     getUniqueUs                        `thenUs` \ u5 ->
 444     let bndr1 = mkSysLocal FSLIT("t1") u1 boolTy in
 445     let bndr2 = mkSysLocal FSLIT("t2") u2 realWorldStatePrimTy in
 446     let bndr3 = mkSysLocal FSLIT("t3") u3 realWorldStatePrimTy in
 447     let tick_e  = mkTickBoxOpId u4 m e in
 448     let tick_t  = mkTickBoxOpId u5 m t in
 449     ASSERT (exprType expr `coreEqType` boolTy)
 450     corePrepExprFloat env $
 451                 Case expr
 452                        bndr1
 453                        ty
 454                        [ (DataAlt falseDataCon, [],
 455                           Case (Var tick_e) bndr2 ty [(DEFAULT,[],falseBranch)])
 456                        , (DataAlt trueDataCon,  [],
 457                           Case (Var tick_t) bndr3 ty [(DEFAULT,[],trueBranch)])
 458                        ]
 459
 460    where
 461        (_,_,trueBranch)  = findAlt (DataAlt trueDataCon) alts
 462        (_,_,falseBranch) = findAlt (DataAlt falseDataCon) alts
 463
 464 corePrepExprFloat env (Case scrut bndr ty alts)
 465   = corePrepExprFloat env scrut         `thenUs` \ (floats1, scrut1) ->
 466     deLamFloat scrut1                   `thenUs` \ (floats2, scrut2) ->
 467     let
 468         bndr1 = bndr `setIdUnfolding` evaldUnfolding
 469         -- Record that the case binder is evaluated in the alternatives
 470     in
 471     cloneBndr env bndr1                 `thenUs` \ (env', bndr2) ->
 472     mapUs (sat_alt env') alts           `thenUs` \ alts' ->
 473     returnUs (floats1 `appendFloats` floats2 , Case scrut2 bndr2 ty alts')
 474   where
 475     sat_alt env (con, bs, rhs)
 476           = cloneBndrs env bs           `thenUs` \ (env2, bs') ->
 477             corePrepAnExpr env2 rhs     `thenUs` \ rhs1 ->
 478             deLam rhs1                  `thenUs` \ rhs2 ->
 479             returnUs (con, bs', rhs2)
 480
 481 corePrepExprFloat env expr@(App _ _)
 482   = collect_args expr 0  `thenUs` \ (app, (head,depth), ty, floats, ss) ->
 483     ASSERT(null ss)     -- make sure we used all the strictness info
 484
 485         -- Now deal with the function
 486     case head of
 487       Var fn_id -> maybeSaturate fn_id app depth floats ty
 488       _other    -> returnUs (floats, app)
 489
 490   where
 491
 492     -- Deconstruct and rebuild the application, floating any non-atomic
 493     -- arguments to the outside.  We collect the type of the expression,
 494     -- the head of the application, and the number of actual value arguments,
 495     -- all of which are used to possibly saturate this application if it
 496     -- has a constructor or primop at the head.
 497
 498     collect_args
 499         :: CoreExpr
 500         -> Int                            -- current app depth
 501         -> UniqSM (CoreExpr,              -- the rebuilt expression
 502                    (CoreExpr,Int),        -- the head of the application,
 503                                           -- and no. of args it was applied to
 504                    Type,                  -- type of the whole expr
 505                    Floats,                -- any floats we pulled out
 506                    [Demand])              -- remaining argument demands
 507
 508     collect_args (App fun arg@(Type arg_ty)) depth
 509         = collect_args fun depth   `thenUs` \ (fun',hd,fun_ty,floats,ss) ->
 510           returnUs (App fun' arg, hd, applyTy fun_ty arg_ty, floats, ss)
 511
 512     collect_args (App fun arg) depth
 513         = collect_args fun (depth+1)   `thenUs` \ (fun',hd,fun_ty,floats,ss) ->
 514           let
 515               (ss1, ss_rest)   = case ss of
 516                                    (ss1:ss_rest) -> (ss1,     ss_rest)
 517                                    []            -> (lazyDmd, [])
 518               (arg_ty, res_ty) = expectJust "corePrepExprFloat:collect_args" $
 519                                  splitFunTy_maybe fun_ty
 520           in
 521           corePrepArg env arg (mkDemTy ss1 arg_ty)      `thenUs` \ (fs, arg') ->
 522           returnUs (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest)
 523
 524     collect_args (Var v) depth
 525         = fiddleCCall v `thenUs` \ v1 ->
 526           let
 527                 v2 = lookupCorePrepEnv env v1
 528           in
 529           returnUs (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts)
 530         where
 531           stricts = case idNewStrictness v of
 532                         StrictSig (DmdType _ demands _)
 533                             | listLengthCmp demands depth /= GT -> demands
 534                                     -- length demands <= depth
 535                             | otherwise                         -> []
 536                 -- If depth < length demands, then we have too few args to
 537                 -- satisfy strictness  info so we have to  ignore all the
 538                 -- strictness info, e.g. + (error "urk")
 539                 -- Here, we can't evaluate the arg strictly, because this
 540                 -- partial application might be seq'd
 541
 542     collect_args (Cast fun co) depth
 543         = let (_ty1,ty2) = coercionKind co in
 544           collect_args fun depth  `thenUs` \ (fun', hd, fun_ty, floats, ss) ->
 545           returnUs (Cast fun' co, hd, ty2, floats, ss)
 546
 547     collect_args (Note note fun) depth
 548         | ignore_note note      -- Drop these notes altogether
 549                                 -- They aren't used by the code generator
 550         = collect_args fun depth   `thenUs` \ (fun', hd, fun_ty, floats, ss) ->
 551           returnUs (fun', hd, fun_ty, floats, ss)
 552
 553         -- N-variable fun, better let-bind it
 554         -- ToDo: perhaps we can case-bind rather than let-bind this closure,
 555         -- since it is sure to be evaluated.
 556     collect_args fun depth
 557         = corePrepExprFloat env fun                     `thenUs` \ (fun_floats, fun') ->
 558           newVar ty                                     `thenUs` \ fn_id ->
 559           mkLocalNonRec fn_id onceDem fun_floats fun'   `thenUs` \ (floats, fn_id') ->
 560           returnUs (Var fn_id', (Var fn_id', depth), ty, floats, [])
 561         where
 562           ty = exprType fun
 563
 564     ignore_note (CoreNote _) = True
 565     ignore_note InlineMe     = True
 566     ignore_note _other       = False
 567         -- We don't ignore SCCs, since they require some code generation
 568
 569 ------------------------------------------------------------------------------
 570 -- Building the saturated syntax
 571 -- ---------------------------------------------------------------------------
 572
 573 -- maybeSaturate deals with saturating primops and constructors
 574 -- The type is the type of the entire application
 575 maybeSaturate :: Id -> CoreExpr -> Int -> Floats -> Type -> UniqSM (Floats, CoreExpr)
 576 maybeSaturate fn expr n_args floats ty
 577   | Just DataToTagOp <- isPrimOpId_maybe fn     -- DataToTag must have an evaluated arg
 578                                                 -- A gruesome special case
 579   = saturate_it         `thenUs` \ sat_expr ->
 580
 581         -- OK, now ensure that the arg is evaluated.
 582         -- But (sigh) take into account the lambdas we've now introduced
 583     let
 584         (eta_bndrs, eta_body) = collectBinders sat_expr
 585     in
 586     eval_data2tag_arg eta_body  `thenUs` \ (eta_floats, eta_body') ->
 587     if null eta_bndrs then
 588         returnUs (floats `appendFloats` eta_floats, eta_body')
 589     else
 590         mkBinds eta_floats eta_body'            `thenUs` \ eta_body'' ->
 591         returnUs (floats, mkLams eta_bndrs eta_body'')
 592
 593   | hasNoBinding fn = saturate_it       `thenUs` \ sat_expr ->
 594                       returnUs (floats, sat_expr)
 595
 596   | otherwise       = returnUs (floats, expr)
 597
 598   where
 599     fn_arity     = idArity fn
 600     excess_arity = fn_arity - n_args
 601
 602     saturate_it :: UniqSM CoreExpr
 603     saturate_it | excess_arity == 0 = returnUs expr
 604                 | otherwise         = getUniquesUs              `thenUs` \ us ->
 605                                       returnUs (etaExpand excess_arity us expr ty)
 606
 607         -- Ensure that the argument of DataToTagOp is evaluated
 608     eval_data2tag_arg :: CoreExpr -> UniqSM (Floats, CoreExpr)
 609     eval_data2tag_arg app@(fun `App` arg)
 610         | exprIsHNF arg         -- Includes nullary constructors
 611         = returnUs (emptyFloats, app)   -- The arg is evaluated
 612         | otherwise                     -- Arg not evaluated, so evaluate it
 613         = newVar (exprType arg)         `thenUs` \ arg_id ->
 614           let
 615              arg_id1 = setIdUnfolding arg_id evaldUnfolding
 616           in
 617           returnUs (unitFloat (FloatCase arg_id1 arg False ),
 618                     fun `App` Var arg_id1)
 619
 620     eval_data2tag_arg (Note note app)   -- Scc notes can appear
 621         = eval_data2tag_arg app         `thenUs` \ (floats, app') ->
 622           returnUs (floats, Note note app')
 623
 624     eval_data2tag_arg other     -- Should not happen
 625         = pprPanic "eval_data2tag" (ppr other)
 626
 627
 628 -- ---------------------------------------------------------------------------
 629 -- Precipitating the floating bindings
 630 -- ---------------------------------------------------------------------------
 631
 632 floatRhs :: TopLevelFlag -> RecFlag
 633          -> Id
 634          -> (Floats, CoreExpr)  -- Rhs: let binds in body
 635          -> UniqSM (Floats,     -- Floats out of this bind
 636                     CoreExpr)   -- Final Rhs
 637
 638 floatRhs top_lvl is_rec bndr (floats, rhs)
 639   | isTopLevel top_lvl || exprIsHNF rhs,        -- Float to expose value or
 640     allLazy top_lvl is_rec floats               -- at top level
 641   =     -- Why the test for allLazy?
 642         --      v = f (x `divInt#` y)
 643         -- we don't want to float the case, even if f has arity 2,
 644         -- because floating the case would make it evaluated too early
 645     returnUs (floats, rhs)
 646
 647   | otherwise
 648         -- Don't float; the RHS isn't a value
 649   = mkBinds floats rhs          `thenUs` \ rhs' ->
 650     returnUs (emptyFloats, rhs')
 651
 652 -- mkLocalNonRec is used only for *nested*, *non-recursive* bindings
 653 mkLocalNonRec :: Id  -> RhsDemand       -- Lhs: id with demand
 654               -> Floats -> CoreExpr     -- Rhs: let binds in body
 655               -> UniqSM (Floats, Id)    -- The new Id may have an evaldUnfolding,
 656                                         -- to record that it's been evaluated
 657
 658 mkLocalNonRec bndr dem floats rhs
 659   | isUnLiftedType (idType bndr)
 660         -- If this is an unlifted binding, we always make a case for it.
 661   = ASSERT( not (isUnboxedTupleType (idType bndr)) )
 662     let
 663         float = FloatCase bndr rhs (exprOkForSpeculation rhs)
 664     in
 665     returnUs (addFloat floats float, evald_bndr)
 666
 667   | isStrict dem
 668         -- It's a strict let so we definitely float all the bindings
 669  = let          -- Don't make a case for a value binding,
 670                 -- even if it's strict.  Otherwise we get
 671                 --      case (\x -> e) of ...!
 672         float | exprIsHNF rhs = FloatLet (NonRec bndr rhs)
 673               | otherwise       = FloatCase bndr rhs (exprOkForSpeculation rhs)
 674     in
 675     returnUs (addFloat floats float, evald_bndr)
 676
 677   | otherwise
 678   = floatRhs NotTopLevel NonRecursive bndr (floats, rhs)        `thenUs` \ (floats', rhs') ->
 679     returnUs (addFloat floats' (FloatLet (NonRec bndr rhs')),
 680               if exprIsHNF rhs' then evald_bndr else bndr)
 681
 682   where
 683     evald_bndr = bndr `setIdUnfolding` evaldUnfolding
 684         -- Record if the binder is evaluated
 685
 686
 687 mkBinds :: Floats -> CoreExpr -> UniqSM CoreExpr
 688 mkBinds (Floats _ binds) body
 689   | isNilOL binds = returnUs body
 690   | otherwise     = deLam body          `thenUs` \ body' ->
 691                         -- Lambdas are not allowed as the body of a 'let'
 692                     returnUs (foldrOL mk_bind body' binds)
 693   where
 694     mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
 695     mk_bind (FloatLet bind)        body = Let bind body
 696
 697 etaExpandRhs bndr rhs
 698   =     -- Eta expand to match the arity claimed by the binder
 699         -- Remember, after CorePrep we must not change arity
 700         --
 701         -- Eta expansion might not have happened already,
 702         -- because it is done by the simplifier only when
 703         -- there at least one lambda already.
 704         --
 705         -- NB1:we could refrain when the RHS is trivial (which can happen
 706         --     for exported things).  This would reduce the amount of code
 707         --     generated (a little) and make things a little words for
 708         --     code compiled without -O.  The case in point is data constructor
 709         --     wrappers.
 710         --
 711         -- NB2: we have to be careful that the result of etaExpand doesn't
 712         --    invalidate any of the assumptions that CorePrep is attempting
 713         --    to establish.  One possible cause is eta expanding inside of
 714         --    an SCC note - we're now careful in etaExpand to make sure the
 715         --    SCC is pushed inside any new lambdas that are generated.
 716         --
 717         -- NB3: It's important to do eta expansion, and *then* ANF-ising
 718         --              f = /\a -> g (h 3)      -- h has arity 2
 719         -- If we ANF first we get
 720         --              f = /\a -> let s = h 3 in g s
 721         -- and now eta expansion gives
 722         --              f = /\a -> \ y -> (let s = h 3 in g s) y
 723         -- which is horrible.
 724         -- Eta expanding first gives
 725         --              f = /\a -> \y -> let s = h 3 in g s y
 726         --
 727     getUniquesUs                `thenUs` \ us ->
 728     returnUs (etaExpand arity us rhs (idType bndr))
 729   where
 730         -- For a GlobalId, take the Arity from the Id.
 731         -- It was set in CoreTidy and must not change
 732         -- For all others, just expand at will
 733     arity | isGlobalId bndr = idArity bndr
 734           | otherwise       = exprArity rhs
 735
 736 -- ---------------------------------------------------------------------------
 737 -- Eliminate Lam as a non-rhs (STG doesn't have such a thing)
 738 -- We arrange that they only show up as the RHS of a let(rec)
 739 -- ---------------------------------------------------------------------------
 740
 741 deLam :: CoreExpr -> UniqSM CoreExpr
 742 -- Takes an expression that may be a lambda,
 743 -- and returns one that definitely isn't:
 744 --      (\x.e) ==>  let f = \x.e in f
 745 deLam expr =
 746   deLamFloat expr   `thenUs` \ (floats, expr) ->
 747   mkBinds floats expr
 748
 749
 750 deLamFloat :: CoreExpr -> UniqSM (Floats, CoreExpr)
 751 -- Remove top level lambdas by let-bindinig
 752
 753 deLamFloat (Note n expr)
 754   =     -- You can get things like
 755         --      case e of { p -> coerce t (\s -> ...) }
 756     deLamFloat expr     `thenUs` \ (floats, expr') ->
 757     returnUs (floats, Note n expr')
 758
 759 deLamFloat (Cast e co)
 760   = deLamFloat e        `thenUs` \ (floats, e') ->
 761     returnUs (floats, Cast e' co)
 762
 763 deLamFloat expr
 764   | null bndrs = returnUs (emptyFloats, expr)
 765   | otherwise
 766   = case tryEta bndrs body of
 767       Just no_lam_result -> returnUs (emptyFloats, no_lam_result)
 768       Nothing            -> newVar (exprType expr)      `thenUs` \ fn ->
 769                             returnUs (unitFloat (FloatLet (NonRec fn expr)),
 770                                       Var fn)
 771   where
 772     (bndrs,body) = collectBinders expr
 773
 774 -- Why try eta reduction?  Hasn't the simplifier already done eta?
 775 -- But the simplifier only eta reduces if that leaves something
 776 -- trivial (like f, or f Int).  But for deLam it would be enough to
 777 -- get to a partial application:
 778 --      \xs. map f xs ==> map f
 779
 780 tryEta bndrs expr@(App _ _)
 781   | ok_to_eta_reduce f &&
 782     n_remaining >= 0 &&
 783     and (zipWith ok bndrs last_args) &&
 784     not (any (`elemVarSet` fvs_remaining) bndrs)
 785   = Just remaining_expr
 786   where
 787     (f, args) = collectArgs expr
 788     remaining_expr = mkApps f remaining_args
 789     fvs_remaining = exprFreeVars remaining_expr
 790     (remaining_args, last_args) = splitAt n_remaining args
 791     n_remaining = length args - length bndrs
 792
 793     ok bndr (Var arg) = bndr == arg
 794     ok bndr other     = False
 795
 796           -- we can't eta reduce something which must be saturated.
 797     ok_to_eta_reduce (Var f) = not (hasNoBinding f)
 798     ok_to_eta_reduce _       = False --safe. ToDo: generalise
 799
 800 tryEta bndrs (Let bind@(NonRec b r) body)
 801   | not (any (`elemVarSet` fvs) bndrs)
 802   = case tryEta bndrs body of
 803         Just e -> Just (Let bind e)
 804         Nothing -> Nothing
 805   where
 806     fvs = exprFreeVars r
 807
 808 tryEta bndrs _ = Nothing
 809 \end{code}
 810
 811
 812 -- -----------------------------------------------------------------------------
 813 -- Demands
 814 -- -----------------------------------------------------------------------------
 815
 816 \begin{code}
 817 data RhsDemand
 818      = RhsDemand { isStrict :: Bool,  -- True => used at least once
 819                    isOnceDem   :: Bool   -- True => used at most once
 820                  }
 821
 822 mkDem :: Demand -> Bool -> RhsDemand
 823 mkDem strict once = RhsDemand (isStrictDmd strict) once
 824
 825 mkDemTy :: Demand -> Type -> RhsDemand
 826 mkDemTy strict ty = RhsDemand (isStrictDmd strict)
 827                               False {- For now -}
 828
 829 bdrDem :: Id -> RhsDemand
 830 bdrDem id = mkDem (idNewDemandInfo id)
 831                   False {- For now -}
 832
 833 -- safeDem :: RhsDemand
 834 -- safeDem = RhsDemand False False  -- always safe to use this
 835
 836 onceDem :: RhsDemand
 837 onceDem = RhsDemand False True   -- used at most once
 838 \end{code}
 839
 840
 841
 842
 843 %************************************************************************
 844 %*                                                                      *
 845 \subsection{Cloning}
 846 %*                                                                      *
 847 %************************************************************************
 848
 849 \begin{code}
 850 -- ---------------------------------------------------------------------------
 851 --                      The environment
 852 -- ---------------------------------------------------------------------------
 853
 854 data CorePrepEnv = CPE (IdEnv Id)       -- Clone local Ids
 855
 856 emptyCorePrepEnv :: CorePrepEnv
 857 emptyCorePrepEnv = CPE emptyVarEnv
 858
 859 extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
 860 extendCorePrepEnv (CPE env) id id' = CPE (extendVarEnv env id id')
 861
 862 lookupCorePrepEnv :: CorePrepEnv -> Id -> Id
 863 lookupCorePrepEnv (CPE env) id
 864   = case lookupVarEnv env id of
 865         Nothing  -> id
 866         Just id' -> id'
 867
 868 ------------------------------------------------------------------------------
 869 -- Cloning binders
 870 -- ---------------------------------------------------------------------------
 871
 872 cloneBndrs :: CorePrepEnv -> [Var] -> UniqSM (CorePrepEnv, [Var])
 873 cloneBndrs env bs = mapAccumLUs cloneBndr env bs
 874
 875 cloneBndr  :: CorePrepEnv -> Var -> UniqSM (CorePrepEnv, Var)
 876 cloneBndr env bndr
 877   | isLocalId bndr
 878   = getUniqueUs   `thenUs` \ uniq ->
 879     let
 880         bndr' = setVarUnique bndr uniq
 881     in
 882     returnUs (extendCorePrepEnv env bndr bndr', bndr')
 883
 884   | otherwise   -- Top level things, which we don't want
 885                 -- to clone, have become GlobalIds by now
 886                 -- And we don't clone tyvars
 887   = returnUs (env, bndr)
 888
 889
 890 ------------------------------------------------------------------------------
 891 -- Cloning ccall Ids; each must have a unique name,
 892 -- to give the code generator a handle to hang it on
 893 -- ---------------------------------------------------------------------------
 894
 895 fiddleCCall :: Id -> UniqSM Id
 896 fiddleCCall id
 897   | isFCallId id = getUniqueUs          `thenUs` \ uniq ->
 898                    returnUs (id `setVarUnique` uniq)
 899   | otherwise    = returnUs id
 900
 901 ------------------------------------------------------------------------------
 902 -- Generating new binders
 903 -- ---------------------------------------------------------------------------
 904
 905 newVar :: Type -> UniqSM Id
 906 newVar ty
 907  = seqType ty                   `seq`
 908    getUniqueUs                  `thenUs` \ uniq ->
 909    returnUs (mkSysLocal FSLIT("sat") uniq ty)
 910 \end{code}