Make dumpIfSet_dyn_or use dumpSDoc
[ghc-hetmet.git] / compiler / coreSyn / CorePrep.lhs
index 049960f..0405716 100644 (file)
@@ -5,27 +5,24 @@
 Core pass to saturate constructors and PrimOps
 
 \begin{code}
-{-# OPTIONS_GHC -w #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and fix
--- any warnings in the module. See
---     http://hackage.haskell.org/trac/ghc/wiki/WorkingConventions#Warnings
--- for details
-
 module CorePrep (
       corePrepPgm, corePrepExpr
   ) where
 
 #include "HsVersions.h"
 
-import CoreUtils hiding (exprIsTrivial)
+import PrelNames       ( lazyIdKey, hasKey )
+import CoreUtils
+import CoreArity
 import CoreFVs
-import CoreLint
+import CoreMonad       ( endPass, CoreToDo(..) )
 import CoreSyn
+import CoreSubst
+import OccurAnal        ( occurAnalyseExpr )
 import Type
 import Coercion
 import TyCon
-import NewDemand
+import Demand
 import Var
 import VarSet
 import VarEnv
@@ -40,7 +37,12 @@ import OrdList
 import ErrUtils
 import DynFlags
 import Util
+import Pair
 import Outputable
+import MonadUtils
+import FastString
+import Data.List       ( mapAccumL )
+import Control.Monad
 \end{code}
 
 -- ---------------------------------------------------------------------------
@@ -66,8 +68,9 @@ The goal of this pass is to prepare for code generation.
     [I'm experimenting with leaving 'ok-for-speculation' 
      rhss in let-form right up to this point.]
 
-4.  Ensure that lambdas only occur as the RHS of a binding
+4.  Ensure that *value* lambdas only occur as the RHS of a binding
     (The code generator can't deal with anything else.)
+    Type lambdas are ok, however, because the code gen discards them.
 
 5.  [Not any more; nuked Jun 2002] Do the seq/par munging.
 
@@ -76,9 +79,9 @@ The goal of this pass is to prepare for code generation.
     weaker guarantee of no clashes which the simplifier provides.
     And that is what the code generator needs.
 
-    We don't clone TyVars. The code gen doesn't need that, 
+    We don't clone TyVars or CoVars. The code gen doesn't need that, 
     and doing so would be tiresome because then we'd need
-    to substitute in types.
+    to substitute in types and coercions.
 
 
 7.  Give each dynamic CCall occurrence a fresh unique; this is
@@ -87,53 +90,152 @@ The goal of this pass is to prepare for code generation.
 8.  Inject bindings for the "implicit" Ids:
        * Constructor wrappers
        * Constructor workers
-       * Record selectors
     We want curried definitions for all of these in case they
     aren't inlined by some caller.
        
+9.  Replace (lazy e) by e.  See Note [lazyId magic] in MkId.lhs
+
 This is all done modulo type applications and abstractions, so that
 when type erasure is done for conversion to STG, we don't end up with
 any trivial or useless bindings.
 
   
+Invariants
+~~~~~~~~~~
+Here is the syntax of the Core produced by CorePrep:
 
--- -----------------------------------------------------------------------------
--- Top level stuff
--- -----------------------------------------------------------------------------
+    Trivial expressions 
+       triv ::= lit |  var  
+              | triv ty  |  /\a. triv 
+              | truv co  |  /\c. triv  |  triv |> co
+
+    Applications
+       app ::= lit  |  var  |  app triv  |  app ty  | app co | app |> co
+
+    Expressions
+       body ::= app  
+              | let(rec) x = rhs in body     -- Boxed only
+              | case body of pat -> body
+             | /\a. body | /\c. body 
+              | body |> co
+
+    Right hand sides (only place where value lambdas can occur)
+       rhs ::= /\a.rhs  |  \x.rhs  |  body
+
+We define a synonym for each of these non-terminals.  Functions
+with the corresponding name produce a result in that syntax.
+
+\begin{code}
+type CpeTriv = CoreExpr           -- Non-terminal 'triv'
+type CpeApp  = CoreExpr           -- Non-terminal 'app'
+type CpeBody = CoreExpr           -- Non-terminal 'body'
+type CpeRhs  = CoreExpr           -- Non-terminal 'rhs'
+\end{code}
+
+%************************************************************************
+%*                                                                     *
+               Top level stuff
+%*                                                                     *
+%************************************************************************
 
 \begin{code}
 corePrepPgm :: DynFlags -> [CoreBind] -> [TyCon] -> IO [CoreBind]
-corePrepPgm dflags binds data_tycons
-  = do showPass dflags "CorePrep"
-       us <- mkSplitUniqSupply 's'
-
-       let implicit_binds = mkDataConWorkers data_tycons
-               -- NB: we must feed mkImplicitBinds through corePrep too
-               -- so that they are suitably cloned and eta-expanded
-
-           binds_out = initUs_ us (
-                         corePrepTopBinds binds        `thenUs` \ floats1 ->
-                         corePrepTopBinds implicit_binds       `thenUs` \ floats2 ->
-                         returnUs (deFloatTop (floats1 `appendFloats` floats2))
-                       )
-           
-        endPass dflags "CorePrep" Opt_D_dump_prep binds_out
-       return binds_out
+corePrepPgm dflags binds data_tycons = do
+    showPass dflags "CorePrep"
+    us <- mkSplitUniqSupply 's'
+
+    let implicit_binds = mkDataConWorkers data_tycons
+            -- NB: we must feed mkImplicitBinds through corePrep too
+            -- so that they are suitably cloned and eta-expanded
+
+        binds_out = initUs_ us $ do
+                      floats1 <- corePrepTopBinds binds
+                      floats2 <- corePrepTopBinds implicit_binds
+                      return (deFloatTop (floats1 `appendFloats` floats2))
+
+    endPass dflags CorePrep binds_out []
+    return binds_out
 
 corePrepExpr :: DynFlags -> CoreExpr -> IO CoreExpr
-corePrepExpr dflags expr
-  = do showPass dflags "CorePrep"
-       us <- mkSplitUniqSupply 's'
-       let new_expr = initUs_ us (corePrepAnExpr emptyCorePrepEnv expr)
-       dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep" 
-                    (ppr new_expr)
-       return new_expr
-\end{code}
+corePrepExpr dflags expr = do
+    showPass dflags "CorePrep"
+    us <- mkSplitUniqSupply 's'
+    let new_expr = initUs_ us (cpeBodyNF emptyCorePrepEnv expr)
+    dumpIfSet_dyn dflags Opt_D_dump_prep "CorePrep" (ppr new_expr)
+    return new_expr
 
--- -----------------------------------------------------------------------------
--- Implicit bindings
--- -----------------------------------------------------------------------------
+corePrepTopBinds :: [CoreBind] -> UniqSM Floats
+-- Note [Floating out of top level bindings]
+corePrepTopBinds binds 
+  = go emptyCorePrepEnv binds
+  where
+    go _   []             = return emptyFloats
+    go env (bind : binds) = do (env', bind') <- cpeBind TopLevel env bind
+                               binds' <- go env' binds
+                               return (bind' `appendFloats` binds')
 
+mkDataConWorkers :: [TyCon] -> [CoreBind]
+-- See Note [Data constructor workers]
+mkDataConWorkers data_tycons
+  = [ NonRec id (Var id)       -- The ice is thin here, but it works
+    | tycon <- data_tycons,    -- CorePrep will eta-expand it
+      data_con <- tyConDataCons tycon,
+      let id = dataConWorkId data_con ]
+\end{code}
+
+Note [Floating out of top level bindings]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+NB: we do need to float out of top-level bindings
+Consider       x = length [True,False]
+We want to get
+               s1 = False : []
+               s2 = True  : s1
+               x  = length s2
+
+We return a *list* of bindings, because we may start with
+       x* = f (g y)
+where x is demanded, in which case we want to finish with
+       a = g y
+       x* = f a
+And then x will actually end up case-bound
+
+Note [CafInfo and floating]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+What happens when we try to float bindings to the top level?  At this
+point all the CafInfo is supposed to be correct, and we must make certain
+that is true of the new top-level bindings.  There are two cases
+to consider
+
+a) The top-level binding is marked asCafRefs.  In that case we are
+   basically fine.  The floated bindings had better all be lazy lets,
+   so they can float to top level, but they'll all have HasCafRefs
+   (the default) which is safe.
+
+b) The top-level binding is marked NoCafRefs.  This really happens
+   Example.  CoreTidy produces
+      $fApplicativeSTM [NoCafRefs] = D:Alternative retry# ...blah...
+   Now CorePrep has to eta-expand to
+      $fApplicativeSTM = let sat = \xy. retry x y
+                         in D:Alternative sat ...blah...
+   So what we *want* is
+      sat [NoCafRefs] = \xy. retry x y
+      $fApplicativeSTM [NoCafRefs] = D:Alternative sat ...blah...
+   
+   So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
+   *and* substutite the modified 'sat' into the old RHS.  
+
+   It should be the case that 'sat' is itself [NoCafRefs] (a value, no
+   cafs) else the original top-level binding would not itself have been
+   marked [NoCafRefs].  The DEBUG check in CoreToStg for
+   consistentCafInfo will find this.
+
+This is all very gruesome and horrible. It would be better to figure
+out CafInfo later, after CorePrep.  We'll do that in due course. 
+Meanwhile this horrible hack works.
+
+
+Note [Data constructor workers]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Create any necessary "implicit" bindings for data con workers.  We
 create the rather strange (non-recursive!) binding
 
