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