Kill off complications in CoreFVs

When doing type-in-type, Richard introduce some substantial
complications in CoreFVs, gathering types and free variables
of type.  In Trac #13160 we decided that this complication was
unnecessary, so this patch removes it.

Unfortnately I then fell down a twisty rabbit hole.  Roughly:

* An apparently-innocuous change in the AnnApp case of
  fiExpr made the fuction a little bit stricter, so we ended
  up peering into the arguments when we didn't before (namely
  when there are no bindings being floated inwards).  I've
  rejigged it so that it's not fragile any more.

* Peering into the arguments was sometimes expensive, becuase
  exprIsExpandable can be expensive because it looks deeply into
  the expression.

* The combination of the two led to a non-linear explosion
  of work when the argument of a function is a deeep nest
  of constructors.  This bug has been lurking for ages.

  I solved it by replacing exprIsExpandable with exprIsHNF
  + exprIsTrivial; see Note [noFloatInto considerations]

* The code around floating case-expressions turned out to be
  very delicate, because can_fail primops (which we want to
  float inwards) can't be floated outwards; so we have to be
  careful never to float them too far. Note [Floating primops]
  has the details

* I ended up refactoring some rather opaque code in
  sepBindsByDropPoint.

Working all this out meant that I rewrote quite a bit of
code, so it's now a reasonably substantial patch.  But it's
a net improvement.

2 files changed, 255 insertions, 219 deletions

diff --git a/compiler/coreSyn/CoreFVs.hs b/compiler/coreSyn/CoreFVs.hs
index 5997a9c4b6..4dc1ed2f4a 100644
--- a/compiler/coreSyn/CoreFVs.hs
+++ b/compiler/coreSyn/CoreFVs.hs
@@ -55,9 +55,7 @@ module CoreFVs (
         freeVars,               -- CoreExpr -> CoreExprWithFVs
         freeVarsBind,           -- CoreBind -> DVarSet -> (DVarSet, CoreBindWithFVs)
         freeVarsOf,             -- CoreExprWithFVs -> DIdSet
-        freeVarsOfType,         -- CoreExprWithFVs -> TyCoVarSet
-        freeVarsOfAnn, freeVarsOfTypeAnn,
-        exprTypeFV              -- CoreExprWithFVs -> Type
+        freeVarsOfAnn
     ) where
 
 #include "HsVersions.h"
@@ -68,7 +66,6 @@ import IdInfo
 import NameSet
 import UniqSet
 import Unique (Uniquable (..))
-import Literal ( literalType )
 import Name
 import VarSet
 import Var
@@ -78,7 +75,6 @@ import TyCon
 import CoAxiom
 import FamInstEnv
 import TysPrim( funTyConName )
-import Coercion
 import Maybes( orElse )
 import Util
 import BasicTypes( Activation )
@@ -539,10 +535,7 @@ The free variable pass annotates every node in the expression with its
 NON-GLOBAL free variables and type variables.
 -}
 
-data FVAnn = FVAnn { fva_fvs    :: DVarSet   -- free in expression
-                   , fva_ty_fvs :: DVarSet   -- free only in expression's type
-                   , fva_ty     :: Type      -- expression's type
-                   }
+type FVAnn = DVarSet
 
 -- | Every node in a binding group annotated with its
 -- (non-global) free variables, both Ids and TyVars, and type.
@@ -558,23 +551,11 @@ type CoreAltWithFVs = AnnAlt Id FVAnn
 
 freeVarsOf :: CoreExprWithFVs -> DIdSet
 -- ^ Inverse function to 'freeVars'
-freeVarsOf (FVAnn { fva_fvs = fvs }, _) = fvs
-
--- | Extract the vars free in an annotated expression's type
-freeVarsOfType :: CoreExprWithFVs -> DTyCoVarSet
-freeVarsOfType (FVAnn { fva_ty_fvs = ty_fvs }, _) = ty_fvs
-
--- | Extract the type of an annotated expression. (This is cheap.)
-exprTypeFV :: CoreExprWithFVs -> Type
-exprTypeFV (FVAnn { fva_ty = ty }, _) = ty
+freeVarsOf (fvs, _) = fvs
 
 -- | Extract the vars reported in a FVAnn
 freeVarsOfAnn :: FVAnn -> DIdSet
-freeVarsOfAnn = fva_fvs
-
--- | Extract the type-level vars reported in a FVAnn
-freeVarsOfTypeAnn :: FVAnn -> DTyCoVarSet
-freeVarsOfTypeAnn = fva_ty_fvs
+freeVarsOfAnn fvs = fvs
 
 noFVs :: VarSet
 noFVs = emptyVarSet
@@ -594,10 +575,9 @@ delBindersFV bs fvs = foldr delBinderFV fvs bs
 delBinderFV :: Var -> DVarSet -> DVarSet
 -- This way round, so we can do it multiple times using foldr
 