@@ -149,214 +251,190 @@ Hmm.  Should we create bindings for dictionary constructors?  They are
 always fully applied, and the bindings are just there to support
 partial applications. But it's easier to let them through.
 
-\begin{code}
-mkDataConWorkers data_tycons
-  = [ NonRec id (Var id)       -- The ice is thin here, but it works
-    | tycon <- data_tycons,    -- CorePrep will eta-expand it
-      data_con <- tyConDataCons tycon,
-      let id = dataConWorkId data_con ]
-\end{code}
-       
 
-\begin{code}
--- ---------------------------------------------------------------------------
--- Dealing with bindings
--- ---------------------------------------------------------------------------
+Note [Dead code in CorePrep]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Imagine that we got an input program like this:
 
-data FloatingBind = FloatLet CoreBind
-                 | FloatCase Id CoreExpr Bool
-                       -- The bool indicates "ok-for-speculation"
+  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
+  f x = (g True (Just x) + g () (Just x), g)
+    where
+      g :: Show a => a -> Maybe Int -> Int
+      g _ Nothing = x
+      g y (Just z) = if z > 100 then g y (Just (z + length (show y))) else g y unknown
 
-data Floats = Floats OkToSpec (OrdList FloatingBind)
+After specialisation and SpecConstr, we would get something like this:
 
