ghc/compiler/coreSyn/CorePrep.lhs

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