--- (b `delBinderFV` s) removes the binder b from the free variable set s,
--- but *adds* to s
---
---      the free variables of b's type
+-- (b `delBinderFV` s)
+--   * removes the binder b from the free variable set s,
+--   * AND *adds* to s the free variables of b's type
 --
 -- This is really important for some lambdas:
 --      In (\x::a -> x) the only mention of "a" is in the binder.
@@ -733,48 +713,33 @@ freeVars = go
   where
     go :: CoreExpr -> CoreExprWithFVs
     go (Var v)
-      = (FVAnn fvs ty_fvs (idType v), AnnVar v)
+      | isLocalVar v = (aFreeVar v `unionFVs` ty_fvs, AnnVar v)
+      | otherwise    = (emptyDVarSet,                 AnnVar v)
       where
-            -- ToDo: insert motivating example for why we *need*
-            -- to include the idSpecVars in the FV list.
-            --      Actually [June 98] I don't think it's necessary
-            -- fvs = fvs_v `unionVarSet` idSpecVars v
-
-        (fvs, ty_fvs)
-            | isLocalVar v = (aFreeVar v `unionFVs` ty_fvs, dVarTypeTyCoVars v)
-            | otherwise    = (emptyDVarSet, emptyDVarSet)
+        ty_fvs = dVarTypeTyCoVars v  -- Do we need this?
 
-    go (Lit lit) = (FVAnn emptyDVarSet emptyDVarSet (literalType lit), AnnLit lit)
+    go (Lit lit) = (emptyDVarSet, AnnLit lit)
     go (Lam b body)
-      = ( FVAnn { fva_fvs    = b_fvs `unionFVs` (b `delBinderFV` body_fvs)
-                , fva_ty_fvs = b_fvs `unionFVs` (b `delBinderFV` body_ty_fvs)
-                , fva_ty     = mkFunTy b_ty body_ty }
+      = ( b_fvs `unionFVs` (b `delBinderFV` body_fvs)
         , AnnLam b body' )
       where
-        body'@(FVAnn { fva_fvs = body_fvs, fva_ty_fvs = body_ty_fvs
-                     , fva_ty = body_ty }, _) = go body
+        body'@(body_fvs, _) = go body
         b_ty  = idType b
         b_fvs = tyCoVarsOfTypeDSet b_ty
 
     go (App fun arg)