--- Can we float these binds out of the rhs of a let?  We cache this decision
--- to avoid having to recompute it in a non-linear way when there are
--- deeply nested lets.
-data OkToSpec
-   = NotOkToSpec       -- definitely not
-   | OkToSpec          -- yes
-   | IfUnboxedOk       -- only if floating an unboxed binding is ok
+  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
+  f x = (g$Bool_True_Just x + g$Unit_Unit_Just x, g)
+    where
+      {-# RULES g $dBool = g$Bool 
+                g $dUnit = g$Unit #-}
+      g = ...
+      {-# RULES forall x. g$Bool True (Just x) = g$Bool_True_Just x #-}
+      g$Bool = ...
+      {-# RULES forall x. g$Unit () (Just x) = g$Unit_Unit_Just x #-}
+      g$Unit = ...
+      g$Bool_True_Just = ...
+      g$Unit_Unit_Just = ...
 
-emptyFloats :: Floats
-emptyFloats = Floats OkToSpec nilOL
+Note that the g$Bool and g$Unit functions are actually dead code: they are only kept
+alive by the occurrence analyser because they are referred to by the rules of g,
+which is being kept alive by the fact that it is used (unspecialised) in the returned pair.
 
-addFloat :: Floats -> FloatingBind -> Floats
-addFloat (Floats ok_to_spec floats) new_float
-  = Floats (combine ok_to_spec (check new_float)) (floats `snocOL` new_float)
-  where
-    check (FloatLet _)               = OkToSpec
-    check (FloatCase _ _ ok_for_spec) 
-       | ok_for_spec  =  IfUnboxedOk
-       | otherwise    =  NotOkToSpec
-       -- The ok-for-speculation flag says that it's safe to
-       -- float this Case out of a let, and thereby do it more eagerly
-       -- We need the top-level flag because it's never ok to float
-       -- an unboxed binding to the top level
+However, at the CorePrep stage there is no way that the rules for g will ever fire,
+and it really seems like a shame to produce an output program that goes to the trouble
+of allocating a closure for the unreachable g$Bool and g$Unit functions.
 
-unitFloat :: FloatingBind -> Floats
-unitFloat = addFloat emptyFloats
+The way we fix this is to:
+ * In cloneBndr, drop all unfoldings/rules
+ * In deFloatTop, run the occurrence analyser on each top-level RHS to drop
+   the dead local bindings
 
-appendFloats :: Floats -> Floats -> Floats
-appendFloats (Floats spec1 floats1) (Floats spec2 floats2)
-  = Floats (combine spec1 spec2) (floats1 `appOL` floats2)
+The reason we don't just OccAnal the whole output of CorePrep is that the tidier
+ensures that all top-level binders are GlobalIds, so they don't show up in the free
+variables any longer. So if you run the occurrence analyser on the output of CoreTidy
+(or later) you e.g. turn this program:
 
-concatFloats :: [Floats] -> Floats
-concatFloats = foldr appendFloats emptyFloats
+  Rec {
+  f = ... f ...
+  }
 
-combine NotOkToSpec _ = NotOkToSpec
-combine _ NotOkToSpec = NotOkToSpec
-combine IfUnboxedOk _ = IfUnboxedOk
-combine _ IfUnboxedOk = IfUnboxedOk
-combine _ _           = OkToSpec
-    
-instance Outputable FloatingBind where
-  ppr (FloatLet bind)        = text "FloatLet" <+> ppr bind
-  ppr (FloatCase b rhs spec) = text "FloatCase" <+> ppr b <+> ppr spec <+> equals <+> ppr rhs
+Into this one:
 
-deFloatTop :: Floats -> [CoreBind]
--- For top level only; we don't expect any FloatCases
-deFloatTop (Floats _ floats)
-  = foldrOL get [] floats
-  where
-    get (FloatLet b) bs = b:bs
-    get b           bs = pprPanic "corePrepPgm" (ppr b)
+  f = ... f ...
 
-allLazy :: TopLevelFlag -> RecFlag -> Floats -> Bool
-allLazy top_lvl is_rec (Floats ok_to_spec _)
-  = case ok_to_spec of
-       OkToSpec    -> True
-       NotOkToSpec -> False
-       IfUnboxedOk -> isNotTopLevel top_lvl && isNonRec is_rec
+(Since f is not considered to be free in its own RHS.)
 
--- ---------------------------------------------------------------------------
---                     Bindings
--- ---------------------------------------------------------------------------
 
-corePrepTopBinds :: [CoreBind] -> UniqSM Floats
-corePrepTopBinds binds 
-  = go emptyCorePrepEnv binds
-  where
-    go env []            = returnUs emptyFloats
-    go env (bind : binds) = corePrepTopBind env bind   `thenUs` \ (env', bind') ->
-                           go env' binds               `thenUs` \ binds' ->
-                           returnUs (bind' `appendFloats` binds')
-
--- NB: we do need to float out of top-level bindings
--- Consider    x = length [True,False]
--- We want to get
---             s1 = False : []
---             s2 = True  : s1
---             x  = length s2
-
--- We return a *list* of bindings, because we may start with
---     x* = f (g y)
--- where x is demanded, in which case we want to finish with
---     a = g y
---     x* = f a
--- And then x will actually end up case-bound
---
--- What happens to the CafInfo on the floated bindings?  By
--- default, all the CafInfos will be set to MayHaveCafRefs,
--- which is safe.
---
--- This might be pessimistic, because eg. s1 & s2
--- might not refer to any CAFs and the GC will end up doing
--- more traversal than is necessary, but it's still better
--- than not floating the bindings at all, because then
--- the GC would have to traverse the structure in the heap
--- instead.  Given this, we decided not to try to get
--- the CafInfo on the floated bindings correct, because
--- it looks difficult.
-
---------------------------------
-corePrepTopBind :: CorePrepEnv -> CoreBind -> UniqSM (CorePrepEnv, Floats)
-corePrepTopBind env (NonRec bndr rhs) 
-  = cloneBndr env bndr                                 `thenUs` \ (env', bndr') ->
-    corePrepRhs TopLevel NonRecursive env (bndr, rhs)  `thenUs` \ (floats, rhs') -> 
-    returnUs (env', addFloat floats (FloatLet (NonRec bndr' rhs')))
-
-corePrepTopBind env (Rec pairs) = corePrepRecPairs TopLevel env pairs
-
---------------------------------
-corePrepBind ::  CorePrepEnv -> CoreBind -> UniqSM (CorePrepEnv, Floats)
-       -- This one is used for *local* bindings
-corePrepBind env (NonRec bndr rhs)
-  = etaExpandRhs bndr rhs                              `thenUs` \ rhs1 ->
-    corePrepExprFloat env rhs1                         `thenUs` \ (floats, rhs2) ->
-    cloneBndr env bndr                                 `thenUs` \ (_, bndr') ->
-    mkLocalNonRec bndr' (bdrDem bndr) floats rhs2      `thenUs` \ (floats', bndr'') ->
-       -- We want bndr'' in the envt, because it records
-       -- the evaluated-ness of the binder
-    returnUs (extendCorePrepEnv env bndr bndr'', floats')
-
-corePrepBind env (Rec pairs) = corePrepRecPairs NotTopLevel env pairs
-
---------------------------------
-corePrepRecPairs :: TopLevelFlag -> CorePrepEnv
-                -> [(Id,CoreExpr)]     -- Recursive bindings
-                -> UniqSM (CorePrepEnv, Floats)
--- Used for all recursive bindings, top level and otherwise
-corePrepRecPairs lvl env pairs
-  = cloneBndrs env (map fst pairs)                             `thenUs` \ (env', bndrs') ->
-    mapAndUnzipUs (corePrepRhs lvl Recursive env') pairs       `thenUs` \ (floats_s, rhss') ->
-    returnUs (env', unitFloat (FloatLet (Rec (flatten (concatFloats floats_s) bndrs' rhss'))))
+%************************************************************************
+%*                                                                     *
+               The main code
+%*                                                                     *
+%************************************************************************
+
+\begin{code}
+cpeBind :: TopLevelFlag
+       -> CorePrepEnv -> CoreBind
+       -> UniqSM (CorePrepEnv, Floats)
+cpeBind top_lvl env (NonRec bndr rhs)
+  = do { (_, bndr1) <- cloneBndr env bndr
+       ; let is_strict   = isStrictDmd (idDemandInfo bndr)
+             is_unlifted = isUnLiftedType (idType bndr)
+       ; (floats, bndr2, rhs2) <- cpePair top_lvl NonRecursive 
+                                                 (is_strict || is_unlifted) 
+                                         env bndr1 rhs
+       ; let new_float = mkFloat is_strict is_unlifted bndr2 rhs2
+
+        -- We want bndr'' in the envt, because it records
+        -- the evaluated-ness of the binder
+       ; return (extendCorePrepEnv env bndr bndr2, 
+                        addFloat floats new_float) }
+
+cpeBind top_lvl env (Rec pairs)
+  = do { let (bndrs,rhss) = unzip pairs
+       ; (env', bndrs1) <- cloneBndrs env (map fst pairs)
+       ; stuff <- zipWithM (cpePair top_lvl Recursive False env') bndrs1 rhss
+
+       ; let (floats_s, bndrs2, rhss2) = unzip3 stuff
+             all_pairs = foldrOL add_float (bndrs2 `zip` rhss2)
+                                          (concatFloats floats_s)
+       ; return (extendCorePrepEnvList env (bndrs `zip` bndrs2),
+                        unitFloat (FloatLet (Rec all_pairs))) }
   where
        -- Flatten all the floats, and the currrent
        -- group into a single giant Rec
-    flatten (Floats _ floats) bndrs rhss = foldrOL get (bndrs `zip` rhss) floats
-
-    get (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
-    get (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
-    get b                      prs2 = pprPanic "corePrepRecPairs" (ppr b)
-
---------------------------------
-corePrepRhs :: TopLevelFlag -> RecFlag
-           -> CorePrepEnv -> (Id, CoreExpr)
-           -> UniqSM (Floats, CoreExpr)
--- Used for top-level bindings, and local recursive bindings
-corePrepRhs top_lvl is_rec env (bndr, rhs)
-  = etaExpandRhs bndr rhs      `thenUs` \ rhs' ->
-    corePrepExprFloat env rhs' `thenUs` \ floats_w_rhs ->
-    floatRhs top_lvl is_rec bndr floats_w_rhs
-
-
--- ---------------------------------------------------------------------------
--- Making arguments atomic (function args & constructor args)
--- ---------------------------------------------------------------------------
-
--- This is where we arrange that a non-trivial argument is let-bound
-corePrepArg :: CorePrepEnv -> CoreArg -> RhsDemand
-          -> UniqSM (Floats, CoreArg)
-corePrepArg env arg dem
-  = corePrepExprFloat env arg          `thenUs` \ (floats, arg') ->
-    if exprIsTrivial arg'
-    then returnUs (floats, arg')
-    else newVar (exprType arg')                        `thenUs` \ v ->
-        mkLocalNonRec v dem floats arg'        `thenUs` \ (floats', v') -> 
-        returnUs (floats', Var v')
-
--- version that doesn't consider an scc annotation to be trivial.
-exprIsTrivial (Var v)                 = True
-exprIsTrivial (Type _)                = True
-exprIsTrivial (Lit lit)               = True
-exprIsTrivial (App e arg)             = isTypeArg arg && exprIsTrivial e
-exprIsTrivial (Note (SCC _) e)                = False
-exprIsTrivial (Note _ e)              = exprIsTrivial e
-exprIsTrivial (Cast e co)              = exprIsTrivial e
-exprIsTrivial (Lam b body) | isTyVar b = exprIsTrivial body
-exprIsTrivial other                   = False
+    add_float (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
+    add_float (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
+    add_float b                       _    = pprPanic "cpeBind" (ppr b)
+
+---------------
+cpePair :: TopLevelFlag -> RecFlag -> RhsDemand
+       -> CorePrepEnv -> Id -> CoreExpr
+       -> UniqSM (Floats, Id, CpeRhs)
+-- Used for all bindings
+cpePair top_lvl is_rec is_strict_or_unlifted env bndr rhs
+  = do { (floats1, rhs1) <- cpeRhsE env rhs
+
+       -- See if we are allowed to float this stuff out of the RHS
+       ; (floats2, rhs2) <- float_from_rhs floats1 rhs1
+
+       -- Make the arity match up
+       ; (floats3, rhs')
+            <- if manifestArity rhs1 <= arity 
+              then return (floats2, cpeEtaExpand arity rhs2)
+              else WARN(True, text "CorePrep: silly extra arguments:" <+> ppr bndr)
+                              -- Note [Silly extra arguments]
+                   (do { v <- newVar (idType bndr)
+                       ; let float = mkFloat False False v rhs2
+                       ; return (addFloat floats2 float, cpeEtaExpand arity (Var v)) })
+
+       -- Record if the binder is evaluated
+       -- and otherwise trim off the unfolding altogether
+       -- It's not used by the code generator; getting rid of it reduces
+       -- heap usage and, since we may be changing uniques, we'd have
+       -- to substitute to keep it right
+       ; let bndr' | exprIsHNF rhs' = bndr `setIdUnfolding` evaldUnfolding
+                          | otherwise      = bndr `setIdUnfolding` noUnfolding
+
+       ; return (floats3, bndr', rhs') }
+  where
+    arity = idArity bndr       -- We must match this arity
+
+    ---------------------
+    float_from_rhs floats rhs
+      | isEmptyFloats floats = return (emptyFloats, rhs)
+      | isTopLevel top_lvl    = float_top    floats rhs
+      | otherwise             = float_nested floats rhs
+
+    ---------------------
+    float_nested floats rhs
+      | wantFloatNested is_rec is_strict_or_unlifted floats rhs
+                  = return (floats, rhs)
+      | otherwise = dont_float floats rhs
+
+    ---------------------
+    float_top floats rhs       -- Urhgh!  See Note [CafInfo and floating]
+      | mayHaveCafRefs (idCafInfo bndr)
+      , allLazyTop floats
+      = return (floats, rhs)
+
+      -- So the top-level binding is marked NoCafRefs
+      | Just (floats', rhs') <- canFloatFromNoCaf floats rhs
+      = return (floats', rhs')
+
+      | otherwise
+      = dont_float floats rhs
+
+    ---------------------
+    dont_float floats rhs
+      -- Non-empty floats, but do not want to float from rhs
+      -- So wrap the rhs in the floats
+      -- But: rhs1 might have lambdas, and we can't
+      --      put them inside a wrapBinds
+      = do { body <- rhsToBodyNF rhs
+          ; return (emptyFloats, wrapBinds floats body) } 
+
+{- Note [Silly extra arguments]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose we had this
+       f{arity=1} = \x\y. e
+We *must* match the arity on the Id, so we have to generate
+        f' = \x\y. e
+       f  = \x. f' x
+
+It's a bizarre case: why is the arity on the Id wrong?  Reason
+(in the days of __inline_me__): 
+        f{arity=0} = __inline_me__ (let v = expensive in \xy. e)
+When InlineMe notes go away this won't happen any more.  But
+it seems good for CorePrep to be robust.
+-}
 
 -- ---------------------------------------------------------------------------
--- Dealing with expressions
+--             CpeRhs: produces a result satisfying CpeRhs
 -- ---------------------------------------------------------------------------
 
-corePrepAnExpr :: CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
-corePrepAnExpr env expr
-  = corePrepExprFloat env expr         `thenUs` \ (floats, expr) ->
-    mkBinds floats expr
-
-
-corePrepExprFloat :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
+cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 -- If
 --     e  ===>  (bs, e')
 -- then        
@@ -365,78 +443,126 @@ corePrepExprFloat :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
 -- For example
 --     f (g x)   ===>   ([v = g x], f v)
 
-corePrepExprFloat env (Var v)
-  = fiddleCCall v                              `thenUs` \ v1 ->
-    let 
-       v2 = lookupCorePrepEnv env v1
-    in
-    maybeSaturate v2 (Var v2) 0 emptyFloats (idType v2)
+cpeRhsE _env expr@(Type {})     = return (emptyFloats, expr)
+cpeRhsE _env expr@(Coercion {}) = return (emptyFloats, expr)
+cpeRhsE _env expr@(Lit {})      = return (emptyFloats, expr)
+cpeRhsE env expr@(Var {})       = cpeApp env expr
 
-corePrepExprFloat env expr@(Type _)
-  = returnUs (emptyFloats, expr)
+cpeRhsE env (Var f `App` _ `App` arg)
+  | f `hasKey` lazyIdKey         -- Replace (lazy a) by a
+  = cpeRhsE env arg              -- See Note [lazyId magic] in MkId
 
-corePrepExprFloat env expr@(Lit lit)
-  = returnUs (emptyFloats, expr)
+cpeRhsE env expr@(App {}) = cpeApp env expr
 
-corePrepExprFloat env (Let bind body)
-  = corePrepBind env bind              `thenUs` \ (env', new_binds) ->
-    corePrepExprFloat env' body                `thenUs` \ (floats, new_body) ->
-    returnUs (new_binds `appendFloats` floats, new_body)
+cpeRhsE env (Let bind expr)
+  = do { (env', new_binds) <- cpeBind NotTopLevel env bind
+       ; (floats, body) <- cpeRhsE env' expr
+       ; return (new_binds `appendFloats` floats, body) }
 
-corePrepExprFloat env (Note n@(SCC _) expr)
-  = corePrepAnExpr env expr            `thenUs` \ expr1 ->
-    deLamFloat expr1                   `thenUs` \ (floats, expr2) ->
-    returnUs (floats, Note n expr2)
+cpeRhsE env (Note note expr)
+  | ignoreNote note
+  = cpeRhsE env expr
+  | otherwise        -- Just SCCs actually
+  = do { body <- cpeBodyNF env expr
+       ; return (emptyFloats, Note note body) }
 
-corePrepExprFloat env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
+cpeRhsE env (Cast expr co)
+   = do { (floats, expr') <- cpeRhsE env expr
+        ; return (floats, Cast expr' co) }
+
+cpeRhsE env expr@(Lam {})
+   = do { let (bndrs,body) = collectBinders expr
+        ; (env', bndrs') <- cloneBndrs env bndrs
+       ; body' <- cpeBodyNF env' body
+       ; return (emptyFloats, mkLams bndrs' body') }
+
+cpeRhsE env (Case (Var id) bndr ty [(DEFAULT,[],expr)])
   | Just (TickBox {}) <- isTickBoxOp_maybe id
-  = corePrepAnExpr env expr            `thenUs` \ expr1 ->
-    deLamFloat expr1                   `thenUs` \ (floats, expr2) ->
-    return (floats, Case (Var id) bndr ty [(DEFAULT,[],expr2)])
-
-corePrepExprFloat env (Note other_note expr)
-  = corePrepExprFloat env expr         `thenUs` \ (floats, expr') ->
-    returnUs (floats, Note other_note expr')
-
-corePrepExprFloat env (Cast expr co)
-  = corePrepExprFloat env expr         `thenUs` \ (floats, expr') ->
-    returnUs (floats, Cast expr' co)
-
-corePrepExprFloat env expr@(Lam _ _)
-  = cloneBndrs env bndrs               `thenUs` \ (env', bndrs') ->
-    corePrepAnExpr env' body           `thenUs` \ body' ->
-    returnUs (emptyFloats, mkLams bndrs' body')
+  = do { body <- cpeBodyNF env expr
+       ; return (emptyFloats, Case (Var id) bndr ty [(DEFAULT,[],body)]) }
+
+cpeRhsE env (Case scrut bndr ty alts)
+  = do { (floats, scrut') <- cpeBody env scrut
+       ; let bndr1 = bndr `setIdUnfolding` evaldUnfolding
+            -- Record that the case binder is evaluated in the alternatives
+       ; (env', bndr2) <- cloneBndr env bndr1
+       ; alts' <- mapM (sat_alt env') alts
+       ; return (floats, Case scrut' bndr2 ty alts') }
   where
-    (bndrs,body) = collectBinders expr
+    sat_alt env (con, bs, rhs)
+       = do { (env2, bs') <- cloneBndrs env bs
+            ; rhs' <- cpeBodyNF env2 rhs
+            ; return (con, bs', rhs') }
+
+-- ---------------------------------------------------------------------------
+--             CpeBody: produces a result satisfying CpeBody
+-- ---------------------------------------------------------------------------
 
-corePrepExprFloat env (Case scrut bndr ty alts)
-  = corePrepExprFloat env scrut                `thenUs` \ (floats1, scrut1) ->
-    deLamFloat scrut1                  `thenUs` \ (floats2, scrut2) ->
-    let
-       bndr1 = bndr `setIdUnfolding` evaldUnfolding
-       -- Record that the case binder is evaluated in the alternatives
-    in
-    cloneBndr env bndr1                        `thenUs` \ (env', bndr2) ->
-    mapUs (sat_alt env') alts          `thenUs` \ alts' ->
-    returnUs (floats1 `appendFloats` floats2 , Case scrut2 bndr2 ty alts')
+cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
+cpeBodyNF env expr 
+  = do { (floats, body) <- cpeBody env expr
+       ; return (wrapBinds floats body) }
+
+--------
+cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
+cpeBody env expr
+  = do { (floats1, rhs) <- cpeRhsE env expr
+       ; (floats2, body) <- rhsToBody rhs
+       ; return (floats1 `appendFloats` floats2, body) }
+
+--------
+rhsToBodyNF :: CpeRhs -> UniqSM CpeBody
+rhsToBodyNF rhs = do { (floats,body) <- rhsToBody rhs
+                    ; return (wrapBinds floats body) }
+
+--------
+rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
+-- Remove top level lambdas by let-binding
+
+rhsToBody (Note n expr)
+        -- You can get things like
+        --      case e of { p -> coerce t (\s -> ...) }
+  = do { (floats, expr') <- rhsToBody expr
+       ; return (floats, Note n expr') }
+
+rhsToBody (Cast e co)
+  = do { (floats, e') <- rhsToBody e
+       ; return (floats, Cast e' co) }
+
+rhsToBody expr@(Lam {})
+  | Just no_lam_result <- tryEtaReducePrep bndrs body
+  = return (emptyFloats, no_lam_result)
+  | all isTyVar bndrs          -- Type lambdas are ok
+  = return (emptyFloats, expr)
+  | otherwise                  -- Some value lambdas
+  = do { fn <- newVar (exprType expr)
+       ; let rhs   = cpeEtaExpand (exprArity expr) expr
+                    float = FloatLet (NonRec fn rhs)
+       ; return (unitFloat float, Var fn) }
   where
-    sat_alt env (con, bs, rhs)
-         = cloneBndrs env bs           `thenUs` \ (env2, bs') ->
-           corePrepAnExpr env2 rhs     `thenUs` \ rhs1 ->
-           deLam rhs1                  `thenUs` \ rhs2 ->
-           returnUs (con, bs', rhs2)
+    (bndrs,body) = collectBinders expr
+
+rhsToBody expr = return (emptyFloats, expr)
+
 
-corePrepExprFloat env expr@(App _ _)
-  = collect_args expr 0  `thenUs` \ (app, (head,depth), ty, floats, ss) ->
-    ASSERT(null ss)    -- make sure we used all the strictness info
+
+-- ---------------------------------------------------------------------------
+--             CpeApp: produces a result satisfying CpeApp
+-- ---------------------------------------------------------------------------
+
+cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
+-- May return a CpeRhs because of saturating primops
+cpeApp env expr 
+  = do { (app, (head,depth), _, floats, ss) <- collect_args expr 0
+       ; MASSERT(null ss)      -- make sure we used all the strictness info
 
        -- Now deal with the function
-    case head of
-      Var fn_id -> maybeSaturate fn_id app depth floats ty
-      _other    -> returnUs (floats, app)
+       ; case head of
+           Var fn_id -> do { sat_app <- maybeSaturate fn_id app depth
+                          ; return (floats, sat_app) }
+           _other    -> return (floats, app) }
 
   where
-
     -- Deconstruct and rebuild the application, floating any non-atomic
     -- arguments to the outside.  We collect the type of the expression,
     -- the head of the application, and the number of actual value arguments,
@@ -445,38 +571,40 @@ corePrepExprFloat env expr@(App _ _)
 
     collect_args
        :: CoreExpr
-       -> Int                            -- current app depth
-       -> UniqSM (CoreExpr,              -- the rebuilt expression
-                  (CoreExpr,Int),        -- the head of the application,
-                                         -- and no. of args it was applied to
-                  Type,                  -- type of the whole expr
-                  Floats,                -- any floats we pulled out
-                  [Demand])              -- remaining argument demands
+       -> Int                     -- Current app depth
+       -> UniqSM (CpeApp,         -- The rebuilt expression
+                  (CoreExpr,Int), -- The head of the application,
+                                  -- and no. of args it was applied to
+                  Type,           -- Type of the whole expr
+                  Floats,         -- Any floats we pulled out
+                  [Demand])       -- Remaining argument demands
 
     collect_args (App fun arg@(Type arg_ty)) depth
-        = collect_args fun depth   `thenUs` \ (fun',hd,fun_ty,floats,ss) ->
-         returnUs (App fun' arg, hd, applyTy fun_ty arg_ty, floats, ss)
+      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
+           ; return (App fun' arg, hd, applyTy fun_ty arg_ty, floats, ss) }
+
+    collect_args (App fun arg@(Coercion arg_co)) depth
+      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun depth
+           ; return (App fun' arg, hd, applyCo fun_ty arg_co, floats, ss) }
 
     collect_args (App fun arg) depth
-        = collect_args fun (depth+1)   `thenUs` \ (fun',hd,fun_ty,floats,ss) ->
-         let
-             (ss1, ss_rest)   = case ss of
-                                  (ss1:ss_rest) -> (ss1,     ss_rest)
-                                  []            -> (lazyDmd, [])
-              (arg_ty, res_ty) = expectJust "corePrepExprFloat:collect_args" $
+      = do { (fun',hd,fun_ty,floats,ss) <- collect_args fun (depth+1)
+          ; let
+              (ss1, ss_rest)   = case ss of
+                                   (ss1:ss_rest) -> (ss1,     ss_rest)
+                                   []            -> (lazyDmd, [])
+              (arg_ty, res_ty) = expectJust "cpeBody:collect_args" $
                                  splitFunTy_maybe fun_ty
-         in
-         corePrepArg env arg (mkDemTy ss1 arg_ty)      `thenUs` \ (fs, arg') ->
-         returnUs (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest)
-
-    collect_args (Var v) depth
-       = fiddleCCall v `thenUs` \ v1 ->
-         let 
-               v2 = lookupCorePrepEnv env v1
-         in
-         returnUs (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts)
+
+           ; (fs, arg') <- cpeArg env (isStrictDmd ss1) arg arg_ty
+           ; return (App fun' arg', hd, res_ty, fs `appendFloats` floats, ss_rest) }
+
+    collect_args (Var v) depth 
+      = do { v1 <- fiddleCCall v
+           ; let v2 = lookupCorePrepEnv env v1
+           ; return (Var v2, (Var v2, depth), idType v2, emptyFloats, stricts) }
        where
-         stricts = case idNewStrictness v of
+         stricts = case idStrictness v of
                        StrictSig (DmdType _ demands _)
                            | listLengthCmp demands depth /= GT -> demands
                                    -- length demands <= depth
@@ -488,244 +616,219 @@ corePrepExprFloat env expr@(App _ _)
                -- partial application might be seq'd
 
     collect_args (Cast fun co) depth
-        = let (_ty1,ty2) = coercionKind co in
-          collect_args fun depth  `thenUs` \ (fun', hd, fun_ty, floats, ss) ->
-         returnUs (Cast fun' co, hd, ty2, floats, ss)
+      = do { let Pair _ty1 ty2 = coercionKind co
+           ; (fun', hd, _, floats, ss) <- collect_args fun depth
+           ; return (Cast fun' co, hd, ty2, floats, ss) }
           
     collect_args (Note note fun) depth
-       | ignore_note note      -- Drop these notes altogether
-                               -- They aren't used by the code generator
-        = collect_args fun depth   `thenUs` \ (fun', hd, fun_ty, floats, ss) ->
-         returnUs (fun', hd, fun_ty, floats, ss)
+      | ignoreNote note         -- Drop these notes altogether
+      = collect_args fun depth  -- They aren't used by the code generator
 
        -- N-variable fun, better let-bind it
-       -- ToDo: perhaps we can case-bind rather than let-bind this closure,
-       -- since it is sure to be evaluated.
     collect_args fun depth
-       = corePrepExprFloat env fun                     `thenUs` \ (fun_floats, fun') ->
-         newVar ty                                     `thenUs` \ fn_id ->
-          mkLocalNonRec fn_id onceDem fun_floats fun'  `thenUs` \ (floats, fn_id') ->
-         returnUs (Var fn_id', (Var fn_id', depth), ty, floats, [])
+      = do { (fun_floats, fun') <- cpeArg env True fun ty
+                         -- The True says that it's sure to be evaluated,
+                         -- so we'll end up case-binding it
+           ; return (fun', (fun', depth), ty, fun_floats, []) }
         where
          ty = exprType fun
 
-    ignore_note        (CoreNote _) = True 
-    ignore_note        InlineMe     = True
-    ignore_note        _other       = False
-       -- We don't ignore SCCs, since they require some code generation
-
-------------------------------------------------------------------------------
--- Building the saturated syntax
+-- ---------------------------------------------------------------------------
+--     CpeArg: produces a result satisfying CpeArg
 -- ---------------------------------------------------------------------------
 
--- maybeSaturate deals with saturating primops and constructors
--- The type is the type of the entire application
-maybeSaturate :: Id -> CoreExpr -> Int -> Floats -> Type -> UniqSM (Floats, CoreExpr)
-maybeSaturate fn expr n_args floats ty
-  | Just DataToTagOp <- isPrimOpId_maybe fn    -- DataToTag must have an evaluated arg
-                                               -- A gruesome special case
-  = saturate_it                `thenUs` \ sat_expr ->
-
-       -- OK, now ensure that the arg is evaluated.
-       -- But (sigh) take into account the lambdas we've now introduced
-    let 
-       (eta_bndrs, eta_body) = collectBinders sat_expr
-    in
-    eval_data2tag_arg eta_body `thenUs` \ (eta_floats, eta_body') -> 
-    if null eta_bndrs then
-       returnUs (floats `appendFloats` eta_floats, eta_body')
-    else
-       mkBinds eta_floats eta_body'            `thenUs` \ eta_body'' ->
-       returnUs (floats, mkLams eta_bndrs eta_body'')
-
-  | hasNoBinding fn = saturate_it      `thenUs` \ sat_expr ->
-                     returnUs (floats, sat_expr)
-
-  | otherwise       = returnUs (floats, expr)
-
+-- This is where we arrange that a non-trivial argument is let-bound
+cpeArg :: CorePrepEnv -> RhsDemand -> CoreArg -> Type
+       -> UniqSM (Floats, CpeTriv)
+cpeArg env is_strict arg arg_ty
+  = do { (floats1, arg1) <- cpeRhsE env arg     -- arg1 can be a lambda
+       ; (floats2, arg2) <- if want_float floats1 arg1 
+                                   then return (floats1, arg1)
+                                   else do { body1 <- rhsToBodyNF arg1
+                                   ; return (emptyFloats, wrapBinds floats1 body1) } 
+               -- Else case: arg1 might have lambdas, and we can't
+               --            put them inside a wrapBinds
+
+       ; if cpe_ExprIsTrivial arg2    -- Do not eta expand a trivial argument
+         then return (floats2, arg2)
+         else do
+       { v <- newVar arg_ty
+       ; let arg3      = cpeEtaExpand (exprArity arg2) arg2
+                    arg_float = mkFloat is_strict is_unlifted v arg3
+       ; return (addFloat floats2 arg_float, Var v) } }
   where
-    fn_arity    = idArity fn
-    excess_arity = fn_arity - n_args
+    is_unlifted = isUnLiftedType arg_ty
+    want_float = wantFloatNested NonRecursive (is_strict || is_unlifted)
+\end{code}
 
-    saturate_it :: UniqSM CoreExpr
-    saturate_it | excess_arity == 0 = returnUs expr
-               | otherwise         = getUniquesUs              `thenUs` \ us ->
-                                     returnUs (etaExpand excess_arity us expr ty)
+Note [Floating unlifted arguments]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider    C (let v* = expensive in v)
 
-       -- Ensure that the argument of DataToTagOp is evaluated
-    eval_data2tag_arg :: CoreExpr -> UniqSM (Floats, CoreExpr)
-    eval_data2tag_arg app@(fun `App` arg)
-       | exprIsHNF arg         -- Includes nullary constructors
-       = returnUs (emptyFloats, app)   -- The arg is evaluated
-       | otherwise                     -- Arg not evaluated, so evaluate it
-       = newVar (exprType arg)         `thenUs` \ arg_id ->
-         let 
-            arg_id1 = setIdUnfolding arg_id evaldUnfolding
-         in
-         returnUs (unitFloat (FloatCase arg_id1 arg False ),
-                   fun `App` Var arg_id1)
+where the "*" indicates "will be demanded".  Usually v will have been
+inlined by now, but let's suppose it hasn't (see Trac #2756).  Then we
+do *not* want to get
 
-    eval_data2tag_arg (Note note app)  -- Scc notes can appear
-       = eval_data2tag_arg app         `thenUs` \ (floats, app') ->
-         returnUs (floats, Note note app')
+     let v* = expensive in C v
 
-    eval_data2tag_arg other    -- Should not happen
-       = pprPanic "eval_data2tag" (ppr other)
+because that has different strictness.  Hence the use of 'allLazy'.
+(NB: the let v* turns into a FloatCase, in mkLocalNonRec.)
 
 
--- ---------------------------------------------------------------------------
--- Precipitating the floating bindings
+------------------------------------------------------------------------------
+-- Building the saturated syntax
 -- ---------------------------------------------------------------------------
 
-floatRhs :: TopLevelFlag -> RecFlag
-        -> Id
-        -> (Floats, CoreExpr)  -- Rhs: let binds in body
-        -> UniqSM (Floats,     -- Floats out of this bind
-                   CoreExpr)   -- Final Rhs
+maybeSaturate deals with saturating primops and constructors
+The type is the type of the entire application
 
-floatRhs top_lvl is_rec bndr (floats, rhs)
-  | isTopLevel top_lvl || exprIsHNF rhs,       -- Float to expose value or 
-    allLazy top_lvl is_rec floats              -- at top level
-  =    -- Why the test for allLazy? 
-       --      v = f (x `divInt#` y)
-       -- we don't want to float the case, even if f has arity 2,
-       -- because floating the case would make it evaluated too early
-    returnUs (floats, rhs)
-    
-  | otherwise
-       -- Don't float; the RHS isn't a value
-  = mkBinds floats rhs         `thenUs` \ rhs' ->
-    returnUs (emptyFloats, rhs')
-
--- mkLocalNonRec is used only for *nested*, *non-recursive* bindings
-mkLocalNonRec :: Id  -> RhsDemand      -- Lhs: id with demand
-             -> Floats -> CoreExpr     -- Rhs: let binds in body
-             -> UniqSM (Floats, Id)    -- The new Id may have an evaldUnfolding, 
-                                       -- to record that it's been evaluated
-
-mkLocalNonRec bndr dem floats rhs
-  | isUnLiftedType (idType bndr)
-       -- If this is an unlifted binding, we always make a case for it.
-  = ASSERT( not (isUnboxedTupleType (idType bndr)) )
-    let
-       float = FloatCase bndr rhs (exprOkForSpeculation rhs)
-    in
-    returnUs (addFloat floats float, evald_bndr)
-
-  | isStrict dem 
-       -- It's a strict let so we definitely float all the bindings
- = let         -- Don't make a case for a value binding,
-               -- even if it's strict.  Otherwise we get
-               --      case (\x -> e) of ...!
-       float | exprIsHNF rhs = FloatLet (NonRec bndr rhs)
-             | otherwise       = FloatCase bndr rhs (exprOkForSpeculation rhs)
-    in
-    returnUs (addFloat floats float, evald_bndr)
+\begin{code}
+maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
+maybeSaturate fn expr n_args
+  | Just DataToTagOp <- isPrimOpId_maybe fn     -- DataToTag must have an evaluated arg
+                                                -- A gruesome special case
+  = saturateDataToTag sat_expr
 
-  | otherwise
-  = floatRhs NotTopLevel NonRecursive bndr (floats, rhs)       `thenUs` \ (floats', rhs') ->
-    returnUs (addFloat floats' (FloatLet (NonRec bndr rhs')),
-             if exprIsHNF rhs' then evald_bndr else bndr)
+  | hasNoBinding fn       -- There's no binding
+  = return sat_expr
 
+  | otherwise 
+  = return expr
   where
-    evald_bndr = bndr `setIdUnfolding` evaldUnfolding
-       -- Record if the binder is evaluated
+    fn_arity    = idArity fn
+    excess_arity = fn_arity - n_args
+    sat_expr     = cpeEtaExpand excess_arity expr
+
+-------------
+saturateDataToTag :: CpeApp -> UniqSM CpeApp
+-- See Note [dataToTag magic]
+saturateDataToTag sat_expr
+  = do { let (eta_bndrs, eta_body) = collectBinders sat_expr
+       ; eta_body' <- eval_data2tag_arg eta_body
+       ; return (mkLams eta_bndrs eta_body') }
+  where
+    eval_data2tag_arg :: CpeApp -> UniqSM CpeBody
+    eval_data2tag_arg app@(fun `App` arg)
+        | exprIsHNF arg         -- Includes nullary constructors
+        = return app           -- The arg is evaluated
+        | otherwise                     -- Arg not evaluated, so evaluate it
+        = do { arg_id <- newVar (exprType arg)
+             ; let arg_id1 = setIdUnfolding arg_id evaldUnfolding
+             ; return (Case arg arg_id1 (exprType app)
+                            [(DEFAULT, [], fun `App` Var arg_id1)]) }
 
+    eval_data2tag_arg (Note note app)  -- Scc notes can appear
+        = do { app' <- eval_data2tag_arg app
+             ; return (Note note app') }
 
-mkBinds :: Floats -> CoreExpr -> UniqSM CoreExpr
-mkBinds (Floats _ binds) body 
-  | isNilOL binds = returnUs body
-  | otherwise    = deLam body          `thenUs` \ body' ->
-                       -- Lambdas are not allowed as the body of a 'let'
-                   returnUs (foldrOL mk_bind body' binds)
-  where
-    mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
-    mk_bind (FloatLet bind)        body = Let bind body
+    eval_data2tag_arg other    -- Should not happen
+       = pprPanic "eval_data2tag" (ppr other)
+\end{code}
 
-etaExpandRhs bndr rhs
-  =    -- Eta expand to match the arity claimed by the binder
-       -- Remember, after CorePrep we must not change arity
-       --
-       -- Eta expansion might not have happened already, 
-       -- because it is done by the simplifier only when 
-       -- there at least one lambda already.
-       -- 
-       -- NB1:we could refrain when the RHS is trivial (which can happen
-       --     for exported things).  This would reduce the amount of code
-       --     generated (a little) and make things a little words for
-       --     code compiled without -O.  The case in point is data constructor
-       --     wrappers.
-       --
-       -- NB2: we have to be careful that the result of etaExpand doesn't
-       --    invalidate any of the assumptions that CorePrep is attempting
-       --    to establish.  One possible cause is eta expanding inside of
-       --    an SCC note - we're now careful in etaExpand to make sure the
-       --    SCC is pushed inside any new lambdas that are generated.
-       --
-       -- NB3: It's important to do eta expansion, and *then* ANF-ising
-       --              f = /\a -> g (h 3)      -- h has arity 2
-       -- If we ANF first we get
-       --              f = /\a -> let s = h 3 in g s
-       -- and now eta expansion gives
-       --              f = /\a -> \ y -> (let s = h 3 in g s) y
-       -- which is horrible.
-       -- Eta expanding first gives
-       --              f = /\a -> \y -> let s = h 3 in g s y
-       --
-    getUniquesUs               `thenUs` \ us ->
-    returnUs (etaExpand arity us rhs (idType bndr))
-  where
-       -- For a GlobalId, take the Arity from the Id.
-       -- It was set in CoreTidy and must not change
-       -- For all others, just expand at will
-    arity | isGlobalId bndr = idArity bndr
-         | otherwise       = exprArity rhs
+Note [dataToTag magic]
+~~~~~~~~~~~~~~~~~~~~~~
+Horrid: we must ensure that the arg of data2TagOp is evaluated
+  (data2tag x) -->  (case x of y -> data2tag y)
+(yuk yuk) take into account the lambdas we've now introduced
 
--- ---------------------------------------------------------------------------
--- Eliminate Lam as a non-rhs (STG doesn't have such a thing)
--- We arrange that they only show up as the RHS of a let(rec)
--- ---------------------------------------------------------------------------
+How might it not be evaluated?  Well, we might have floated it out
+of the scope of a `seq`, or dropped the `seq` altogether.
 
-deLam :: CoreExpr -> UniqSM CoreExpr
--- Takes an expression that may be a lambda, 
--- and returns one that definitely isn't:
---     (\x.e) ==>  let f = \x.e in f
-deLam expr = 
-  deLamFloat expr   `thenUs` \ (floats, expr) ->
-  mkBinds floats expr
 
+%************************************************************************
+%*                                                                     *
+               Simple CoreSyn operations
+%*                                                                     *
+%************************************************************************
+
+\begin{code}
+       -- We don't ignore SCCs, since they require some code generation
+ignoreNote :: Note -> Bool
+-- Tells which notes to drop altogether; they are ignored by code generation
+-- Do not ignore SCCs!
+-- It's important that we do drop InlineMe notes; for example
+--    unzip = __inline_me__ (/\ab. foldr (..) (..))
+-- Here unzip gets arity 1 so we'll eta-expand it. But we don't
+-- want to get this:
+--     unzip = /\ab \xs. (__inline_me__ ...) a b xs
+ignoreNote (CoreNote _) = True 
+ignoreNote _other       = False
+
+
+cpe_ExprIsTrivial :: CoreExpr -> Bool
+-- Version that doesn't consider an scc annotation to be trivial.
+cpe_ExprIsTrivial (Var _)                  = True
+cpe_ExprIsTrivial (Type _)                 = True
+cpe_ExprIsTrivial (Coercion _)             = True
+cpe_ExprIsTrivial (Lit _)                  = True
+cpe_ExprIsTrivial (App e arg)              = isTypeArg arg && cpe_ExprIsTrivial e
+cpe_ExprIsTrivial (Note n e)               = notSccNote n  && cpe_ExprIsTrivial e
+cpe_ExprIsTrivial (Cast e _)               = cpe_ExprIsTrivial e
+cpe_ExprIsTrivial (Lam b body) | isTyVar b = cpe_ExprIsTrivial body
+cpe_ExprIsTrivial _                        = False
+\end{code}
 
-deLamFloat :: CoreExpr -> UniqSM (Floats, CoreExpr)
--- Remove top level lambdas by let-bindinig
+-- -----------------------------------------------------------------------------
+--     Eta reduction
+-- -----------------------------------------------------------------------------
 
-deLamFloat (Note n expr)
-  =    -- You can get things like
-       --      case e of { p -> coerce t (\s -> ...) }
-    deLamFloat expr    `thenUs` \ (floats, expr') ->
-    returnUs (floats, Note n expr')
+Note [Eta expansion]
+~~~~~~~~~~~~~~~~~~~~~
+Eta expand to match the arity claimed by the binder Remember,
+CorePrep must not change arity
+
+Eta expansion might not have happened already, because it is done by
+the simplifier only when there at least one lambda already.
+
+NB1:we could refrain when the RHS is trivial (which can happen
+    for exported things).  This would reduce the amount of code
+    generated (a little) and make things a little words for
+    code compiled without -O.  The case in point is data constructor
+    wrappers.
+
+NB2: we have to be careful that the result of etaExpand doesn't
+   invalidate any of the assumptions that CorePrep is attempting
+   to establish.  One possible cause is eta expanding inside of
+   an SCC note - we're now careful in etaExpand to make sure the
+   SCC is pushed inside any new lambdas that are generated.
+
+Note [Eta expansion and the CorePrep invariants]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+It turns out to be much much easier to do eta expansion
+*after* the main CorePrep stuff.  But that places constraints
+on the eta expander: given a CpeRhs, it must return a CpeRhs.
+
+For example here is what we do not want:
+               f = /\a -> g (h 3)      -- h has arity 2
+After ANFing we get
+               f = /\a -> let s = h 3 in g s
+and now we do NOT want eta expansion to give
+               f = /\a -> \ y -> (let s = h 3 in g s) y
+
+Instead CoreArity.etaExpand gives
+               f = /\a -> \y -> let s = h 3 in g s y
 
-deLamFloat (Cast e co)
-  = deLamFloat e       `thenUs` \ (floats, e') ->
-    returnUs (floats, Cast e' co)
+\begin{code}
+cpeEtaExpand :: Arity -> CpeRhs -> CpeRhs
+cpeEtaExpand arity expr
+  | arity == 0 = expr
+  | otherwise  = etaExpand arity expr
+\end{code}
 
-deLamFloat expr 
-  | null bndrs = returnUs (emptyFloats, expr)
-  | otherwise 
-  = case tryEta bndrs body of
-      Just no_lam_result -> returnUs (emptyFloats, no_lam_result)
-      Nothing           -> newVar (exprType expr)      `thenUs` \ fn ->
-                           returnUs (unitFloat (FloatLet (NonRec fn expr)), 
-                                     Var fn)
-  where
-    (bndrs,body) = collectBinders expr
+-- -----------------------------------------------------------------------------
+--     Eta reduction
+-- -----------------------------------------------------------------------------
 
--- Why try eta reduction?  Hasn't the simplifier already done eta?
--- But the simplifier only eta reduces if that leaves something
--- trivial (like f, or f Int).  But for deLam it would be enough to
--- get to a partial application:
---     \xs. map f xs ==> map f
+Why try eta reduction?  Hasn't the simplifier already done eta?
+But the simplifier only eta reduces if that leaves something
+trivial (like f, or f Int).  But for deLam it would be enough to
+get to a partial application:
+       case x of { p -> \xs. map f xs }
+    ==> case x of { p -> map f }
 
-tryEta bndrs expr@(App _ _)
+\begin{code}
+tryEtaReducePrep :: [CoreBndr] -> CoreExpr -> Maybe CoreExpr
+tryEtaReducePrep bndrs expr@(App _ _)
   | ok_to_eta_reduce f &&
     n_remaining >= 0 &&
     and (zipWith ok bndrs last_args) &&
@@ -739,21 +842,21 @@ tryEta bndrs expr@(App _ _)
     n_remaining = length args - length bndrs
 
     ok bndr (Var arg) = bndr == arg
-    ok bndr other     = False
+    ok _    _         = False
 
          -- we can't eta reduce something which must be saturated.
     ok_to_eta_reduce (Var f) = not (hasNoBinding f)
     ok_to_eta_reduce _       = False --safe. ToDo: generalise
 
-tryEta bndrs (Let bind@(NonRec b r) body)
+tryEtaReducePrep bndrs (Let bind@(NonRec _ r) body)
   | not (any (`elemVarSet` fvs) bndrs)
-  = case tryEta bndrs body of
+  = case tryEtaReducePrep bndrs body of
        Just e -> Just (Let bind e)
        Nothing -> Nothing
   where
     fvs = exprFreeVars r
 
-tryEta bndrs _ = Nothing
+tryEtaReducePrep _ _ = Nothing
 \end{code}
 
 
@@ -762,35 +865,186 @@ tryEta bndrs _ = Nothing
 -- -----------------------------------------------------------------------------
 
 \begin{code}
-data RhsDemand
-     = RhsDemand { isStrict :: Bool,  -- True => used at least once
-                   isOnceDem   :: Bool   -- True => used at most once
-                 }
+type RhsDemand = Bool  -- True => used strictly; hence not top-level, non-recursive
+\end{code}
+
+%************************************************************************
+%*                                                                     *
+               Floats
+%*                                                                     *
+%************************************************************************
+
+\begin{code}
+data FloatingBind 
+  = FloatLet CoreBind   -- Rhs of bindings are CpeRhss
+                        -- They are always of lifted type;
+                        -- unlifted ones are done with FloatCase
+ | FloatCase 
+      Id CpeBody 
+      Bool             -- The bool indicates "ok-for-speculation"
 
-mkDem :: Demand -> Bool -> RhsDemand
-mkDem strict once = RhsDemand (isStrictDmd strict) once
+data Floats = Floats OkToSpec (OrdList FloatingBind)
 
-mkDemTy :: Demand -> Type -> RhsDemand
-mkDemTy strict ty = RhsDemand (isStrictDmd strict) 
-                             False {- For now -}
+instance Outputable FloatingBind where
+  ppr (FloatLet b) = ppr b
+  ppr (FloatCase b r ok) = brackets (ppr ok) <+> ppr b <+> equals <+> ppr r
+
+instance Outputable Floats where
+  ppr (Floats flag fs) = ptext (sLit "Floats") <> brackets (ppr flag) <+>
+                         braces (vcat (map ppr (fromOL fs)))
+
+instance Outputable OkToSpec where
+  ppr OkToSpec    = ptext (sLit "OkToSpec")
+  ppr IfUnboxedOk = ptext (sLit "IfUnboxedOk")
+  ppr NotOkToSpec = ptext (sLit "NotOkToSpec")
+-- Can we float these binds out of the rhs of a let?  We cache this decision
+-- to avoid having to recompute it in a non-linear way when there are
+-- deeply nested lets.
+data OkToSpec
+   = OkToSpec          -- Lazy bindings of lifted type
+   | IfUnboxedOk       -- A mixture of lazy lifted bindings and n
+                       -- ok-to-speculate unlifted bindings
+   | NotOkToSpec       -- Some not-ok-to-speculate unlifted bindings
+
+mkFloat :: Bool -> Bool -> Id -> CpeRhs -> FloatingBind
+mkFloat is_strict is_unlifted bndr rhs
+  | use_case  = FloatCase bndr rhs (exprOkForSpeculation rhs)
+  | otherwise = FloatLet (NonRec bndr rhs)
+  where
+    use_case = is_unlifted || is_strict && not (exprIsHNF rhs)
+               -- Don't make a case for a value binding,
+               -- even if it's strict.  Otherwise we get
+               --      case (\x -> e) of ...!
+             
+emptyFloats :: Floats
+emptyFloats = Floats OkToSpec nilOL
 
-bdrDem :: Id -> RhsDemand
-bdrDem id = mkDem (idNewDemandInfo id)
-                 False {- For now -}
+isEmptyFloats :: Floats -> Bool
+isEmptyFloats (Floats _ bs) = isNilOL bs
 
--- safeDem :: RhsDemand
--- safeDem = RhsDemand False False  -- always safe to use this
+wrapBinds :: Floats -> CpeBody -> CpeBody
+wrapBinds (Floats _ binds) body
+  = foldrOL mk_bind body binds
+  where
+    mk_bind (FloatCase bndr rhs _) body = Case rhs bndr (exprType body) [(DEFAULT, [], body)]
+    mk_bind (FloatLet bind)        body = Let bind body
 
-onceDem :: RhsDemand
-onceDem = RhsDemand False True   -- used at most once
-\end{code}
+addFloat :: Floats -> FloatingBind -> Floats
+addFloat (Floats ok_to_spec floats) new_float
+  = Floats (combine ok_to_spec (check new_float)) (floats `snocOL` new_float)
+  where
+    check (FloatLet _) = OkToSpec
+    check (FloatCase _ _ ok_for_spec) 
+       | ok_for_spec  =  IfUnboxedOk
+       | otherwise    =  NotOkToSpec
+       -- The ok-for-speculation flag says that it's safe to
+       -- float this Case out of a let, and thereby do it more eagerly
+       -- We need the top-level flag because it's never ok to float
+       -- an unboxed binding to the top level
 
+unitFloat :: FloatingBind -> Floats
+unitFloat = addFloat emptyFloats
+
+appendFloats :: Floats -> Floats -> Floats
+appendFloats (Floats spec1 floats1) (Floats spec2 floats2)
+  = Floats (combine spec1 spec2) (floats1 `appOL` floats2)
+
+concatFloats :: [Floats] -> OrdList FloatingBind
+concatFloats = foldr (\ (Floats _ bs1) bs2 -> appOL bs1 bs2) nilOL
+
+combine :: OkToSpec -> OkToSpec -> OkToSpec
+combine NotOkToSpec _ = NotOkToSpec
+combine _ NotOkToSpec = NotOkToSpec
+combine IfUnboxedOk _ = IfUnboxedOk
+combine _ IfUnboxedOk = IfUnboxedOk
+combine _ _           = OkToSpec
+    
+deFloatTop :: Floats -> [CoreBind]
+-- For top level only; we don't expect any FloatCases
+deFloatTop (Floats _ floats)
+  = foldrOL get [] floats
+  where
+    get (FloatLet b) bs = occurAnalyseRHSs b : bs
+    get b            _  = pprPanic "corePrepPgm" (ppr b)
+    
+    -- See Note [Dead code in CorePrep]
+    occurAnalyseRHSs (NonRec x e) = NonRec x (occurAnalyseExpr e)
+    occurAnalyseRHSs (Rec xes)    = Rec [(x, occurAnalyseExpr e) | (x, e) <- xes]
+
+-------------------------------------------
+canFloatFromNoCaf ::  Floats -> CpeRhs -> Maybe (Floats, CpeRhs)
+       -- Note [CafInfo and floating]
+canFloatFromNoCaf (Floats ok_to_spec fs) rhs
+  | OkToSpec <- ok_to_spec          -- Worth trying
+  , Just (subst, fs') <- go (emptySubst, nilOL) (fromOL fs)
+  = Just (Floats OkToSpec fs', subst_expr subst rhs)
+  | otherwise              
+  = Nothing
+  where
+    subst_expr = substExpr (text "CorePrep")
 
+    go :: (Subst, OrdList FloatingBind) -> [FloatingBind]
+       -> Maybe (Subst, OrdList FloatingBind)
+
+    go (subst, fbs_out) [] = Just (subst, fbs_out)
+    
+    go (subst, fbs_out) (FloatLet (NonRec b r) : fbs_in) 
+      | rhs_ok r
+      = go (subst', fbs_out `snocOL` new_fb) fbs_in
+      where
+        (subst', b') = set_nocaf_bndr subst b
+        new_fb = FloatLet (NonRec b' (subst_expr subst r))
+
+    go (subst, fbs_out) (FloatLet (Rec prs) : fbs_in)
+      | all rhs_ok rs
+      = go (subst', fbs_out `snocOL` new_fb) fbs_in
+      where
+        (bs,rs) = unzip prs
+        (subst', bs') = mapAccumL set_nocaf_bndr subst bs
+        rs' = map (subst_expr subst') rs
+        new_fb = FloatLet (Rec (bs' `zip` rs'))
+
+    go _ _ = Nothing     -- Encountered a caffy binding
+
+    ------------
+    set_nocaf_bndr subst bndr 
+      = (extendIdSubst subst bndr (Var bndr'), bndr')
+      where
+        bndr' = bndr `setIdCafInfo` NoCafRefs
+
+    ------------
+    rhs_ok :: CoreExpr -> Bool
+    -- We can only float to top level from a NoCaf thing if
+    -- the new binding is static. However it can't mention
+    -- any non-static things or it would *already* be Caffy
+    rhs_ok = rhsIsStatic (\_ -> False)
+
+wantFloatNested :: RecFlag -> Bool -> Floats -> CpeRhs -> Bool
+wantFloatNested is_rec strict_or_unlifted floats rhs
+  =  isEmptyFloats floats
+  || strict_or_unlifted
+  || (allLazyNested is_rec floats && exprIsHNF rhs)
+       -- Why the test for allLazyNested? 
+       --      v = f (x `divInt#` y)
+       -- we don't want to float the case, even if f has arity 2,
+       -- because floating the case would make it evaluated too early
+
+allLazyTop :: Floats -> Bool
+allLazyTop (Floats OkToSpec _) = True
+allLazyTop _                  = False
+
+allLazyNested :: RecFlag -> Floats -> Bool
+allLazyNested _      (Floats OkToSpec    _) = True
+allLazyNested _      (Floats NotOkToSpec _) = False
+allLazyNested is_rec (Floats IfUnboxedOk _) = isNonRec is_rec
+\end{code}
 
 
 %************************************************************************
 %*                                                                     *
-\subsection{Cloning}
+               Cloning
 %*                                                                     *
 %************************************************************************
 
@@ -807,6 +1061,9 @@ emptyCorePrepEnv = CPE emptyVarEnv
 extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
 extendCorePrepEnv (CPE env) id id' = CPE (extendVarEnv env id id')
 
+extendCorePrepEnvList :: CorePrepEnv -> [(Id,Id)] -> CorePrepEnv
+extendCorePrepEnvList (CPE env) prs = CPE (extendVarEnvList env prs)
+
 lookupCorePrepEnv :: CorePrepEnv -> Id -> Id
 lookupCorePrepEnv (CPE env) id
   = case lookupVarEnv env id of
@@ -818,21 +1075,24 @@ lookupCorePrepEnv (CPE env) id
 -- ---------------------------------------------------------------------------
 
 cloneBndrs :: CorePrepEnv -> [Var] -> UniqSM (CorePrepEnv, [Var])
-cloneBndrs env bs = mapAccumLUs cloneBndr env bs
+cloneBndrs env bs = mapAccumLM cloneBndr env bs
 
 cloneBndr  :: CorePrepEnv -> Var -> UniqSM (CorePrepEnv, Var)
 cloneBndr env bndr
-  | isLocalId bndr
-  = getUniqueUs   `thenUs` \ uniq ->
-    let
-       bndr' = setVarUnique bndr uniq
-    in
-    returnUs (extendCorePrepEnv env bndr bndr', bndr')
+  | isLocalId bndr, not (isCoVar bndr)
+  = do bndr' <- setVarUnique bndr <$> getUniqueM
+       
+       -- We are going to OccAnal soon, so drop (now-useless) rules/unfoldings
+       -- so that we can drop more stuff as dead code.
+       -- See also Note [Dead code in CorePrep]
+       let bndr'' = bndr' `setIdUnfolding` noUnfolding
+                          `setIdSpecialisation` emptySpecInfo
+       return (extendCorePrepEnv env bndr bndr'', bndr'')
 
   | otherwise  -- Top level things, which we don't want
                -- to clone, have become GlobalIds by now
-               -- And we don't clone tyvars
-  = returnUs (env, bndr)
+               -- And we don't clone tyvars, or coercion variables
+  = return (env, bndr)
   
 
 ------------------------------------------------------------------------------
@@ -842,9 +1102,8 @@ cloneBndr env bndr
 
 fiddleCCall :: Id -> UniqSM Id
 fiddleCCall id 
-  | isFCallId id = getUniqueUs         `thenUs` \ uniq ->
-                  returnUs (id `setVarUnique` uniq)
-  | otherwise    = returnUs id
+  | isFCallId id = (id `setVarUnique`) <$> getUniqueM
+  | otherwise    = return id
 
 ------------------------------------------------------------------------------
 -- Generating new binders
@@ -852,7 +1111,7 @@ fiddleCCall id
 
 newVar :: Type -> UniqSM Id
 newVar ty
- = seqType ty                  `seq`
-   getUniqueUs                 `thenUs` \ uniq ->
-   returnUs (mkSysLocal FSLIT("sat") uniq ty)
+ = seqType ty `seq` do
+     uniq <- getUniqueM
+     return (mkSysLocal (fsLit "sat") uniq ty)
 \end{code}