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