-      = ( FVAnn { fva_fvs    = freeVarsOf fun' `unionFVs` freeVarsOf arg'
-                , fva_ty_fvs = tyCoVarsOfTypeDSet res_ty
-                , fva_ty     = res_ty }
+      = ( freeVarsOf fun' `unionFVs` freeVarsOf arg'
         , AnnApp fun' arg' )
       where
         fun'   = go fun
-        fun_ty = exprTypeFV fun'
         arg'   = go arg
-        res_ty = applyTypeToArg fun_ty arg
 
     go (Case scrut bndr ty alts)
-      = ( FVAnn { fva_fvs = (bndr `delBinderFV` alts_fvs)
-                            `unionFVs` freeVarsOf scrut2
-                            `unionFVs` tyCoVarsOfTypeDSet ty
-                           -- don't need to look at (idType bndr)
-                           -- b/c that's redundant with scrut
-                , fva_ty_fvs = tyCoVarsOfTypeDSet ty
-                , fva_ty     = ty }
+      = ( (bndr `delBinderFV` alts_fvs)
+           `unionFVs` freeVarsOf scrut2
+           `unionFVs` tyCoVarsOfTypeDSet ty
+          -- don't need to look at (idType bndr)
+          -- b/c that's redundant with scrut
         , AnnCase scrut2 bndr ty alts2 )
       where
         scrut2 = go scrut
@@ -788,45 +753,25 @@ freeVars = go
                                  rhs2 = go rhs
 
     go (Let bind body)
-      = ( FVAnn { fva_fvs    = bind_fvs
-                , fva_ty_fvs = freeVarsOfType body2
-                , fva_ty     = exprTypeFV body2 }
-        , AnnLet bind2 body2 )
+      = (bind_fvs, AnnLet bind2 body2)
       where
         (bind2, bind_fvs) = freeVarsBind bind (freeVarsOf body2)
         body2             = go body
 
     go (Cast expr co)
-      = ( FVAnn (freeVarsOf expr2 `unionFVs` cfvs) (tyCoVarsOfTypeDSet to_ty) to_ty
-        , AnnCast expr2 (c_ann, co) )
+      = ( freeVarsOf expr2 `unionFVs` cfvs
+        , AnnCast expr2 (cfvs, co) )
       where
         expr2 = go expr
         cfvs  = tyCoVarsOfCoDSet co
-        c_ann = FVAnn cfvs (tyCoVarsOfTypeDSet co_ki) co_ki
-        co_ki = coercionType co
-        Just (_, to_ty) = splitCoercionType_maybe co_ki
-
 
     go (Tick tickish expr)
-      = ( FVAnn { fva_fvs    = tickishFVs tickish `unionFVs` freeVarsOf expr2
-                , fva_ty_fvs = freeVarsOfType expr2
-                , fva_ty     = exprTypeFV expr2 }
+      = ( tickishFVs tickish `unionFVs` freeVarsOf expr2
         , AnnTick tickish expr2 )
       where
         expr2 = go expr
         tickishFVs (Breakpoint _ ids) = mkDVarSet ids
         tickishFVs _                  = emptyDVarSet
 
-    go (Type ty) = ( FVAnn (tyCoVarsOfTypeDSet ty)
-                           (tyCoVarsOfTypeDSet ki)
-                           ki
-                   , AnnType ty)
-      where
-        ki = typeKind ty
-
-    go (Coercion co) = ( FVAnn (tyCoVarsOfCoDSet co)
-                               (tyCoVarsOfTypeDSet ki)
-                               ki
-                       , AnnCoercion co)
-      where
-        ki = coercionType co
+    go (Type ty)     = (tyCoVarsOfTypeDSet ty, AnnType ty)
+    go (Coercion co) = (tyCoVarsOfCoDSet co, AnnCoercion co)
diff --git a/compiler/simplCore/FloatIn.hs b/compiler/simplCore/FloatIn.hs
index 4d5a564257..34b0b42365 100644
--- a/compiler/simplCore/FloatIn.hs
+++ b/compiler/simplCore/FloatIn.hs
@@ -13,6 +13,7 @@ then discover that they aren't needed in the chosen branch.
 -}
 
 {-# LANGUAGE CPP #-}
+{-# OPTIONS_GHC -fprof-auto #-}
 
 module FloatIn ( floatInwards ) where
 
@@ -21,18 +22,17 @@ module FloatIn ( floatInwards ) where
 import CoreSyn
 import MkCore
 import HscTypes         ( ModGuts(..) )
-import CoreUtils        ( exprIsDupable, exprIsExpandable,
-                          exprOkForSideEffects, mkTicks )
+import CoreUtils
 import CoreFVs
 import CoreMonad        ( CoreM )
 import Id               ( isOneShotBndr, idType, isJoinId, isJoinId_maybe )
 import Var
-import Type             ( isUnliftedType )
+import Type
 import VarSet
 import Util
 import DynFlags
 import Outputable
-import Data.List        ( mapAccumL )
+-- import Data.List        ( mapAccumL )
 import BasicTypes       ( RecFlag(..), isRec )
 
 {-
@@ -151,7 +151,6 @@ fiExpr dflags to_drop (_, AnnCast expr (co_ann, co))
     [drop_here, e_drop, co_drop]
       = sepBindsByDropPoint dflags False
           [freeVarsOf expr, freeVarsOfAnn co_ann]
-          (freeVarsOfType expr `unionDVarSet` freeVarsOfTypeAnn co_ann)
           to_drop
 
 {-
@@ -161,33 +160,45 @@ pull out any silly ones.
 -}
 
 fiExpr dflags to_drop ann_expr@(_,AnnApp {})
-  = mkTicks ticks $ wrapFloats drop_here $ wrapFloats extra_drop $
+  = wrapFloats drop_here $ wrapFloats extra_drop $
+    mkTicks ticks $
     mkApps (fiExpr dflags fun_drop ann_fun)
            (zipWith (fiExpr dflags) arg_drops ann_args)
   where
     (ann_fun, ann_args, ticks) = collectAnnArgsTicks tickishFloatable ann_expr
-    (extra_fvs0, fun_fvs)
-      | (_, AnnVar _) <- ann_fun = (freeVarsOf ann_fun, emptyDVarSet)
+    fun_ty  = exprType (deAnnotate ann_fun)
+    fun_fvs = freeVarsOf ann_fun
+    arg_fvs = map freeVarsOf ann_args
+
+    (drop_here : extra_drop : fun_drop : arg_drops)
+       = sepBindsByDropPoint dflags False
+                             (extra_fvs : fun_fvs : arg_fvs)
+                             to_drop
+         -- Shortcut behaviour: if to_drop is empty,
+         -- sepBindsByDropPoint returns a suitable bunch of empty
+         -- lists without evaluating extra_fvs, and hence without
+         -- peering into each argument
+
+    (_, extra_fvs) = foldl add_arg (fun_ty, extra_fvs0) ann_args
+    extra_fvs0 = case ann_fun of
+                   (_, AnnVar _) -> fun_fvs
+                   _             -> emptyDVarSet
           -- Don't float the binding for f into f x y z; see Note [Join points]
           -- for why we *can't* do it when f is a join point. (If f isn't a
           -- join point, floating it in isn't especially harmful but it's
           -- useless since the simplifier will immediately float it back out.)
-      | otherwise                = (emptyDVarSet, freeVarsOf ann_fun)
-    (extra_fvs, arg_fvs) = mapAccumL mk_arg_fvs extra_fvs0 ann_args
 
-    mk_arg_fvs :: FreeVarSet -> CoreExprWithFVs -> (FreeVarSet, FreeVarSet)
-    mk_arg_fvs extra_fvs ann_arg
-      | noFloatIntoRhs False NonRecursive ann_arg
-      = (extra_fvs `unionDVarSet` freeVarsOf ann_arg, emptyDVarSet)
-      | otherwise
-      = (extra_fvs, freeVarsOf ann_arg)
+    add_arg :: (Type,FreeVarSet) -> CoreExprWithFVs -> (Type,FreeVarSet)
+    add_arg (fun_ty, extra_fvs) (_, AnnType ty)
+      = (piResultTy fun_ty ty, extra_fvs)
 
-    drop_here : extra_drop : fun_drop : arg_drops
-      = sepBindsByDropPoint dflags False
-          (extra_fvs : fun_fvs : arg_fvs)
-          (freeVarsOfType ann_fun `unionDVarSet`
-           mapUnionDVarSet freeVarsOfType ann_args)
-          to_drop
+    add_arg (fun_ty, extra_fvs) (arg_fvs, arg)
+      | noFloatIntoArg arg arg_ty
+      = (res_ty, extra_fvs `unionDVarSet` arg_fvs)
+      | otherwise
+      = (res_ty, extra_fvs)
+      where
+       (arg_ty, res_ty) = splitFunTy fun_ty
 
 {-
 Note [Do not destroy the let/app invariant]
@@ -260,16 +271,37 @@ it's non-recursive, so we float only into non-recursive join points.)
 
 Urk! if all are tyvars, and we don't float in, we may miss an
       opportunity to float inside a nested case branch
+
+
+Note [Floating coercions]
+~~~~~~~~~~~~~~~~~~~~~~~~~
+We could, in principle, have a coercion binding like
+   case f x of co { DEFAULT -> e1 e2 }
+It's not common to have a function that returns a coercion, but nothing
+in Core prohibits it.  If so, 'co' might be mentioned in e1 or e2
+/only in a type/.  E.g. suppose e1 was
+  let (x :: Int |> co) = blah in blah2
+
+
+But, with coercions appearing in types, there is a complication: we
+might be floating in a "strict let" -- that is, a case. Case expressions
+mention their return type. We absolutely can't float a coercion binding
+inward to the point that the type of the expression it's about to wrap
+mentions the coercion. So we include the union of the sets of free variables
+of the types of all the drop points involved. If any of the floaters
+bind a coercion variable mentioned in any of the types, that binder must
+be dropped right away.
+
 -}
 
 fiExpr dflags to_drop lam@(_, AnnLam _ _)
-  | okToFloatInside bndrs       -- Float in
+  | noFloatIntoLam bndrs       -- Dump it all here
      -- NB: Must line up with noFloatIntoRhs (AnnLam...); see Trac #7088
-  = mkLams bndrs (fiExpr dflags to_drop body)
-
-  | otherwise           -- Dump it all here
   = wrapFloats to_drop (mkLams bndrs (fiExpr dflags [] body))
 
+  | otherwise           -- Float inside
+  = mkLams bndrs (fiExpr dflags to_drop body)
+
   where
     (bndrs, body) = collectAnnBndrs lam
 
@@ -342,42 +374,69 @@ fiExpr dflags to_drop (_,AnnLet bind body)
   = fiExpr dflags (after ++ new_float : before) body
            -- to_drop is in reverse dependency order
   where
-    (before, new_float, after) = fiBind dflags to_drop bind body_fvs body_ty_fvs
+    (before, new_float, after) = fiBind dflags to_drop bind body_fvs
     body_fvs    = freeVarsOf body
-    body_ty_fvs = freeVarsOfType body
 
-{-
-For @Case@, the possible ``drop points'' for the \tr{to_drop}
-bindings are: (a)~inside the scrutinee, (b)~inside one of the
-alternatives/default [default FVs always {\em first}!].
-
-Floating case expressions inward was added to fix Trac #5658: strict bindings
-not floated in. In particular, this change allows array indexing operations,
-which have a single DEFAULT alternative without any binders, to be floated
-inward. SIMD primops for unpacking SIMD vectors into an unboxed tuple of unboxed
+{- Note [Floating primops]
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+We try to float-in a case expression over an unlifted type.  The
+motivating example was Trac #5658: in particular, this change allows
+array indexing operations, which have a single DEFAULT alternative
+without any binders, to be floated inward.
+
+SIMD primops for unpacking SIMD vectors into an unboxed tuple of unboxed
 scalars also need to be floated inward, but unpacks have a single non-DEFAULT
 alternative that binds the elements of the tuple. We now therefore also support
 floating in cases with a single alternative that may bind values.
+
+But there are wrinkles
+
+* Which unlifted cases do we float? See PrimOp.hs
+  Note [PrimOp can_fail and has_side_effects] which explains:
+   - We can float-in can_fail primops, but we can't float them out.
+   - But we can float a has_side_effects primop, but NOT inside a lambda,
+     so for now we don't float them at all.
+  Hence exprOkForSideEffects
+
+* Because we can float can-fail primops (array indexing, division) inwards
+  but not outwards, we must be careful not to transform
+     case a /# b of r -> f (F# r)
+  ===>
+    f (case a /# b of r -> F# r)
+  because that creates a new thunk that wasn't there before.  And
+  because it can't be floated out (can_fail), the thunk will stay
+  there.  Disaster!  (This happened in nofib 'simple' and 'scs'.)
+
+  Solution: only float cases into the branches of other cases, and
+  not into the arguments of an application, or the RHS of a let. This
+  is somewhat conservative, but it's simple.  And it stil hits the
+  cases like Trac #5658.   This is implemented in sepBindsByJoinPoint;
+  if is_case is False we dump all floating cases right here.
+
+For @Case@, the possible drop points for the 'to_drop'
+bindings are:
+  (a) inside the scrutinee
+  (b) inside one of the alternatives/default (default FVs always /first/!).
+
 -}
 
 fiExpr dflags to_drop (_, AnnCase scrut case_bndr _ [(con,alt_bndrs,rhs)])
   | isUnliftedType (idType case_bndr)
   , exprOkForSideEffects (deAnnotate scrut)
-      -- See PrimOp, Note [PrimOp can_fail and has_side_effects]
+      -- See Note [Floating primops]
   = wrapFloats shared_binds $
     fiExpr dflags (case_float : rhs_binds) rhs
   where
     case_float = FB (mkDVarSet (case_bndr : alt_bndrs)) scrut_fvs
                     (FloatCase scrut' case_bndr con alt_bndrs)
-    scrut' = fiExpr dflags scrut_binds scrut
+    scrut'     = fiExpr dflags scrut_binds scrut
+    rhs_fvs    = freeVarsOf rhs `delDVarSetList` (case_bndr : alt_bndrs)
+    scrut_fvs  = freeVarsOf scrut
+
     [shared_binds, scrut_binds, rhs_binds]
        = sepBindsByDropPoint dflags False
            [scrut_fvs, rhs_fvs]
-           (freeVarsOfType scrut `unionDVarSet` rhs_ty_fvs)
            to_drop
-    rhs_fvs    = freeVarsOf rhs `delDVarSetList` (case_bndr : alt_bndrs)
-    rhs_ty_fvs = freeVarsOfType rhs `delDVarSetList` (case_bndr : alt_bndrs)
-    scrut_fvs  = freeVarsOf scrut
 
 fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
   = wrapFloats drop_here1 $
@@ -389,25 +448,20 @@ fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
     [drop_here1, scrut_drops, alts_drops]
        = sepBindsByDropPoint dflags False
            [scrut_fvs, all_alts_fvs]
-           (freeVarsOfType scrut `unionDVarSet` all_alts_ty_fvs)
            to_drop
 
         -- Float into the alts with the is_case flag set
     (drop_here2 : alts_drops_s)
-      = sepBindsByDropPoint dflags True alts_fvs all_alts_ty_fvs
-                            alts_drops
-
-    scrut_fvs       = freeVarsOf scrut
-    alts_fvs        = map alt_fvs alts
-    all_alts_fvs    = unionDVarSets alts_fvs
-    alts_ty_fvs     = map alt_ty_fvs alts
-    all_alts_ty_fvs = unionDVarSets alts_ty_fvs
+      | [ _ ] <- alts = [] : [alts_drops]
+      | otherwise     = sepBindsByDropPoint dflags True alts_fvs alts_drops
+
+    scrut_fvs    = freeVarsOf scrut
+    alts_fvs     = map alt_fvs alts
+    all_alts_fvs = unionDVarSets alts_fvs
     alt_fvs (_con, args, rhs)
-      = foldl delDVarSet (freeVarsOf rhs)     (case_bndr:args)
-    alt_ty_fvs (_con, args, rhs)
-      = foldl delDVarSet (freeVarsOfType rhs) (case_bndr:args)
-                                -- Delete case_bndr and args from free vars of rhs
-                                -- to get free vars of alt
+      = foldl delDVarSet (freeVarsOf rhs) (case_bndr:args)
+           -- Delete case_bndr and args from free vars of rhs
+           -- to get free vars of alt
 
     fi_alt to_drop (con, args, rhs) = (con, args, fiExpr dflags to_drop rhs)
 
@@ -417,12 +471,11 @@ fiBind :: DynFlags
                             -- as far "inwards" as possible
        -> CoreBindWithFVs   -- Input binding
        -> DVarSet           -- Free in scope of binding
-       -> DVarSet           -- Free in type of body of binding
        -> ( FloatInBinds    -- Land these before
           , FloatInBind     -- The binding itself
           , FloatInBinds)   -- Land these after
 
-fiBind dflags to_drop (AnnNonRec id rhs) body_fvs body_ty_fvs
+fiBind dflags to_drop (AnnNonRec id ann_rhs@(rhs_fvs, rhs)) body_fvs
   = ( extra_binds ++ shared_binds          -- Land these before
                                            -- See Note [extra_fvs (1,2)]
     , FB (unitDVarSet id) rhs_fvs'         -- The new binding itself
@@ -431,11 +484,10 @@ fiBind dflags to_drop (AnnNonRec id rhs) body_fvs body_ty_fvs
 
   where
     body_fvs2 = body_fvs `delDVarSet` id
-    rhs_fvs   = freeVarsOf rhs
 
     rule_fvs = idRuleAndUnfoldingVarsDSet id        -- See Note [extra_fvs (2): free variables of rules]
-    extra_fvs | noFloatIntoRhs (isJoinId id) NonRecursive rhs
-              = rule_fvs `unionDVarSet` freeVarsOf rhs
+    extra_fvs | noFloatIntoRhs NonRecursive id rhs
+              = rule_fvs `unionDVarSet` rhs_fvs
               | otherwise
               = rule_fvs
         -- See Note [extra_fvs (1): avoid floating into RHS]
@@ -446,15 +498,14 @@ fiBind dflags to_drop (AnnNonRec id rhs) body_fvs body_ty_fvs
     [shared_binds, extra_binds, rhs_binds, body_binds]
         = sepBindsByDropPoint dflags False
             [extra_fvs, rhs_fvs, body_fvs2]
-            (freeVarsOfType rhs `unionDVarSet` body_ty_fvs)
             to_drop
 
         -- Push rhs_binds into the right hand side of the binding
-    rhs'     = fiRhs dflags rhs_binds id rhs
+    rhs'     = fiRhs dflags rhs_binds id ann_rhs
     rhs_fvs' = rhs_fvs `unionDVarSet` floatedBindsFVs rhs_binds `unionDVarSet` rule_fvs
                         -- Don't forget the rule_fvs; the binding mentions them!
 
-fiBind dflags to_drop (AnnRec bindings) body_fvs body_ty_fvs
+fiBind dflags to_drop (AnnRec bindings) body_fvs
   = ( extra_binds ++ shared_binds
     , FB (mkDVarSet ids) rhs_fvs'
          (FloatLet (Rec (fi_bind rhss_binds bindings)))
@@ -466,13 +517,12 @@ fiBind dflags to_drop (AnnRec bindings) body_fvs body_ty_fvs
         -- See Note [extra_fvs (1,2)]
     rule_fvs = mapUnionDVarSet idRuleAndUnfoldingVarsDSet ids
     extra_fvs = rule_fvs `unionDVarSet`
-                unionDVarSets [ freeVarsOf rhs | (bndr, rhs) <- bindings
-                              , noFloatIntoRhs (isJoinId bndr) Recursive rhs ]
+                unionDVarSets [ rhs_fvs | (bndr, (rhs_fvs, rhs)) <- bindings
+                              , noFloatIntoRhs Recursive bndr rhs ]
 
     (shared_binds:extra_binds:body_binds:rhss_binds)
         = sepBindsByDropPoint dflags False
             (extra_fvs:body_fvs:rhss_fvs)
-            (body_ty_fvs `unionDVarSet` mapUnionDVarSet freeVarsOfType rhss)
             to_drop
 
     rhs_fvs' = unionDVarSets rhss_fvs `unionDVarSet`
@@ -498,45 +548,75 @@ fiRhs dflags to_drop bndr rhs
   = fiExpr dflags to_drop rhs
 
 ------------------
-okToFloatInside :: [Var] -> Bool
-okToFloatInside bndrs = all ok bndrs
+noFloatIntoLam :: [Var] -> Bool
+noFloatIntoLam bndrs = any bad bndrs
   where
-    ok b = not (isId b) || isOneShotBndr b
-    -- Push the floats inside there are no non-one-shot value binders
+    bad b = isId b && not (isOneShotBndr b)
+    -- Don't float inside a non-one-shot lambda
 
-noFloatIntoRhs :: Bool -> RecFlag -> CoreExprWithFVs -> Bool
+noFloatIntoRhs :: RecFlag -> Id -> CoreExprWithFVs' -> Bool
 -- ^ True if it's a bad idea to float bindings into this RHS
--- Preconditio:  rhs :: rhs_ty
-noFloatIntoRhs is_join is_rec rhs@(_, rhs')
-  |  is_join
-  =  isRec is_rec -- Joins are one-shot iff non-recursive
-  |  otherwise
-  =  isUnliftedType rhs_ty
-       -- See Note [Do not destroy the let/app invariant]
-  || noFloatIntoExpr rhs'
+noFloatIntoRhs is_rec bndr rhs
+  | isJoinId bndr
+  = isRec is_rec -- Joins are one-shot iff non-recursive
+
+  | otherwise
+  = noFloatIntoArg rhs (idType bndr)
+
+noFloatIntoArg :: CoreExprWithFVs' -> Type -> Bool
+noFloatIntoArg expr expr_ty
+  | isUnliftedType expr_ty
+  = True  -- See Note [Do not destroy the let/app invariant]
+
+   | AnnLam bndr e <- expr
+   , (bndrs, _) <- collectAnnBndrs e
+   =  noFloatIntoLam (bndr:bndrs)  -- Wrinkle 1 (a)
+   || all isTyVar (bndr:bndrs)     -- Wrinkle 1 (b)
+      -- See Note [noFloatInto considerations] wrinkle 2
+
+  | otherwise  -- Note [noFloatInto considerations] wrinkle 2
+  = exprIsTrivial deann_expr || exprIsHNF deann_expr
   where
-    rhs_ty = exprTypeFV rhs
-
-noFloatIntoExpr :: CoreExprWithFVs' -> Bool
-noFloatIntoExpr (AnnLam bndr e)
-   = not (okToFloatInside (bndr:bndrs))
-     -- NB: Must line up with fiExpr (AnnLam...); see Trac #7088
-   where
-     (bndrs, _) = collectAnnBndrs e
-        -- IMPORTANT: don't say 'True' for a RHS with a one-shot lambda at the top.
-        -- This makes a big difference for things like
-        --      f x# = let x = I# x#
-        --             in let j = \() -> ...x...
-        --                in if <condition> then normal-path else j ()
-        -- If x is used only in the error case join point, j, we must float the
-        -- boxing constructor into it, else we box it every time which is very bad
-        -- news indeed.
-
-noFloatIntoExpr rhs = exprIsExpandable (deAnnotate' rhs)
-       -- We'd just float right back out again...
-       -- Should match the test in SimplEnv.doFloatFromRhs
+    deann_expr = deAnnotate' expr
+
+{- Note [noFloatInto considerations]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When do we want to float bindings into
+   - noFloatIntoRHs: the RHS of a let-binding
+   - noFloatIntoArg: the argument of a function application
+
+Definitely don't float in if it has unlifted type; that
+would destroy the let/app invariant.
+
+* Wrinkle 1: do not float in if
+     (a) any non-one-shot value lambdas
+  or (b) all type lambdas
+  In both cases we'll float straight back out again
+  NB: Must line up with fiExpr (AnnLam...); see Trac #7088
+
+  (a) is important: we /must/ float into a one-shot lambda group
+  (which includes join points). This makes a big difference
+  for things like
+     f x# = let x = I# x#
+            in let j = \() -> ...x...
+               in if <condition> then normal-path else j ()
+  If x is used only in the error case join point, j, we must float the
+  boxing constructor into it, else we box it every time which is very
+  bad news indeed.
+
+* Wrinkle 2: for RHSs, do not float into a HNF; we'll just float right
+  back out again... not tragic, but a waste of time.
+
+  For function arguments we will still end up with this
+  in-then-out stuff; consider
+    letrec x = e in f x
+  Here x is not a HNF, so we'll produce
+    f (letrec x = e in x)
+  which is OK... it's not that common, and we'll end up
+  floating out again, in CorePrep if not earlier.
+  Still, we use exprIsTrivial to catch this case (sigh)
+
 
-{-
 ************************************************************************
 *                                                                      *
 \subsection{@sepBindsByDropPoint@}
@@ -556,27 +636,21 @@ in it goes.  If a binding is used inside {\em multiple} drop points,
 then it has to go in a you-must-drop-it-above-all-these-drop-points
 point.
 
-But, with coercions appearing in types, there is a complication: we
-might be floating in a "strict let" -- that is, a case. Case expressions
-mention their return type. We absolutely can't float a coercion binding
-inward to the point that the type of the expression it's about to wrap
-mentions the coercion. So we include the union of the sets of free variables
-of the types of all the drop points involved. If any of the floaters
-bind a coercion variable mentioned in any of the types, that binder must
-be dropped right away.
-
 We have to maintain the order on these drop-point-related lists.
 -}
 
+-- pprFIB :: FloatInBinds -> SDoc
+-- pprFIB fibs = text "FIB:" <+> ppr [b | FB _ _ b <- fibs]
+
 sepBindsByDropPoint
     :: DynFlags
-    -> Bool             -- True <=> is case expression
-    -> [FreeVarSet]         -- One set of FVs per drop point
-    -> FreeVarSet           -- Vars free in all the types of the drop points
-    -> FloatInBinds         -- Candidate floaters
+    -> Bool                -- True <=> is case expression
+    -> [FreeVarSet]        -- One set of FVs per drop point
+                           -- Always at least two long!
+    -> FloatInBinds        -- Candidate floaters
     -> [FloatInBinds]      -- FIRST one is bindings which must not be floated
-                            -- inside any drop point; the rest correspond
-                            -- one-to-one with the input list of FV sets
+                           -- inside any drop point; the rest correspond
+                           -- one-to-one with the input list of FV sets
 
 -- Every input floater is returned somewhere in the result;
 -- none are dropped, not even ones which don't seem to be
@@ -586,12 +660,16 @@ sepBindsByDropPoint
 
 type DropBox = (FreeVarSet, FloatInBinds)
 
-sepBindsByDropPoint _ _is_case drop_pts _ty_fvs []
-  = [] : [[] | _ <- drop_pts]   -- cut to the chase scene; it happens
+sepBindsByDropPoint dflags is_case drop_pts floaters
+  | null floaters  -- Shortcut common case
+  = [] : [[] | _ <- drop_pts]
 
-sepBindsByDropPoint dflags is_case drop_pts ty_fvs floaters
-  = go floaters (map (\fvs -> (fvs, [])) (emptyDVarSet : drop_pts))
+  | otherwise
+  = ASSERT( length drop_pts >= 2 )
+    go floaters (map (\fvs -> (fvs, [])) (emptyDVarSet : drop_pts))
   where
+    n_alts = length drop_pts
+
     go :: FloatInBinds -> [DropBox] -> [FloatInBinds]
         -- The *first* one in the argument list is the drop_here set
         -- The FloatInBinds in the lists are in the reverse of
@@ -606,32 +684,25 @@ sepBindsByDropPoint dflags is_case drop_pts ty_fvs floaters
 
           (used_here : used_in_flags) = [ fvs `intersectsDVarSet` bndrs
                                         | (fvs, _) <- drop_boxes]
-          used_in_ty = ty_fvs `intersectsDVarSet` bndrs
 
-          drop_here = used_here || not can_push || used_in_ty
+          drop_here = used_here || cant_push
 
-                -- For case expressions we duplicate the binding if it is
-                -- reasonably small, and if it is not used in all the RHSs
-                -- This is good for situations like
-                --      let x = I# y in
-                --      case e of
-                --        C -> error x
-                --        D -> error x
-                --        E -> ...not mentioning x...
-
-          n_alts      = length used_in_flags
           n_used_alts = count id used_in_flags -- returns number of Trues in list.
 
-          can_push = n_used_alts == 1           -- Used in just one branch
-                   || (is_case &&               -- We are looking at case alternatives
-                       n_used_alts > 1 &&       -- It's used in more than one
-                       n_used_alts < n_alts &&  -- ...but not all
-                       floatIsDupable dflags bind) -- and we can duplicate the binding
+          cant_push
+            | is_case   = n_used_alts == n_alts   -- Used in all, don't push
+                                                  -- Remember n_alts > 1
+                          || (n_used_alts > 1 && not (floatIsDupable dflags bind))
+                             -- floatIsDupable: see Note [Duplicating floats]
+
+            | otherwise = floatIsCase bind || n_used_alts > 1
+                             -- floatIsCase: see Note [Floating primops]
 
           new_boxes | drop_here = (insert here_box : fork_boxes)
                     | otherwise = (here_box : new_fork_boxes)
 
-          new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe fork_boxes used_in_flags
+          new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe
+                                        fork_boxes used_in_flags
 
           insert :: DropBox -> DropBox
           insert (fvs,drops) = (fvs `unionDVarSet` bind_fvs, bind_w_fvs:drops)
@@ -642,6 +713,22 @@ sepBindsByDropPoint dflags is_case drop_pts ty_fvs floaters
     go _ _ = panic "sepBindsByDropPoint/go"
 
 
+{- Note [Duplicating floats]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For case expressions we duplicate the binding if it is reasonably
+small, and if it is not used in all the RHSs This is good for
+situations like
+     let x = I# y in
+     case e of
+       C -> error x
+       D -> error x
+       E -> ...not mentioning x...
+
+If the thing is used in all RHSs there is nothing gained,
+so we don't duplicate then.
+-}
+
 floatedBindsFVs :: FloatInBinds -> FreeVarSet
 floatedBindsFVs binds = mapUnionDVarSet fbFVs binds
 
@@ -657,3 +744,7 @@ floatIsDupable :: DynFlags -> FloatBind -> Bool
 floatIsDupable dflags (FloatCase scrut _ _ _) = exprIsDupable dflags scrut
 floatIsDupable dflags (FloatLet (Rec prs))    = all (exprIsDupable dflags . snd) prs
 floatIsDupable dflags (FloatLet (NonRec _ r)) = exprIsDupable dflags r
+
+floatIsCase :: FloatBind -> Bool
+floatIsCase (FloatCase {}) = True
+floatIsCase (FloatLet {})  = False