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 happense 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        ; (floats2, rhs2)
 301             <- if manifestArity rhs1 <= arity
 302                then return (floats1, cpeEtaExpand arity rhs1)
 303                else WARN(True, text "CorePrep: silly extra arguments:" <+> ppr bndr)
 304                                -- Note [Silly extra arguments]
 305                     (do { v <- newVar (idType bndr)
 306                         ; let float = mkFloat False False v rhs1
 307                         ; return (addFloat floats1 float, cpeEtaExpand arity (Var v)) })
 308
 309        ; (floats3, rhs') <- float_from_rhs floats2 rhs2
 310
 311                 -- Record if the binder is evaluated
 312        ; let bndr' | exprIsHNF rhs' = bndr `setIdUnfolding` evaldUnfolding
 313                    | otherwise      = bndr
 314
 315        ; return (floats3, bndr', rhs') }
 316   where
 317     arity = idArity bndr        -- We must match this arity
 318
 319     ---------------------
 320     float_from_rhs floats2 rhs2
 321       | isEmptyFloats floats2 = return (emptyFloats, rhs2)
 322       | isTopLevel top_lvl    = float_top    floats2 rhs2
 323       | otherwise             = float_nested floats2 rhs2
 324
 325     ---------------------
 326     float_nested floats2 rhs2
 327       | wantFloatNested is_rec is_strict_or_unlifted floats2 rhs2
 328                   = return (floats2, rhs2)
 329       | otherwise = dont_float floats2 rhs2
 330
 331     ---------------------
 332     float_top floats2 rhs2      -- Urhgh!  See Note [CafInfo and floating]
 333       | mayHaveCafRefs (idCafInfo bndr)
 334       = if allLazyTop floats2
 335         then return (floats2, rhs2)
 336         else dont_float floats2 rhs2
 337
 338       | otherwise
 339       = case canFloatFromNoCaf floats2 rhs2 of
 340           Just (floats2', rhs2') -> return (floats2', rhs2')
 341           Nothing -> pprPanic "cpePair" (ppr bndr $$ ppr rhs2 $$ ppr floats2)
 342
 343     ---------------------
 344     dont_float floats2 rhs2
 345       -- Non-empty floats, but do not want to float from rhs
 346       -- So wrap the rhs in the floats
 347       -- But: rhs1 might have lambdas, and we can't
 348       --      put them inside a wrapBinds
 349       = do { body2 <- rhsToBodyNF rhs2
 350            ; return (emptyFloats, wrapBinds floats2 body2) }
 351
 352 {- Note [Silly extra arguments]
 353 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 354 Suppose we had this
 355         f{arity=1} = \x\y. e
 356 We *must* match the arity on the Id, so we have to generate
 357         f' = \x\y. e
 358         f  = \x. f' x
 359
 360 It's a bizarre case: why is the arity on the Id wrong?  Reason
 361 (in the days of __inline_me__):
 362         f{arity=0} = __inline_me__ (let v = expensive in \xy. e)
 363 When InlineMe notes go away this won't happen any more.  But
 364 it seems good for CorePrep to be robust.
 365 -}
 366
 367 -- ---------------------------------------------------------------------------
 368 --              CpeRhs: produces a result satisfying CpeRhs
 369 -- ---------------------------------------------------------------------------
 370
 371 cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 372 -- If
 373 --      e  ===>  (bs, e')
 374 -- then
 375 --      e = let bs in e'        (semantically, that is!)
 376 --
 377 -- For example
 378 --      f (g x)   ===>   ([v = g x], f v)
 379
 380 cpeRhsE _env expr@(Type _) = return (emptyFloats, expr)
 381 cpeRhsE _env expr@(Lit _)  = return (emptyFloats, expr)
 382 cpeRhsE env expr@(Var {})  = cpeApp env expr
 383
 384 cpeRhsE env (Var f `App` _ `App` arg)
 385   | f `hasKey` lazyIdKey          -- Replace (lazy a) by a
 386   = cpeRhsE env arg               -- See Note [lazyId magic] in MkId
 387
 388 cpeRhsE env expr@(App {}) = cpeApp env expr
 389
 390 cpeRhsE env (Let bind expr)
 391   = do { (env', new_binds) <- cpeBind NotTopLevel env bind
 392        ; (floats, body) <- cpeRhsE env' expr
 393        ; return (new_binds `appendFloats` floats, body) }
 394
 395 cpeRhsE env (Note note expr)
 396   | ignoreNote note
 397   = cpeRhsE env expr
 398   | otherwise         -- Just SCCs actually
 399   = do { body <- cpeBodyNF env expr
 400        ; return (emptyFloats, Note note body) }
 401
 402 cpeRhsE env (Cast expr co)
 403    = do { (floats, expr') <- cpeRhsE env expr
 404         ; return (floats, Cast expr' co) }
 405
 406 cpeRhsE env expr@(Lam {})
 407    = do { let (bndrs,body) = collectBinders expr
 408         ; (env', bndrs') <- cloneBndrs env bndrs
 409         ; body' <- cpeBodyNF env' body
 410         ; return (emptyFloats, mkLams bndrs' body') }
 411
 412 cpeRhsE env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
 413   | Just (TickBox {}) <- isTickBoxOp_maybe id
 414   = do { body <- cpeBodyNF env expr
 415        ; return (emptyFloats, Case (Var id) bndr ty [(DEFAULT,[],body)]) }
 416
 417 cpeRhsE env (Case scrut bndr ty alts)
 418   = do { (floats, scrut') <- cpeBody env scrut
 419        ; let bndr1 = bndr `setIdUnfolding` evaldUnfolding
 420             -- Record that the case binder is evaluated in the alternatives
 421        ; (env', bndr2) <- cloneBndr env bndr1
 422        ; alts' <- mapM (sat_alt env') alts
 423        ; return (floats, Case scrut' bndr2 ty alts') }
 424   where
 425     sat_alt env (con, bs, rhs)
 426        = do { (env2, bs') <- cloneBndrs env bs
 427             ; rhs' <- cpeBodyNF env2 rhs
 428             ; return (con, bs', rhs') }
 429
 430 -- ---------------------------------------------------------------------------
 431 --              CpeBody: produces a result satisfying CpeBody
 432 -- ---------------------------------------------------------------------------
 433
 434 cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
 435 cpeBodyNF env expr
 436   = do { (floats, body) <- cpeBody env expr
 437        ; return (wrapBinds floats body) }
 438
 439 --------
 440 cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
 441 cpeBody env expr
 442   = do { (floats1, rhs) <- cpeRhsE env expr
 443        ; (floats2, body) <- rhsToBody rhs
 444        ; return (floats1 `appendFloats` floats2, body) }
 445
 446 --------
 447 rhsToBodyNF :: CpeRhs -> UniqSM CpeBody
 448 rhsToBodyNF rhs = do { (floats,body) <- rhsToBody rhs
 449                      ; return (wrapBinds floats body) }
 450
 451 --------
 452 rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
 453 -- Remove top level lambdas by let-binding
 454
 455 rhsToBody (Note n expr)
 456         -- You can get things like
 457         --      case e of { p -> coerce t (\s -> ...) }
 458   = do { (floats, expr') <- rhsToBody expr
 459        ; return (floats, Note n expr') }
 460
 461 rhsToBody (Cast e co)
 462   = do { (floats, e') <- rhsToBody e
 463        ; return (floats, Cast e' co) }
 464
 465 rhsToBody expr@(Lam {})
 466   | Just no_lam_result <- tryEtaReducePrep bndrs body
 467   = return (emptyFloats, no_lam_result)
 468   | all isTyCoVar bndrs         -- Type lambdas are ok
 469   = return (emptyFloats, expr)
 470   | otherwise                   -- Some value lambdas
 471   = do { fn <- newVar (exprType expr)
 472        ; let rhs   = cpeEtaExpand (exprArity expr) expr
 473              float = FloatLet (NonRec fn rhs)
 474        ; return (unitFloat float, Var fn) }
 475   where
 476     (bndrs,body) = collectBinders expr
 477
 478 rhsToBody expr = return (emptyFloats, expr)
 479
 480
 481
 482 -- ---------------------------------------------------------------------------
 483 --              CpeApp: produces a result satisfying CpeApp
 484 -- ---------------------------------------------------------------------------
 485
 486 cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 487 -- May return a CpeRhs because of saturating primops
 488 cpeApp env expr
 489   = do { (app, (head,depth), _, floats, ss) <- collect_args expr 0
 490        ; MASSERT(null ss)       -- make sure we used all the strictness info
 491
 492         -- Now deal with the function
 493        ; case head of
 494            Var fn_id -> do { sat_app <- maybeSaturate fn_id app depth
 495                            ; return (floats, sat_app) }
 496            _other    -> return (floats, app) }
 497
 498   where
 499     -- Deconstruct and rebuild the application, floating any non-atomic
 500     -- arguments to the outside.  We collect the type of the expression,
 501     -- the head of the application, and the number of actual value arguments,
 502     -- all of which are used to possibly saturate this application if it
 503     -- has a constructor or primop at the head.
 504
 505     collect_args
 506         :: CoreExpr
 507         -> Int                     -- Current app depth
 508         -> UniqSM (CpeApp,         -- The rebuilt expression
 509                    (CoreExpr,Int), -- The head of the application,
 510                                    -- and no. of args it was applied to
 511                    Type,           -- Type of the whole expr
 512                    Floats,         -- Any floats we pulled out
 513                    [Demand])       -- Remaining argument demands
 514
 515     collect_args (App fun arg@(Type arg_ty)) depth
 516       = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
 517            ; return (App fun' arg, hd, applyTy fun_ty arg_ty, floats, ss) }
 518
 519     collect_args (App fun arg) depth
 520       = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun (depth+1)
 521            ; let
 522               (ss1, ss_rest)   = case ss of
 523                                    (ss1:ss_rest) -> (ss1,     ss_rest)
 524                                    []            -> (lazyDmd, [])
 525               (arg_ty, res_ty) = expectJust "cpeBody:collect_args" $
 526                                  splitFunTy_maybe fun_ty
 527
 528            ; (fs, arg') <- cpeArg env (isStrictDmd ss1) arg arg_ty
 529            ; return (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest) }
 530
 531     collect_args (Var v) depth
 532       = do { v1 <- fiddleCCall v
 533            ; let v2 = lookupCorePrepEnv env v1
 534            ; return (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts) }
 535         where
 536           stricts = case idStrictness v of
 537                         StrictSig (DmdType _ demands _)
 538                             | listLengthCmp demands depth /= GT -> demands
 539                                     -- length demands <= depth
 540                             | otherwise                         -> []
 541                 -- If depth < length demands, then we have too few args to
 542                 -- satisfy strictness  info so we have to  ignore all the
 543                 -- strictness info, e.g. + (error "urk")
 544                 -- Here, we can't evaluate the arg strictly, because this
 545                 -- partial application might be seq'd
 546
 547     collect_args (Cast fun co) depth
 548       = do { let (_ty1,ty2) = coercionKind co
 549            ; (fun', hd, _, floats, ss) <- collect_args fun depth
 550            ; return (Cast fun' co, hd, ty2, floats, ss) }
 551
 552     collect_args (Note note fun) depth
 553       | ignoreNote note         -- Drop these notes altogether
 554       = collect_args fun depth  -- They aren't used by the code generator
 555
 556         -- N-variable fun, better let-bind it
 557     collect_args fun depth
 558       = do { (fun_floats, fun') <- cpeArg env True fun ty
 559                           -- The True says that it's sure to be evaluated,
 560                           -- so we'll end up case-binding it
 561            ; return (fun', (fun', depth), ty, fun_floats, []) }
 562         where
 563           ty = exprType fun
 564
 565 -- ---------------------------------------------------------------------------
 566 --      CpeArg: produces a result satisfying CpeArg
 567 -- ---------------------------------------------------------------------------
 568
 569 -- This is where we arrange that a non-trivial argument is let-bound
 570 cpeArg :: CorePrepEnv -> RhsDemand -> CoreArg -> Type
 571        -> UniqSM (Floats, CpeTriv)
 572 cpeArg env is_strict arg arg_ty
 573   | cpe_ExprIsTrivial arg   -- Do not eta expand etc a trivial argument
 574   = cpeBody env arg         -- Must still do substitution though
 575   | otherwise
 576   = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
 577        ; (floats2, arg2) <- if want_float floats1 arg1
 578                             then return (floats1, arg1)
 579                             else do { body1 <- rhsToBodyNF arg1
 580                                     ; return (emptyFloats, wrapBinds floats1 body1) }
 581                 -- Else case: arg1 might have lambdas, and we can't
 582                 --            put them inside a wrapBinds
 583
 584        ; v <- newVar arg_ty
 585        ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
 586              arg_float = mkFloat is_strict is_unlifted v arg3
 587        ; return (addFloat floats2 arg_float, Var v) }
 588   where
 589     is_unlifted = isUnLiftedType arg_ty
 590     want_float = wantFloatNested NonRecursive (is_strict || is_unlifted)
 591 \end{code}
 592
 593 Note [Floating unlifted arguments]
 594 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 595 Consider    C (let v* = expensive in v)
 596
 597 where the "*" indicates "will be demanded".  Usually v will have been
 598 inlined by now, but let's suppose it hasn't (see Trac #2756).  Then we
 599 do *not* want to get
 600
 601      let v* = expensive in C v
 602
 603 because that has different strictness.  Hence the use of 'allLazy'.
 604 (NB: the let v* turns into a FloatCase, in mkLocalNonRec.)
 605
 606
 607 ------------------------------------------------------------------------------
 608 -- Building the saturated syntax
 609 -- ---------------------------------------------------------------------------
 610
 611 maybeSaturate deals with saturating primops and constructors
 612 The type is the type of the entire application
 613
 614 \begin{code}
 615 maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
 616 maybeSaturate fn expr n_args
 617   | Just DataToTagOp <- isPrimOpId_maybe fn     -- DataToTag must have an evaluated arg
 618                                                 -- A gruesome special case
 619   = saturateDataToTag sat_expr
 620
 621   | hasNoBinding fn        -- There's no binding
 622   = return sat_expr
 623
 624   | otherwise
 625   = return expr
 626   where
 627     fn_arity     = idArity fn
 628     excess_arity = fn_arity - n_args
 629     sat_expr     = cpeEtaExpand excess_arity expr
 630
 631 -------------
 632 saturateDataToTag :: CpeApp -> UniqSM CpeApp
 633 -- Horrid: ensure that the arg of data2TagOp is evaluated
 634 --   (data2tag x) -->  (case x of y -> data2tag y)
 635 -- (yuk yuk) take into account the lambdas we've now introduced
 636 saturateDataToTag sat_expr
 637   = do { let (eta_bndrs, eta_body) = collectBinders sat_expr
 638        ; eta_body' <- eval_data2tag_arg eta_body
 639        ; return (mkLams eta_bndrs eta_body') }
 640   where
 641     eval_data2tag_arg :: CpeApp -> UniqSM CpeBody
 642     eval_data2tag_arg app@(fun `App` arg)
 643         | exprIsHNF arg         -- Includes nullary constructors
 644         = return app            -- The arg is evaluated
 645         | otherwise                     -- Arg not evaluated, so evaluate it
 646         = do { arg_id <- newVar (exprType arg)
 647              ; let arg_id1 = setIdUnfolding arg_id evaldUnfolding
 648              ; return (Case arg arg_id1 (exprType app)
 649                             [(DEFAULT, [], fun `App` Var arg_id1)]) }
 650
 651     eval_data2tag_arg (Note note app)   -- Scc notes can appear
 652         = do { app' <- eval_data2tag_arg app
 653              ; return (Note note app') }
 654
 655     eval_data2tag_arg other     -- Should not happen
 656         = pprPanic "eval_data2tag" (ppr other)
 657 \end{code}
 658
 659
 660
 661
 662 %************************************************************************
 663 %*                                                                      *
 664                 Simple CoreSyn operations
 665 %*                                                                      *
 666 %************************************************************************
 667
 668 \begin{code}
 669         -- We don't ignore SCCs, since they require some code generation
 670 ignoreNote :: Note -> Bool
 671 -- Tells which notes to drop altogether; they are ignored by code generation
 672 -- Do not ignore SCCs!
 673 -- It's important that we do drop InlineMe notes; for example
 674 --    unzip = __inline_me__ (/\ab. foldr (..) (..))
 675 -- Here unzip gets arity 1 so we'll eta-expand it. But we don't
 676 -- want to get this:
 677 --     unzip = /\ab \xs. (__inline_me__ ...) a b xs
 678 ignoreNote (CoreNote _) = True
 679 ignoreNote _other       = False
 680
 681
 682 cpe_ExprIsTrivial :: CoreExpr -> Bool
 683 -- Version that doesn't consider an scc annotation to be trivial.
 684 cpe_ExprIsTrivial (Var _)                  = True
 685 cpe_ExprIsTrivial (Type _)                 = True
 686 cpe_ExprIsTrivial (Lit _)                  = True
 687 cpe_ExprIsTrivial (App e arg)              = isTypeArg arg && cpe_ExprIsTrivial e
 688 cpe_ExprIsTrivial (Note (SCC _) _)         = False
 689 cpe_ExprIsTrivial (Note _ e)               = cpe_ExprIsTrivial e
 690 cpe_ExprIsTrivial (Cast e _)               = cpe_ExprIsTrivial e
 691 cpe_ExprIsTrivial (Lam b body) | isTyCoVar b = cpe_ExprIsTrivial body
 692 cpe_ExprIsTrivial _                        = False
 693 \end{code}
 694
 695 -- -----------------------------------------------------------------------------
 696 --      Eta reduction
 697 -- -----------------------------------------------------------------------------
 698
 699 Note [Eta expansion]
 700 ~~~~~~~~~~~~~~~~~~~~~
 701 Eta expand to match the arity claimed by the binder Remember,
 702 CorePrep must not change arity
 703
 704 Eta expansion might not have happened already, because it is done by
 705 the simplifier only when there at least one lambda already.
 706
 707 NB1:we could refrain when the RHS is trivial (which can happen
 708     for exported things).  This would reduce the amount of code
 709     generated (a little) and make things a little words for
 710     code compiled without -O.  The case in point is data constructor
 711     wrappers.
 712
 713 NB2: we have to be careful that the result of etaExpand doesn't
 714    invalidate any of the assumptions that CorePrep is attempting
 715    to establish.  One possible cause is eta expanding inside of
 716    an SCC note - we're now careful in etaExpand to make sure the
 717    SCC is pushed inside any new lambdas that are generated.
 718
 719 Note [Eta expansion and the CorePrep invariants]
 720 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 721 It turns out to be much much easier to do eta expansion
 722 *after* the main CorePrep stuff.  But that places constraints
 723 on the eta expander: given a CpeRhs, it must return a CpeRhs.
 724
 725 For example here is what we do not want:
 726                 f = /\a -> g (h 3)      -- h has arity 2
 727 After ANFing we get
 728                 f = /\a -> let s = h 3 in g s
 729 and now we do NOT want eta expansion to give
 730                 f = /\a -> \ y -> (let s = h 3 in g s) y
 731
 732 Instead CoreArity.etaExpand gives
 733                 f = /\a -> \y -> let s = h 3 in g s y
 734
 735 \begin{code}
 736 cpeEtaExpand :: Arity -> CpeRhs -> CpeRhs
 737 cpeEtaExpand arity expr
 738   | arity == 0 = expr
 739   | otherwise  = etaExpand arity expr
 740 \end{code}
 741
 742 -- -----------------------------------------------------------------------------
 743 --      Eta reduction
 744 -- -----------------------------------------------------------------------------
 745
 746 Why try eta reduction?  Hasn't the simplifier already done eta?
 747 But the simplifier only eta reduces if that leaves something
 748 trivial (like f, or f Int).  But for deLam it would be enough to
 749 get to a partial application:
 750         case x of { p -> \xs. map f xs }
 751     ==> case x of { p -> map f }
 752
 753 \begin{code}
 754 tryEtaReducePrep :: [CoreBndr] -> CoreExpr -> Maybe CoreExpr
 755 tryEtaReducePrep bndrs expr@(App _ _)
 756   | ok_to_eta_reduce f &&
 757     n_remaining >= 0 &&
 758     and (zipWith ok bndrs last_args) &&
 759     not (any (`elemVarSet` fvs_remaining) bndrs)
 760   = Just remaining_expr
 761   where
 762     (f, args) = collectArgs expr
 763     remaining_expr = mkApps f remaining_args
 764     fvs_remaining = exprFreeVars remaining_expr
 765     (remaining_args, last_args) = splitAt n_remaining args
 766     n_remaining = length args - length bndrs
 767
 768     ok bndr (Var arg) = bndr == arg
 769     ok _    _         = False
 770
 771           -- we can't eta reduce something which must be saturated.
 772     ok_to_eta_reduce (Var f) = not (hasNoBinding f)
 773     ok_to_eta_reduce _       = False --safe. ToDo: generalise
 774
 775 tryEtaReducePrep bndrs (Let bind@(NonRec _ r) body)
 776   | not (any (`elemVarSet` fvs) bndrs)
 777   = case tryEtaReducePrep bndrs body of
 778         Just e -> Just (Let bind e)
 779         Nothing -> Nothing
 780   where
 781     fvs = exprFreeVars r
 782
 783 tryEtaReducePrep _ _ = Nothing
 784 \end{code}
 785
 786
 787 -- -----------------------------------------------------------------------------
 788 -- Demands
 789 -- -----------------------------------------------------------------------------
 790
 791 \begin{code}
 792 type RhsDemand = Bool  -- True => used strictly; hence not top-level, non-recursive
 793 \end{code}
 794
 795 %************************************************************************
 796 %*                                                                      *
 797                 Floats
 798 %*                                                                      *
 799 %************************************************************************
 800
 801 \begin{code}
 802 data FloatingBind
 803   = FloatLet CoreBind    -- Rhs of bindings are CpeRhss
 804                          -- They are always of lifted type;
 805                          -- unlifted ones are done with FloatCase
 806
 807  | FloatCase
 808       Id CpeBody
 809       Bool              -- The bool indicates "ok-for-speculation"
 810
 811 data Floats = Floats OkToSpec (OrdList FloatingBind)
 812
 813 instance Outputable FloatingBind where
 814   ppr (FloatLet b) = ppr b
 815   ppr (FloatCase b r ok) = brackets (ppr ok) <+> ppr b <+> equals <+> ppr r
 816
 817 instance Outputable Floats where
 818   ppr (Floats flag fs) = ptext (sLit "Floats") <> brackets (ppr flag) <+>
 819                          braces (vcat (map ppr (fromOL fs)))
 820
 821 instance Outputable OkToSpec where
 822   ppr OkToSpec    = ptext (sLit "OkToSpec")
 823   ppr IfUnboxedOk = ptext (sLit "IfUnboxedOk")
 824   ppr NotOkToSpec = ptext (sLit "NotOkToSpec")
 825
 826 -- Can we float these binds out of the rhs of a let?  We cache this decision
 827 -- to avoid having to recompute it in a non-linear way when there are
 828 -- deeply nested lets.
 829 data OkToSpec
 830    = OkToSpec           -- Lazy bindings of lifted type
 831    | IfUnboxedOk        -- A mixture of lazy lifted bindings and n
 832                         -- ok-to-speculate unlifted bindings
 833    | NotOkToSpec        -- Some not-ok-to-speculate unlifted bindings
 834
 835 mkFloat :: Bool -> Bool -> Id -> CpeRhs -> FloatingBind
 836 mkFloat is_strict is_unlifted bndr rhs
 837   | use_case  = FloatCase bndr rhs (exprOkForSpeculation rhs)
 838   | otherwise = FloatLet (NonRec bndr rhs)
 839   where
 840     use_case = is_unlifted || is_strict && not (exprIsHNF rhs)
 841                 -- Don't make a case for a value binding,
 842                 -- even if it's strict.  Otherwise we get
 843                 --      case (\x -> e) of ...!
 844
 845 emptyFloats :: Floats
 846 emptyFloats = Floats OkToSpec nilOL
 847
 848 isEmptyFloats :: Floats -> Bool
 849 isEmptyFloats (Floats _ bs) = isNilOL bs
 850
 851 wrapBinds :: Floats -> CpeBody -> CpeBody
 852 wrapBinds (Floats _ binds) body
 853   = foldrOL mk_bind body binds
 854   where
 855     mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
 856     mk_bind (FloatLet bind)        body = Let bind body
 857
 858 addFloat :: Floats -> FloatingBind -> Floats
 859 addFloat (Floats ok_to_spec floats) new_float
 860   = Floats (combine ok_to_spec (check new_float)) (floats `snocOL` new_float)
 861   where
 862     check (FloatLet _) = OkToSpec
 863     check (FloatCase _ _ ok_for_spec)
 864         | ok_for_spec  =  IfUnboxedOk
 865         | otherwise    =  NotOkToSpec
 866         -- The ok-for-speculation flag says that it's safe to
 867         -- float this Case out of a let, and thereby do it more eagerly
 868         -- We need the top-level flag because it's never ok to float
 869         -- an unboxed binding to the top level
 870
 871 unitFloat :: FloatingBind -> Floats
 872 unitFloat = addFloat emptyFloats
 873
 874 appendFloats :: Floats -> Floats -> Floats
 875 appendFloats (Floats spec1 floats1) (Floats spec2 floats2)
 876   = Floats (combine spec1 spec2) (floats1 `appOL` floats2)
 877
 878 concatFloats :: [Floats] -> OrdList FloatingBind
 879 concatFloats = foldr (\ (Floats _ bs1) bs2 -> appOL bs1 bs2) nilOL
 880
 881 combine :: OkToSpec -> OkToSpec -> OkToSpec
 882 combine NotOkToSpec _ = NotOkToSpec
 883 combine _ NotOkToSpec = NotOkToSpec
 884 combine IfUnboxedOk _ = IfUnboxedOk
 885 combine _ IfUnboxedOk = IfUnboxedOk
 886 combine _ _           = OkToSpec
 887
 888 deFloatTop :: Floats -> [CoreBind]
 889 -- For top level only; we don't expect any FloatCases
 890 deFloatTop (Floats _ floats)
 891   = foldrOL get [] floats
 892   where
 893     get (FloatLet b) bs = b:bs
 894     get b            _  = pprPanic "corePrepPgm" (ppr b)
 895
 896 -------------------------------------------
 897 canFloatFromNoCaf ::  Floats -> CpeRhs -> Maybe (Floats, CpeRhs)
 898        -- Note [CafInfo and floating]
 899 canFloatFromNoCaf (Floats ok_to_spec fs) rhs
 900   | OkToSpec <- ok_to_spec
 901   = Just (Floats OkToSpec (toOL fs'), subst_expr subst rhs)
 902   | otherwise
 903   = Nothing
 904   where
 905     (subst, fs') = mapAccumL set_nocaf emptySubst (fromOL fs)
 906
 907     subst_expr = substExpr (text "CorePrep")
 908
 909     set_nocaf _ (FloatCase {})
 910       = panic "canFloatFromNoCaf"
 911
 912     set_nocaf subst (FloatLet (NonRec b r))
 913       = (subst', FloatLet (NonRec b' (subst_expr subst r)))
 914       where
 915         (subst', b') = set_nocaf_bndr subst b
 916
 917     set_nocaf subst (FloatLet (Rec prs))
 918       = (subst', FloatLet (Rec (bs' `zip` rs')))
 919       where
 920         (bs,rs) = unzip prs
 921         (subst', bs') = mapAccumL set_nocaf_bndr subst bs
 922         rs' = map (subst_expr subst') rs
 923
 924     set_nocaf_bndr subst bndr
 925       = (extendIdSubst subst bndr (Var bndr'), bndr')
 926       where
 927         bndr' = bndr `setIdCafInfo` NoCafRefs
 928
 929 wantFloatNested :: RecFlag -> Bool -> Floats -> CpeRhs -> Bool
 930 wantFloatNested is_rec strict_or_unlifted floats rhs
 931   =  isEmptyFloats floats
 932   || strict_or_unlifted
 933   || (allLazyNested is_rec floats && exprIsHNF rhs)
 934         -- Why the test for allLazyNested?
 935         --      v = f (x `divInt#` y)
 936         -- we don't want to float the case, even if f has arity 2,
 937         -- because floating the case would make it evaluated too early
 938
 939 allLazyTop :: Floats -> Bool
 940 allLazyTop (Floats OkToSpec _) = True
 941 allLazyTop _                   = False
 942
 943 allLazyNested :: RecFlag -> Floats -> Bool
 944 allLazyNested _      (Floats OkToSpec    _) = True
 945 allLazyNested _      (Floats NotOkToSpec _) = False
 946 allLazyNested is_rec (Floats IfUnboxedOk _) = isNonRec is_rec
 947 \end{code}
 948
 949
 950 %************************************************************************
 951 %*                                                                      *
 952                 Cloning
 953 %*                                                                      *
 954 %************************************************************************
 955
 956 \begin{code}
 957 -- ---------------------------------------------------------------------------
 958 --                      The environment
 959 -- ---------------------------------------------------------------------------
 960
 961 data CorePrepEnv = CPE (IdEnv Id)       -- Clone local Ids
 962
 963 emptyCorePrepEnv :: CorePrepEnv
 964 emptyCorePrepEnv = CPE emptyVarEnv
 965
 966 extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
 967 extendCorePrepEnv (CPE env) id id' = CPE (extendVarEnv env id id')
 968
 969 extendCorePrepEnvList :: CorePrepEnv -> [(Id,Id)] -> CorePrepEnv
 970 extendCorePrepEnvList (CPE env) prs = CPE (extendVarEnvList env prs)
 971
 972 lookupCorePrepEnv :: CorePrepEnv -> Id -> Id
 973 lookupCorePrepEnv (CPE env) id
 974   = case lookupVarEnv env id of
 975         Nothing  -> id
 976         Just id' -> id'
 977
 978 ------------------------------------------------------------------------------
 979 -- Cloning binders
 980 -- ---------------------------------------------------------------------------
 981
 982 cloneBndrs :: CorePrepEnv -> [Var] -> UniqSM (CorePrepEnv, [Var])
 983 cloneBndrs env bs = mapAccumLM cloneBndr env bs
 984
 985 cloneBndr  :: CorePrepEnv -> Var -> UniqSM (CorePrepEnv, Var)
 986 cloneBndr env bndr
 987   | isLocalId bndr
 988   = do bndr' <- setVarUnique bndr <$> getUniqueM
 989        return (extendCorePrepEnv env bndr bndr', bndr')
 990
 991   | otherwise   -- Top level things, which we don't want
 992                 -- to clone, have become GlobalIds by now
 993                 -- And we don't clone tyvars
 994   = return (env, bndr)
 995
 996
 997 ------------------------------------------------------------------------------
 998 -- Cloning ccall Ids; each must have a unique name,
 999 -- to give the code generator a handle to hang it on
1000 -- ---------------------------------------------------------------------------
1001
1002 fiddleCCall :: Id -> UniqSM Id
1003 fiddleCCall id
1004   | isFCallId id = (id `setVarUnique`) <$> getUniqueM
1005   | otherwise    = return id
1006
1007 ------------------------------------------------------------------------------
1008 -- Generating new binders
1009 -- ---------------------------------------------------------------------------
1010
1011 newVar :: Type -> UniqSM Id
1012 newVar ty
1013  = seqType ty `seq` do
1014      uniq <- getUniqueM
1015      return (mkSysLocal (fsLit "sat") uniq ty)
1016 \end{code}