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 PrelNames        ( lazyIdKey, hasKey )
  15 import CoreUtils
  16 import CoreArity
  17 import CoreFVs
  18 import CoreMonad        ( endPass, CoreToDo(..) )
  19 import CoreSyn
  20 import CoreSubst
  21 import Type
  22 import Coercion
  23 import TyCon
  24 import Demand
  25 import Var
  26 import VarSet
  27 import VarEnv
  28 import Id
  29 import IdInfo
  30 import DataCon
  31 import PrimOp
  32 import BasicTypes
  33 import UniqSupply
  34 import Maybes
  35 import OrdList
  36 import ErrUtils
  37 import DynFlags
  38 import Util
  39 import Outputable
  40 import MonadUtils
  41 import FastString
  42 import Data.List        ( mapAccumL )
  43 import Control.Monad
  44 \end{code}
  45
  46 -- ---------------------------------------------------------------------------
  47 -- Overview
  48 -- ---------------------------------------------------------------------------
  49
  50 The goal of this pass is to prepare for code generation.
  51
  52 1.  Saturate constructor and primop applications.
  53
  54 2.  Convert to A-normal form; that is, function arguments
  55     are always variables.
  56
  57     * Use case for strict arguments:
  58         f E ==> case E of x -> f x
  59         (where f is strict)
  60
  61     * Use let for non-trivial lazy arguments
  62         f E ==> let x = E in f x
  63         (were f is lazy and x is non-trivial)
  64
  65 3.  Similarly, convert any unboxed lets into cases.
  66     [I'm experimenting with leaving 'ok-for-speculation'
  67      rhss in let-form right up to this point.]
  68
  69 4.  Ensure that *value* lambdas only occur as the RHS of a binding
  70     (The code generator can't deal with anything else.)
  71     Type lambdas are ok, however, because the code gen discards them.
  72
  73 5.  [Not any more; nuked Jun 2002] Do the seq/par munging.
  74
  75 6.  Clone all local Ids.
  76     This means that all such Ids are unique, rather than the
  77     weaker guarantee of no clashes which the simplifier provides.
  78     And that is what the code generator needs.
  79
  80     We don't clone TyVars. The code gen doesn't need that,
  81     and doing so would be tiresome because then we'd need
  82     to substitute in types.
  83
  84
  85 7.  Give each dynamic CCall occurrence a fresh unique; this is
  86     rather like the cloning step above.
  87
  88 8.  Inject bindings for the "implicit" Ids:
  89         * Constructor wrappers
  90         * Constructor workers
  91     We want curried definitions for all of these in case they
  92     aren't inlined by some caller.
  93
  94 9.  Replace (lazy e) by e.  See Note [lazyId magic] in MkId.lhs
  95
  96 This is all done modulo type applications and abstractions, so that
  97 when type erasure is done for conversion to STG, we don't end up with
  98 any trivial or useless bindings.
  99
 100
 101 Invariants
 102 ~~~~~~~~~~
 103 Here is the syntax of the Core produced by CorePrep:
 104
 105     Trivial expressions
 106        triv ::= lit |  var  | triv ty  |  /\a. triv  |  triv |> co
 107
 108     Applications
 109        app ::= lit  |  var  |  app triv  |  app ty  |  app |> co
 110
 111     Expressions
 112        body ::= app
 113               | let(rec) x = rhs in body     -- Boxed only
 114               | case body of pat -> body
 115               | /\a. body
 116               | body |> co
 117
 118     Right hand sides (only place where lambdas can occur)
 119        rhs ::= /\a.rhs  |  \x.rhs  |  body
 120
 121 We define a synonym for each of these non-terminals.  Functions
 122 with the corresponding name produce a result in that syntax.
 123
 124 \begin{code}
 125 type CpeTriv = CoreExpr    -- Non-terminal 'triv'
 126 type CpeApp  = CoreExpr    -- Non-terminal 'app'
 127 type CpeBody = CoreExpr    -- Non-terminal 'body'
 128 type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'
 129 \end{code}
 130
 131 %************************************************************************
 132 %*                                                                      *
 133                 Top level stuff
 134 %*                                                                      *
 135 %************************************************************************
 136
 137 \begin{code}
 138 corePrepPgm :: DynFlags -> [CoreBind] -> [TyCon] -> IO [CoreBind]
 139 corePrepPgm dflags binds data_tycons = do
 140     showPass dflags "CorePrep"
 141     us <- mkSplitUniqSupply 's'
 142
 143     let implicit_binds = mkDataConWorkers data_tycons
 144             -- NB: we must feed mkImplicitBinds through corePrep too
 145             -- so that they are suitably cloned and eta-expanded
 146
 147         binds_out = initUs_ us $ do
 148                       floats1 <- corePrepTopBinds binds
 149                       floats2 <- corePrepTopBinds implicit_binds
 150                       return (deFloatTop (floats1 `appendFloats` floats2))
 151
 152     endPass dflags CorePrep binds_out []
 153     return binds_out
 154
 155 corePrepExpr :: DynFlags -> CoreExpr -> IO CoreExpr
 156 corePrepExpr dflags expr = do
 157     showPass dflags "CorePrep"
 158     us <- mkSplitUniqSupply 's'
 159     let new_expr = initUs_ us (cpeBodyNF emptyCorePrepEnv expr)
 160     dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep" (ppr new_expr)
 161     return new_expr
 162
 163 corePrepTopBinds :: [CoreBind] -> UniqSM Floats
 164 -- Note [Floating out of top level bindings]
 165 corePrepTopBinds binds
 166   = go emptyCorePrepEnv binds
 167   where
 168     go _   []             = return emptyFloats
 169     go env (bind : binds) = do (env', bind') <- cpeBind TopLevel env bind
 170                                binds' <- go env' binds
 171                                return (bind' `appendFloats` binds')
 172
 173 mkDataConWorkers :: [TyCon] -> [CoreBind]
 174 -- See Note [Data constructor workers]
 175 mkDataConWorkers data_tycons
 176   = [ NonRec id (Var id)        -- The ice is thin here, but it works
 177     | tycon <- data_tycons,     -- CorePrep will eta-expand it
 178       data_con <- tyConDataCons tycon,
 179       let id = dataConWorkId data_con ]
 180 \end{code}
 181
 182 Note [Floating out of top level bindings]
 183 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 184 NB: we do need to float out of top-level bindings
 185 Consider        x = length [True,False]
 186 We want to get
 187                 s1 = False : []
 188                 s2 = True  : s1
 189                 x  = length s2
 190
 191 We return a *list* of bindings, because we may start with
 192         x* = f (g y)
 193 where x is demanded, in which case we want to finish with
 194         a = g y
 195         x* = f a
 196 And then x will actually end up case-bound
 197
 198 Note [CafInfo and floating]
 199 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 200 What happens when we try to float bindings to the top level?  At this
 201 point all the CafInfo is supposed to be correct, and we must make certain
 202 that is true of the new top-level bindings.  There are two cases
 203 to consider
 204
 205 a) The top-level binding is marked asCafRefs.  In that case we are
 206    basically fine.  The floated bindings had better all be lazy lets,
 207    so they can float to top level, but they'll all have HasCafRefs
 208    (the default) which is safe.
 209
 210 b) The top-level binding is marked NoCafRefs.  This really happens
 211    Example.  CoreTidy produces
 212       $fApplicativeSTM [NoCafRefs] = D:Alternative retry# ...blah...
 213    Now CorePrep has to eta-expand to
 214       $fApplicativeSTM = let sat = \xy. retry x y
 215                          in D:Alternative sat ...blah...
 216    So what we *want* is
 217       sat [NoCafRefs] = \xy. retry x y
 218       $fApplicativeSTM [NoCafRefs] = D:Alternative sat ...blah...
 219
 220    So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
 221    *and* substutite the modified 'sat' into the old RHS.
 222
 223    It should be the case that 'sat' is itself [NoCafRefs] (a value, no
 224    cafs) else the original top-level binding would not itself have been
 225    marked [NoCafRefs].  The DEBUG check in CoreToStg for
 226    consistentCafInfo will find this.
 227
 228 This is all very gruesome and horrible. It would be better to figure
 229 out CafInfo later, after CorePrep.  We'll do that in due course.
 230 Meanwhile this horrible hack works.
 231
 232
 233 Note [Data constructor workers]
 234 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 235 Create any necessary "implicit" bindings for data con workers.  We
 236 create the rather strange (non-recursive!) binding
 237
 238         $wC = \x y -> $wC x y
 239
 240 i.e. a curried constructor that allocates.  This means that we can
 241 treat the worker for a constructor like any other function in the rest
 242 of the compiler.  The point here is that CoreToStg will generate a
 243 StgConApp for the RHS, rather than a call to the worker (which would
 244 give a loop).  As Lennart says: the ice is thin here, but it works.
 245
 246 Hmm.  Should we create bindings for dictionary constructors?  They are
 247 always fully applied, and the bindings are just there to support
 248 partial applications. But it's easier to let them through.
 249
 250
 251 %************************************************************************
 252 %*                                                                      *
 253                 The main code
 254 %*                                                                      *
 255 %************************************************************************
 256
 257 \begin{code}
 258 cpeBind :: TopLevelFlag
 259         -> CorePrepEnv -> CoreBind
 260         -> UniqSM (CorePrepEnv, Floats)
 261 cpeBind top_lvl env (NonRec bndr rhs)
 262   = do { (_, bndr1) <- cloneBndr env bndr
 263        ; let is_strict   = isStrictDmd (idDemandInfo bndr)
 264              is_unlifted = isUnLiftedType (idType bndr)
 265        ; (floats, bndr2, rhs2) <- cpePair top_lvl NonRecursive
 266                                           (is_strict || is_unlifted)
 267                                           env bndr1 rhs
 268        ; let new_float = mkFloat is_strict is_unlifted bndr2 rhs2
 269
 270         -- We want bndr'' in the envt, because it records
 271         -- the evaluated-ness of the binder
 272        ; return (extendCorePrepEnv env bndr bndr2,
 273                  addFloat floats new_float) }
 274
 275 cpeBind top_lvl env (Rec pairs)
 276   = do { let (bndrs,rhss) = unzip pairs
 277        ; (env', bndrs1) <- cloneBndrs env (map fst pairs)
 278        ; stuff <- zipWithM (cpePair top_lvl Recursive False env') bndrs1 rhss
 279
 280        ; let (floats_s, bndrs2, rhss2) = unzip3 stuff
 281              all_pairs = foldrOL add_float (bndrs1 `zip` rhss2)
 282                                            (concatFloats floats_s)
 283        ; return (extendCorePrepEnvList env (bndrs `zip` bndrs2),
 284                  unitFloat (FloatLet (Rec all_pairs))) }
 285   where
 286         -- Flatten all the floats, and the currrent
 287         -- group into a single giant Rec
 288     add_float (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
 289     add_float (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
 290     add_float b                       _    = pprPanic "cpeBind" (ppr b)
 291
 292 ---------------
 293 cpePair :: TopLevelFlag -> RecFlag -> RhsDemand
 294         -> CorePrepEnv -> Id -> CoreExpr
 295         -> UniqSM (Floats, Id, CpeRhs)
 296 -- Used for all bindings
 297 cpePair top_lvl is_rec is_strict_or_unlifted env bndr rhs
 298   = do { (floats1, rhs1) <- cpeRhsE env rhs
 299
 300        -- See if we are allowed to float this stuff out of the RHS
 301        ; (floats2, rhs2) <- float_from_rhs floats1 rhs1
 302
 303        -- Make the arity match up
 304        ; (floats3, rhs')
 305             <- if manifestArity rhs1 <= arity
 306                then return (floats2, cpeEtaExpand arity rhs2)
 307                else WARN(True, text "CorePrep: silly extra arguments:" <+> ppr bndr)
 308                                -- Note [Silly extra arguments]
 309                     (do { v <- newVar (idType bndr)
 310                         ; let float = mkFloat False False v rhs2
 311                         ; return (addFloat floats2 float, cpeEtaExpand arity (Var v)) })
 312
 313                 -- Record if the binder is evaluated
 314        ; let bndr' | exprIsHNF rhs' = bndr `setIdUnfolding` evaldUnfolding
 315                    | otherwise      = bndr
 316
 317        ; return (floats3, bndr', rhs') }
 318   where
 319     arity = idArity bndr        -- We must match this arity
 320
 321     ---------------------
 322     float_from_rhs floats rhs
 323       | isEmptyFloats floats = return (emptyFloats, rhs)
 324       | isTopLevel top_lvl    = float_top    floats rhs
 325       | otherwise             = float_nested floats rhs
 326
 327     ---------------------
 328     float_nested floats rhs
 329       | wantFloatNested is_rec is_strict_or_unlifted floats rhs
 330                   = return (floats, rhs)
 331       | otherwise = dont_float floats rhs
 332
 333     ---------------------
 334     float_top floats rhs        -- Urhgh!  See Note [CafInfo and floating]
 335       | mayHaveCafRefs (idCafInfo bndr)
 336       , allLazyTop floats
 337       = return (floats, rhs)
 338
 339       -- So the top-level binding is marked NoCafRefs
 340       | Just (floats', rhs') <- canFloatFromNoCaf floats rhs
 341       = return (floats', rhs')
 342
 343       | otherwise
 344       = dont_float floats rhs
 345
 346     ---------------------
 347     dont_float floats rhs
 348       -- Non-empty floats, but do not want to float from rhs
 349       -- So wrap the rhs in the floats
 350       -- But: rhs1 might have lambdas, and we can't
 351       --      put them inside a wrapBinds
 352       = do { body <- rhsToBodyNF rhs
 353            ; return (emptyFloats, wrapBinds floats body) }
 354
 355 {- Note [Silly extra arguments]
 356 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 357 Suppose we had this
 358         f{arity=1} = \x\y. e
 359 We *must* match the arity on the Id, so we have to generate
 360         f' = \x\y. e
 361         f  = \x. f' x
 362
 363 It's a bizarre case: why is the arity on the Id wrong?  Reason
 364 (in the days of __inline_me__):
 365         f{arity=0} = __inline_me__ (let v = expensive in \xy. e)
 366 When InlineMe notes go away this won't happen any more.  But
 367 it seems good for CorePrep to be robust.
 368 -}
 369
 370 -- ---------------------------------------------------------------------------
 371 --              CpeRhs: produces a result satisfying CpeRhs
 372 -- ---------------------------------------------------------------------------
 373
 374 cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 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 cpeRhsE _env expr@(Type _) = return (emptyFloats, expr)
 384 cpeRhsE _env expr@(Lit _)  = return (emptyFloats, expr)
 385 cpeRhsE env expr@(Var {})  = cpeApp env expr
 386
 387 cpeRhsE env (Var f `App` _ `App` arg)
 388   | f `hasKey` lazyIdKey          -- Replace (lazy a) by a
 389   = cpeRhsE env arg               -- See Note [lazyId magic] in MkId
 390
 391 cpeRhsE env expr@(App {}) = cpeApp env expr
 392
 393 cpeRhsE env (Let bind expr)
 394   = do { (env', new_binds) <- cpeBind NotTopLevel env bind
 395        ; (floats, body) <- cpeRhsE env' expr
 396        ; return (new_binds `appendFloats` floats, body) }
 397
 398 cpeRhsE env (Note note expr)
 399   | ignoreNote note
 400   = cpeRhsE env expr
 401   | otherwise         -- Just SCCs actually
 402   = do { body <- cpeBodyNF env expr
 403        ; return (emptyFloats, Note note body) }
 404
 405 cpeRhsE env (Cast expr co)
 406    = do { (floats, expr') <- cpeRhsE env expr
 407         ; return (floats, Cast expr' co) }
 408
 409 cpeRhsE env expr@(Lam {})
 410    = do { let (bndrs,body) = collectBinders expr
 411         ; (env', bndrs') <- cloneBndrs env bndrs
 412         ; body' <- cpeBodyNF env' body
 413         ; return (emptyFloats, mkLams bndrs' body') }
 414
 415 cpeRhsE env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
 416   | Just (TickBox {}) <- isTickBoxOp_maybe id
 417   = do { body <- cpeBodyNF env expr
 418        ; return (emptyFloats, Case (Var id) bndr ty [(DEFAULT,[],body)]) }
 419
 420 cpeRhsE env (Case scrut bndr ty alts)
 421   = do { (floats, scrut') <- cpeBody env scrut
 422        ; let bndr1 = bndr `setIdUnfolding` evaldUnfolding
 423             -- Record that the case binder is evaluated in the alternatives
 424        ; (env', bndr2) <- cloneBndr env bndr1
 425        ; alts' <- mapM (sat_alt env') alts
 426        ; return (floats, Case scrut' bndr2 ty alts') }
 427   where
 428     sat_alt env (con, bs, rhs)
 429        = do { (env2, bs') <- cloneBndrs env bs
 430             ; rhs' <- cpeBodyNF env2 rhs
 431             ; return (con, bs', rhs') }
 432
 433 -- ---------------------------------------------------------------------------
 434 --              CpeBody: produces a result satisfying CpeBody
 435 -- ---------------------------------------------------------------------------
 436
 437 cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
 438 cpeBodyNF env expr
 439   = do { (floats, body) <- cpeBody env expr
 440        ; return (wrapBinds floats body) }
 441
 442 --------
 443 cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
 444 cpeBody env expr
 445   = do { (floats1, rhs) <- cpeRhsE env expr
 446        ; (floats2, body) <- rhsToBody rhs
 447        ; return (floats1 `appendFloats` floats2, body) }
 448
 449 --------
 450 rhsToBodyNF :: CpeRhs -> UniqSM CpeBody
 451 rhsToBodyNF rhs = do { (floats,body) <- rhsToBody rhs
 452                      ; return (wrapBinds floats body) }
 453
 454 --------
 455 rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
 456 -- Remove top level lambdas by let-binding
 457
 458 rhsToBody (Note n expr)
 459         -- You can get things like
 460         --      case e of { p -> coerce t (\s -> ...) }
 461   = do { (floats, expr') <- rhsToBody expr
 462        ; return (floats, Note n expr') }
 463
 464 rhsToBody (Cast e co)
 465   = do { (floats, e') <- rhsToBody e
 466        ; return (floats, Cast e' co) }
 467
 468 rhsToBody expr@(Lam {})
 469   | Just no_lam_result <- tryEtaReducePrep bndrs body
 470   = return (emptyFloats, no_lam_result)
 471   | all isTyCoVar bndrs         -- Type lambdas are ok
 472   = return (emptyFloats, expr)
 473   | otherwise                   -- Some value lambdas
 474   = do { fn <- newVar (exprType expr)
 475        ; let rhs   = cpeEtaExpand (exprArity expr) expr
 476              float = FloatLet (NonRec fn rhs)
 477        ; return (unitFloat float, Var fn) }
 478   where
 479     (bndrs,body) = collectBinders expr
 480
 481 rhsToBody expr = return (emptyFloats, expr)
 482
 483
 484
 485 -- ---------------------------------------------------------------------------
 486 --              CpeApp: produces a result satisfying CpeApp
 487 -- ---------------------------------------------------------------------------
 488
 489 cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 490 -- May return a CpeRhs because of saturating primops
 491 cpeApp env expr
 492   = do { (app, (head,depth), _, floats, ss) <- collect_args expr 0
 493        ; MASSERT(null ss)       -- make sure we used all the strictness info
 494
 495         -- Now deal with the function
 496        ; case head of
 497            Var fn_id -> do { sat_app <- maybeSaturate fn_id app depth
 498                            ; return (floats, sat_app) }
 499            _other    -> return (floats, app) }
 500
 501   where
 502     -- Deconstruct and rebuild the application, floating any non-atomic
 503     -- arguments to the outside.  We collect the type of the expression,
 504     -- the head of the application, and the number of actual value arguments,
 505     -- all of which are used to possibly saturate this application if it
 506     -- has a constructor or primop at the head.
 507
 508     collect_args
 509         :: CoreExpr
 510         -> Int                     -- Current app depth
 511         -> UniqSM (CpeApp,         -- The rebuilt expression
 512                    (CoreExpr,Int), -- The head of the application,
 513                                    -- and no. of args it was applied to
 514                    Type,           -- Type of the whole expr
 515                    Floats,         -- Any floats we pulled out
 516                    [Demand])       -- Remaining argument demands
 517
 518     collect_args (App fun arg@(Type arg_ty)) depth
 519       = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
 520            ; return (App fun' arg, hd, applyTy fun_ty arg_ty, floats, ss) }
 521
 522     collect_args (App fun arg) depth
 523       = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun (depth+1)
 524            ; let
 525               (ss1, ss_rest)   = case ss of
 526                                    (ss1:ss_rest) -> (ss1,     ss_rest)
 527                                    []            -> (lazyDmd, [])
 528               (arg_ty, res_ty) = expectJust "cpeBody:collect_args" $
 529                                  splitFunTy_maybe fun_ty
 530
 531            ; (fs, arg') <- cpeArg env (isStrictDmd ss1) arg arg_ty
 532            ; return (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest) }
 533
 534     collect_args (Var v) depth
 535       = do { v1 <- fiddleCCall v
 536            ; let v2 = lookupCorePrepEnv env v1
 537            ; return (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts) }
 538         where
 539           stricts = case idStrictness v of
 540                         StrictSig (DmdType _ demands _)
 541                             | listLengthCmp demands depth /= GT -> demands
 542                                     -- length demands <= depth
 543                             | otherwise                         -> []
 544                 -- If depth < length demands, then we have too few args to
 545                 -- satisfy strictness  info so we have to  ignore all the
 546                 -- strictness info, e.g. + (error "urk")
 547                 -- Here, we can't evaluate the arg strictly, because this
 548                 -- partial application might be seq'd
 549
 550     collect_args (Cast fun co) depth
 551       = do { let (_ty1,ty2) = coercionKind co
 552            ; (fun', hd, _, floats, ss) <- collect_args fun depth
 553            ; return (Cast fun' co, hd, ty2, floats, ss) }
 554
 555     collect_args (Note note fun) depth
 556       | ignoreNote note         -- Drop these notes altogether
 557       = collect_args fun depth  -- They aren't used by the code generator
 558
 559         -- N-variable fun, better let-bind it
 560     collect_args fun depth
 561       = do { (fun_floats, fun') <- cpeArg env True fun ty
 562                           -- The True says that it's sure to be evaluated,
 563                           -- so we'll end up case-binding it
 564            ; return (fun', (fun', depth), ty, fun_floats, []) }
 565         where
 566           ty = exprType fun
 567
 568 -- ---------------------------------------------------------------------------
 569 --      CpeArg: produces a result satisfying CpeArg
 570 -- ---------------------------------------------------------------------------
 571
 572 -- This is where we arrange that a non-trivial argument is let-bound
 573 cpeArg :: CorePrepEnv -> RhsDemand -> CoreArg -> Type
 574        -> UniqSM (Floats, CpeTriv)
 575 cpeArg env is_strict arg arg_ty
 576   | cpe_ExprIsTrivial arg   -- Do not eta expand etc a trivial argument
 577   = cpeBody env arg         -- Must still do substitution though
 578   | otherwise
 579   = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
 580        ; (floats2, arg2) <- if want_float floats1 arg1
 581                             then return (floats1, arg1)
 582                             else do { body1 <- rhsToBodyNF arg1
 583                                     ; return (emptyFloats, wrapBinds floats1 body1) }
 584                 -- Else case: arg1 might have lambdas, and we can't
 585                 --            put them inside a wrapBinds
 586
 587        ; v <- newVar arg_ty
 588        ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
 589              arg_float = mkFloat is_strict is_unlifted v arg3
 590        ; return (addFloat floats2 arg_float, Var v) }
 591   where
 592     is_unlifted = isUnLiftedType arg_ty
 593     want_float = wantFloatNested NonRecursive (is_strict || is_unlifted)
 594 \end{code}
 595
 596 Note [Floating unlifted arguments]
 597 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 598 Consider    C (let v* = expensive in v)
 599
 600 where the "*" indicates "will be demanded".  Usually v will have been
 601 inlined by now, but let's suppose it hasn't (see Trac #2756).  Then we
 602 do *not* want to get
 603
 604      let v* = expensive in C v
 605
 606 because that has different strictness.  Hence the use of 'allLazy'.
 607 (NB: the let v* turns into a FloatCase, in mkLocalNonRec.)
 608
 609
 610 ------------------------------------------------------------------------------
 611 -- Building the saturated syntax
 612 -- ---------------------------------------------------------------------------
 613
 614 maybeSaturate deals with saturating primops and constructors
 615 The type is the type of the entire application
 616
 617 \begin{code}
 618 maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
 619 maybeSaturate fn expr n_args
 620   | Just DataToTagOp <- isPrimOpId_maybe fn     -- DataToTag must have an evaluated arg
 621                                                 -- A gruesome special case
 622   = saturateDataToTag sat_expr
 623
 624   | hasNoBinding fn        -- There's no binding
 625   = return sat_expr
 626
 627   | otherwise
 628   = return expr
 629   where
 630     fn_arity     = idArity fn
 631     excess_arity = fn_arity - n_args
 632     sat_expr     = cpeEtaExpand excess_arity expr
 633
 634 -------------
 635 saturateDataToTag :: CpeApp -> UniqSM CpeApp
 636 -- Horrid: ensure that the arg of data2TagOp is evaluated
 637 --   (data2tag x) -->  (case x of y -> data2tag y)
 638 -- (yuk yuk) take into account the lambdas we've now introduced
 639 saturateDataToTag sat_expr
 640   = do { let (eta_bndrs, eta_body) = collectBinders sat_expr
 641        ; eta_body' <- eval_data2tag_arg eta_body
 642        ; return (mkLams eta_bndrs eta_body') }
 643   where
 644     eval_data2tag_arg :: CpeApp -> UniqSM CpeBody
 645     eval_data2tag_arg app@(fun `App` arg)
 646         | exprIsHNF arg         -- Includes nullary constructors
 647         = return app            -- The arg is evaluated
 648         | otherwise                     -- Arg not evaluated, so evaluate it
 649         = do { arg_id <- newVar (exprType arg)
 650              ; let arg_id1 = setIdUnfolding arg_id evaldUnfolding
 651              ; return (Case arg arg_id1 (exprType app)
 652                             [(DEFAULT, [], fun `App` Var arg_id1)]) }
 653
 654     eval_data2tag_arg (Note note app)   -- Scc notes can appear
 655         = do { app' <- eval_data2tag_arg app
 656              ; return (Note note app') }
 657
 658     eval_data2tag_arg other     -- Should not happen
 659         = pprPanic "eval_data2tag" (ppr other)
 660 \end{code}
 661
 662
 663
 664
 665 %************************************************************************
 666 %*                                                                      *
 667                 Simple CoreSyn operations
 668 %*                                                                      *
 669 %************************************************************************
 670
 671 \begin{code}
 672         -- We don't ignore SCCs, since they require some code generation
 673 ignoreNote :: Note -> Bool
 674 -- Tells which notes to drop altogether; they are ignored by code generation
 675 -- Do not ignore SCCs!
 676 -- It's important that we do drop InlineMe notes; for example
 677 --    unzip = __inline_me__ (/\ab. foldr (..) (..))
 678 -- Here unzip gets arity 1 so we'll eta-expand it. But we don't
 679 -- want to get this:
 680 --     unzip = /\ab \xs. (__inline_me__ ...) a b xs
 681 ignoreNote (CoreNote _) = True
 682 ignoreNote _other       = False
 683
 684
 685 cpe_ExprIsTrivial :: CoreExpr -> Bool
 686 -- Version that doesn't consider an scc annotation to be trivial.
 687 cpe_ExprIsTrivial (Var _)                  = True
 688 cpe_ExprIsTrivial (Type _)                 = True
 689 cpe_ExprIsTrivial (Lit _)                  = True
 690 cpe_ExprIsTrivial (App e arg)              = isTypeArg arg && cpe_ExprIsTrivial e
 691 cpe_ExprIsTrivial (Note (SCC _) _)         = False
 692 cpe_ExprIsTrivial (Note _ e)               = cpe_ExprIsTrivial e
 693 cpe_ExprIsTrivial (Cast e _)               = cpe_ExprIsTrivial e
 694 cpe_ExprIsTrivial (Lam b body) | isTyCoVar b = cpe_ExprIsTrivial body
 695 cpe_ExprIsTrivial _                        = False
 696 \end{code}
 697
 698 -- -----------------------------------------------------------------------------
 699 --      Eta reduction
 700 -- -----------------------------------------------------------------------------
 701
 702 Note [Eta expansion]
 703 ~~~~~~~~~~~~~~~~~~~~~
 704 Eta expand to match the arity claimed by the binder Remember,
 705 CorePrep must not change arity
 706
 707 Eta expansion might not have happened already, because it is done by
 708 the simplifier only when there at least one lambda already.
 709
 710 NB1:we could refrain when the RHS is trivial (which can happen
 711     for exported things).  This would reduce the amount of code
 712     generated (a little) and make things a little words for
 713     code compiled without -O.  The case in point is data constructor
 714     wrappers.
 715
 716 NB2: we have to be careful that the result of etaExpand doesn't
 717    invalidate any of the assumptions that CorePrep is attempting
 718    to establish.  One possible cause is eta expanding inside of
 719    an SCC note - we're now careful in etaExpand to make sure the
 720    SCC is pushed inside any new lambdas that are generated.
 721
 722 Note [Eta expansion and the CorePrep invariants]
 723 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 724 It turns out to be much much easier to do eta expansion
 725 *after* the main CorePrep stuff.  But that places constraints
 726 on the eta expander: given a CpeRhs, it must return a CpeRhs.
 727
 728 For example here is what we do not want:
 729                 f = /\a -> g (h 3)      -- h has arity 2
 730 After ANFing we get
 731                 f = /\a -> let s = h 3 in g s
 732 and now we do NOT want eta expansion to give
 733                 f = /\a -> \ y -> (let s = h 3 in g s) y
 734
 735 Instead CoreArity.etaExpand gives
 736                 f = /\a -> \y -> let s = h 3 in g s y
 737
 738 \begin{code}
 739 cpeEtaExpand :: Arity -> CpeRhs -> CpeRhs
 740 cpeEtaExpand arity expr
 741   | arity == 0 = expr
 742   | otherwise  = etaExpand arity expr
 743 \end{code}
 744
 745 -- -----------------------------------------------------------------------------
 746 --      Eta reduction
 747 -- -----------------------------------------------------------------------------
 748
 749 Why try eta reduction?  Hasn't the simplifier already done eta?
 750 But the simplifier only eta reduces if that leaves something
 751 trivial (like f, or f Int).  But for deLam it would be enough to
 752 get to a partial application:
 753         case x of { p -> \xs. map f xs }
 754     ==> case x of { p -> map f }
 755
 756 \begin{code}
 757 tryEtaReducePrep :: [CoreBndr] -> CoreExpr -> Maybe CoreExpr
 758 tryEtaReducePrep bndrs expr@(App _ _)
 759   | ok_to_eta_reduce f &&
 760     n_remaining >= 0 &&
 761     and (zipWith ok bndrs last_args) &&
 762     not (any (`elemVarSet` fvs_remaining) bndrs)
 763   = Just remaining_expr
 764   where
 765     (f, args) = collectArgs expr
 766     remaining_expr = mkApps f remaining_args
 767     fvs_remaining = exprFreeVars remaining_expr
 768     (remaining_args, last_args) = splitAt n_remaining args
 769     n_remaining = length args - length bndrs
 770
 771     ok bndr (Var arg) = bndr == arg
 772     ok _    _         = False
 773
 774           -- we can't eta reduce something which must be saturated.
 775     ok_to_eta_reduce (Var f) = not (hasNoBinding f)
 776     ok_to_eta_reduce _       = False --safe. ToDo: generalise
 777
 778 tryEtaReducePrep bndrs (Let bind@(NonRec _ r) body)
 779   | not (any (`elemVarSet` fvs) bndrs)
 780   = case tryEtaReducePrep bndrs body of
 781         Just e -> Just (Let bind e)
 782         Nothing -> Nothing
 783   where
 784     fvs = exprFreeVars r
 785
 786 tryEtaReducePrep _ _ = Nothing
 787 \end{code}
 788
 789
 790 -- -----------------------------------------------------------------------------
 791 -- Demands
 792 -- -----------------------------------------------------------------------------
 793
 794 \begin{code}
 795 type RhsDemand = Bool  -- True => used strictly; hence not top-level, non-recursive
 796 \end{code}
 797
 798 %************************************************************************
 799 %*                                                                      *
 800                 Floats
 801 %*                                                                      *
 802 %************************************************************************
 803
 804 \begin{code}
 805 data FloatingBind
 806   = FloatLet CoreBind    -- Rhs of bindings are CpeRhss
 807                          -- They are always of lifted type;
 808                          -- unlifted ones are done with FloatCase
 809
 810  | FloatCase
 811       Id CpeBody
 812       Bool              -- The bool indicates "ok-for-speculation"
 813
 814 data Floats = Floats OkToSpec (OrdList FloatingBind)
 815
 816 instance Outputable FloatingBind where
 817   ppr (FloatLet b) = ppr b
 818   ppr (FloatCase b r ok) = brackets (ppr ok) <+> ppr b <+> equals <+> ppr r
 819
 820 instance Outputable Floats where
 821   ppr (Floats flag fs) = ptext (sLit "Floats") <> brackets (ppr flag) <+>
 822                          braces (vcat (map ppr (fromOL fs)))
 823
 824 instance Outputable OkToSpec where
 825   ppr OkToSpec    = ptext (sLit "OkToSpec")
 826   ppr IfUnboxedOk = ptext (sLit "IfUnboxedOk")
 827   ppr NotOkToSpec = ptext (sLit "NotOkToSpec")
 828
 829 -- Can we float these binds out of the rhs of a let?  We cache this decision
 830 -- to avoid having to recompute it in a non-linear way when there are
 831 -- deeply nested lets.
 832 data OkToSpec
 833    = OkToSpec           -- Lazy bindings of lifted type
 834    | IfUnboxedOk        -- A mixture of lazy lifted bindings and n
 835                         -- ok-to-speculate unlifted bindings
 836    | NotOkToSpec        -- Some not-ok-to-speculate unlifted bindings
 837
 838 mkFloat :: Bool -> Bool -> Id -> CpeRhs -> FloatingBind
 839 mkFloat is_strict is_unlifted bndr rhs
 840   | use_case  = FloatCase bndr rhs (exprOkForSpeculation rhs)
 841   | otherwise = FloatLet (NonRec bndr rhs)
 842   where
 843     use_case = is_unlifted || is_strict && not (exprIsHNF rhs)
 844                 -- Don't make a case for a value binding,
 845                 -- even if it's strict.  Otherwise we get
 846                 --      case (\x -> e) of ...!
 847
 848 emptyFloats :: Floats
 849 emptyFloats = Floats OkToSpec nilOL
 850
 851 isEmptyFloats :: Floats -> Bool
 852 isEmptyFloats (Floats _ bs) = isNilOL bs
 853
 854 wrapBinds :: Floats -> CpeBody -> CpeBody
 855 wrapBinds (Floats _ binds) body
 856   = foldrOL mk_bind body binds
 857   where
 858     mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
 859     mk_bind (FloatLet bind)        body = Let bind body
 860
 861 addFloat :: Floats -> FloatingBind -> Floats
 862 addFloat (Floats ok_to_spec floats) new_float
 863   = Floats (combine ok_to_spec (check new_float)) (floats `snocOL` new_float)
 864   where
 865     check (FloatLet _) = OkToSpec
 866     check (FloatCase _ _ ok_for_spec)
 867         | ok_for_spec  =  IfUnboxedOk
 868         | otherwise    =  NotOkToSpec
 869         -- The ok-for-speculation flag says that it's safe to
 870         -- float this Case out of a let, and thereby do it more eagerly
 871         -- We need the top-level flag because it's never ok to float
 872         -- an unboxed binding to the top level
 873
 874 unitFloat :: FloatingBind -> Floats
 875 unitFloat = addFloat emptyFloats
 876
 877 appendFloats :: Floats -> Floats -> Floats
 878 appendFloats (Floats spec1 floats1) (Floats spec2 floats2)
 879   = Floats (combine spec1 spec2) (floats1 `appOL` floats2)
 880
 881 concatFloats :: [Floats] -> OrdList FloatingBind
 882 concatFloats = foldr (\ (Floats _ bs1) bs2 -> appOL bs1 bs2) nilOL
 883
 884 combine :: OkToSpec -> OkToSpec -> OkToSpec
 885 combine NotOkToSpec _ = NotOkToSpec
 886 combine _ NotOkToSpec = NotOkToSpec
 887 combine IfUnboxedOk _ = IfUnboxedOk
 888 combine _ IfUnboxedOk = IfUnboxedOk
 889 combine _ _           = OkToSpec
 890
 891 deFloatTop :: Floats -> [CoreBind]
 892 -- For top level only; we don't expect any FloatCases
 893 deFloatTop (Floats _ floats)
 894   = foldrOL get [] floats
 895   where
 896     get (FloatLet b) bs = b:bs
 897     get b            _  = pprPanic "corePrepPgm" (ppr b)
 898
 899 -------------------------------------------
 900 canFloatFromNoCaf ::  Floats -> CpeRhs -> Maybe (Floats, CpeRhs)
 901        -- Note [CafInfo and floating]
 902 canFloatFromNoCaf (Floats ok_to_spec fs) rhs
 903   | OkToSpec <- ok_to_spec           -- Worth trying
 904   , Just (subst, fs') <- go (emptySubst, nilOL) (fromOL fs)
 905   = Just (Floats OkToSpec fs', subst_expr subst rhs)
 906   | otherwise
 907   = Nothing
 908   where
 909     subst_expr = substExpr (text "CorePrep")
 910
 911     go :: (Subst, OrdList FloatingBind) -> [FloatingBind]
 912        -> Maybe (Subst, OrdList FloatingBind)
 913
 914     go (subst, fbs_out) [] = Just (subst, fbs_out)
 915
 916     go (subst, fbs_out) (FloatLet (NonRec b r) : fbs_in)
 917       | rhs_ok r
 918       = go (subst', fbs_out `snocOL` new_fb) fbs_in
 919       where
 920         (subst', b') = set_nocaf_bndr subst b
 921         new_fb = FloatLet (NonRec b' (subst_expr subst r))
 922
 923     go (subst, fbs_out) (FloatLet (Rec prs) : fbs_in)
 924       | all rhs_ok rs
 925       = go (subst', fbs_out `snocOL` new_fb) fbs_in
 926       where
 927         (bs,rs) = unzip prs
 928         (subst', bs') = mapAccumL set_nocaf_bndr subst bs
 929         rs' = map (subst_expr subst') rs
 930         new_fb = FloatLet (Rec (bs' `zip` rs'))
 931
 932     go _ _ = Nothing      -- Encountered a caffy binding
 933
 934     ------------
 935     set_nocaf_bndr subst bndr
 936       = (extendIdSubst subst bndr (Var bndr'), bndr')
 937       where
 938         bndr' = bndr `setIdCafInfo` NoCafRefs
 939
 940     ------------
 941     rhs_ok :: CoreExpr -> Bool
 942     -- We can only float to top level from a NoCaf thing if
 943     -- the new binding is static. However it can't mention
 944     -- any non-static things or it would *already* be Caffy
 945     rhs_ok = rhsIsStatic (\_ -> False)
 946
 947 wantFloatNested :: RecFlag -> Bool -> Floats -> CpeRhs -> Bool
 948 wantFloatNested is_rec strict_or_unlifted floats rhs
 949   =  isEmptyFloats floats
 950   || strict_or_unlifted
 951   || (allLazyNested is_rec floats && exprIsHNF rhs)
 952         -- Why the test for allLazyNested?
 953         --      v = f (x `divInt#` y)
 954         -- we don't want to float the case, even if f has arity 2,
 955         -- because floating the case would make it evaluated too early
 956
 957 allLazyTop :: Floats -> Bool
 958 allLazyTop (Floats OkToSpec _) = True
 959 allLazyTop _                   = False
 960
 961 allLazyNested :: RecFlag -> Floats -> Bool
 962 allLazyNested _      (Floats OkToSpec    _) = True
 963 allLazyNested _      (Floats NotOkToSpec _) = False
 964 allLazyNested is_rec (Floats IfUnboxedOk _) = isNonRec is_rec
 965 \end{code}
 966
 967
 968 %************************************************************************
 969 %*                                                                      *
 970                 Cloning
 971 %*                                                                      *
 972 %************************************************************************
 973
 974 \begin{code}
 975 -- ---------------------------------------------------------------------------
 976 --                      The environment
 977 -- ---------------------------------------------------------------------------
 978
 979 data CorePrepEnv = CPE (IdEnv Id)       -- Clone local Ids
 980
 981 emptyCorePrepEnv :: CorePrepEnv
 982 emptyCorePrepEnv = CPE emptyVarEnv
 983
 984 extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
 985 extendCorePrepEnv (CPE env) id id' = CPE (extendVarEnv env id id')
 986
 987 extendCorePrepEnvList :: CorePrepEnv -> [(Id,Id)] -> CorePrepEnv
 988 extendCorePrepEnvList (CPE env) prs = CPE (extendVarEnvList env prs)
 989
 990 lookupCorePrepEnv :: CorePrepEnv -> Id -> Id
 991 lookupCorePrepEnv (CPE env) id
 992   = case lookupVarEnv env id of
 993         Nothing  -> id
 994         Just id' -> id'
 995
 996 ------------------------------------------------------------------------------
 997 -- Cloning binders
 998 -- ---------------------------------------------------------------------------
 999
1000 cloneBndrs :: CorePrepEnv -> [Var] -> UniqSM (CorePrepEnv, [Var])
1001 cloneBndrs env bs = mapAccumLM cloneBndr env bs
1002
1003 cloneBndr  :: CorePrepEnv -> Var -> UniqSM (CorePrepEnv, Var)
1004 cloneBndr env bndr
1005   | isLocalId bndr
1006   = do bndr' <- setVarUnique bndr <$> getUniqueM
1007        return (extendCorePrepEnv env bndr bndr', bndr')
1008
1009   | otherwise   -- Top level things, which we don't want
1010                 -- to clone, have become GlobalIds by now
1011                 -- And we don't clone tyvars
1012   = return (env, bndr)
1013
1014
1015 ------------------------------------------------------------------------------
1016 -- Cloning ccall Ids; each must have a unique name,
1017 -- to give the code generator a handle to hang it on
1018 -- ---------------------------------------------------------------------------
1019
1020 fiddleCCall :: Id -> UniqSM Id
1021 fiddleCCall id
1022   | isFCallId id = (id `setVarUnique`) <$> getUniqueM
1023   | otherwise    = return id
1024
1025 ------------------------------------------------------------------------------
1026 -- Generating new binders
1027 -- ---------------------------------------------------------------------------
1028
1029 newVar :: Type -> UniqSM Id
1030 newVar ty
1031  = seqType ty `seq` do
1032      uniq <- getUniqueM
1033      return (mkSysLocal (fsLit "sat") uniq ty)
1034 \end{code}