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