Don't invoke dataConSrcToImplBang on newtypes
[ghc.git] / compiler / simplCore / FloatIn.hs
index c1147eb..e8c7ef2 100644 (file)
@@ -13,39 +13,46 @@ then discover that they aren't needed in the chosen branch.
 -}
 
 {-# LANGUAGE CPP #-}
+{-# OPTIONS_GHC -fprof-auto #-}
 
 module FloatIn ( floatInwards ) where
 
 #include "HsVersions.h"
 
+import GhcPrelude
+
 import CoreSyn
 import MkCore
-import CoreUtils        ( exprIsDupable, exprIsExpandable, exprType,
-                          exprOkForSideEffects, mkTicks )
-import CoreFVs          ( CoreExprWithFVs, freeVars, freeVarsOf, idRuleAndUnfoldingVars )
-import Id               ( isOneShotBndr, idType )
+import HscTypes         ( ModGuts(..) )
+import CoreUtils
+import CoreFVs
+import CoreMonad        ( CoreM )
+import Id               ( isOneShotBndr, idType, isJoinId, isJoinId_maybe )
 import Var
-import Type             ( Type, isUnLiftedType, isFunTy, splitFunTy, applyTy )
+import Type
 import VarSet
 import Util
-import UniqDFM (UniqDFM, udfmToUfm)
 import DynFlags
 import Outputable
-import Data.List( mapAccumL )
+-- import Data.List        ( mapAccumL )
+import BasicTypes       ( RecFlag(..), isRec )
 
 {-
 Top-level interface function, @floatInwards@.  Note that we do not
 actually float any bindings downwards from the top-level.
 -}
 
-floatInwards :: DynFlags -> CoreProgram -> CoreProgram
-floatInwards dflags = map fi_top_bind
+floatInwards :: ModGuts -> CoreM ModGuts
+floatInwards pgm@(ModGuts { mg_binds = binds })
+  = do { dflags <- getDynFlags
+       ; return (pgm { mg_binds = map (fi_top_bind dflags) binds }) }
   where
-    fi_top_bind (NonRec binder rhs)
+    fi_top_bind dflags (NonRec binder rhs)
       = NonRec binder (fiExpr dflags [] (freeVars rhs))
-    fi_top_bind (Rec pairs)
+    fi_top_bind dflags (Rec pairs)
       = Rec [ (b, fiExpr dflags [] (freeVars rhs)) | (b, rhs) <- pairs ]
 
+
 {-
 ************************************************************************
 *                                                                      *
@@ -118,8 +125,8 @@ the closure for a is not built.
 ************************************************************************
 -}
 
-type FreeVarSet  = IdSet
-type BoundVarSet = IdSet
+type FreeVarSet  = DIdSet
+type BoundVarSet = DIdSet
 
 data FloatInBind = FB BoundVarSet FreeVarSet FloatBind
         -- The FreeVarSet is the free variables of the binding.  In the case
@@ -135,15 +142,19 @@ fiExpr :: DynFlags
        -> CoreExprWithFVs   -- Input expr
        -> CoreExpr          -- Result
 
-fiExpr _ to_drop (_, AnnLit lit)     = ASSERT( null to_drop ) Lit lit
+fiExpr _ to_drop (_, AnnLit lit)     = wrapFloats to_drop (Lit lit)
+                                       -- See Note [Dead bindings]
 fiExpr _ to_drop (_, AnnType ty)     = ASSERT( null to_drop ) Type ty
 fiExpr _ to_drop (_, AnnVar v)       = wrapFloats to_drop (Var v)
 fiExpr _ to_drop (_, AnnCoercion co) = wrapFloats to_drop (Coercion co)
-fiExpr dflags to_drop (_, AnnCast expr (fvs_co, co))
+fiExpr dflags to_drop (_, AnnCast expr (co_ann, co))
   = wrapFloats (drop_here ++ co_drop) $
     Cast (fiExpr dflags e_drop expr) co
   where
-    [drop_here, e_drop, co_drop] = sepBindsByDropPoint dflags False [udfmToUfm $ freeVarsOf expr, udfmToUfm fvs_co] to_drop
+    [drop_here, e_drop, co_drop]
+      = sepBindsByDropPoint dflags False
+          [freeVarsOf expr, freeVarsOfAnn co_ann]
+          to_drop
 
 {-
 Applications: we do float inside applications, mainly because we
@@ -152,34 +163,55 @@ 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@(fun_fvs, _), ann_args, ticks)
-           = collectAnnArgsTicks tickishFloatable ann_expr
-    fun_ty = exprType (deAnnotate ann_fun)
-    ((_,extra_fvs), arg_fvs) = mapAccumL mk_arg_fvs (fun_ty, emptyVarSet) ann_args
-
-    -- All this faffing about is so that we can get hold of
-    -- the types of the arguments, to pass to noFloatIntoRhs
-    mk_arg_fvs :: (Type, FreeVarSet) -> CoreExprWithFVs -> ((Type, FreeVarSet), FreeVarSet)
-    mk_arg_fvs (fun_ty, extra_fvs) (_, AnnType ty)
-      = ((applyTy fun_ty ty, extra_fvs), emptyVarSet)
-
-    mk_arg_fvs (fun_ty, extra_fvs) (arg_dfvs, ann_arg)
-      | ASSERT( isFunTy fun_ty ) noFloatIntoRhs ann_arg arg_ty
-      = ((res_ty, extra_fvs `unionVarSet` arg_fvs), emptyVarSet)
+    (ann_fun, ann_args, ticks) = collectAnnArgsTicks tickishFloatable ann_expr
+    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.)
+
+    add_arg :: (Type,FreeVarSet) -> CoreExprWithFVs -> (Type,FreeVarSet)
+    add_arg (fun_ty, extra_fvs) (_, AnnType ty)
+      = (piResultTy fun_ty ty, extra_fvs)
+
+    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), arg_fvs)
+      = (res_ty, extra_fvs)
       where
-       arg_fvs = udfmToUfm arg_dfvs
        (arg_ty, res_ty) = splitFunTy fun_ty
 
-    drop_here : extra_drop : fun_drop : arg_drops
-      = sepBindsByDropPoint dflags False (extra_fvs : udfmToUfm fun_fvs : arg_fvs) to_drop
+{- Note [Dead bindings]
+~~~~~~~~~~~~~~~~~~~~~~~
+At a literal we won't usually have any floated bindings; the
+only way that can happen is if the binding wrapped the literal
+/in the original input program/.  e.g.
+   case x of { DEFAULT -> 1# }
+But, while this may be unusual it is not actually wrong, and it did
+once happen (Trac #15696).
 
-{-
 Note [Do not destroy the let/app invariant]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Watch out for
@@ -189,6 +221,27 @@ We don't want to float bindings into here
 because that might destroy the let/app invariant, which requires
 unlifted function arguments to be ok-for-speculation.
 
+Note [Join points]
+~~~~~~~~~~~~~~~~~~
+Generally, we don't need to worry about join points - there are places we're
+not allowed to float them, but since they can't have occurrences in those
+places, we're not tempted.
+
+We do need to be careful about jumps, however:
+
+  joinrec j x y z = ... in
+  jump j a b c
+
+Previous versions often floated the definition of a recursive function into its
+only non-recursive occurrence. But for a join point, this is a disaster:
+
+  (joinrec j x y z = ... in
+  jump j) a b c -- wrong!
+
+Every jump must be exact, so the jump to j must have three arguments. Hence
+we're careful not to float into the target of a jump (though we can float into
+the arguments just fine).
+
 Note [Floating in past a lambda group]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 * We must be careful about floating inside a value lambda.
@@ -224,18 +277,42 @@ So we treat lambda in groups, using the following rule:
 
 This is what the 'go' function in the AnnLam case is doing.
 
+(Join points are handled similarly: a join point is considered one-shot iff
+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
 
@@ -304,97 +381,83 @@ idRuleAndUnfoldingVars of x.  No need for type variables, hence not using
 idFreeVars.
 -}
 
-fiExpr dflags to_drop (_,AnnLet (AnnNonRec id rhs@(rhs_dfvs, ann_rhs)) body)
-  = fiExpr dflags new_to_drop body
+fiExpr dflags to_drop (_,AnnLet bind body)
+  = fiExpr dflags (after ++ new_float : before) body
+           -- to_drop is in reverse dependency order
   where
-    body_fvs = udfmToUfm (freeVarsOf body) `delVarSet` id
-    rhs_ty = idType id
-    rhs_fvs = udfmToUfm rhs_dfvs
-    rule_fvs = idRuleAndUnfoldingVars id        -- See Note [extra_fvs (2): free variables of rules]
-    extra_fvs | noFloatIntoRhs ann_rhs rhs_ty = rule_fvs `unionVarSet` rhs_fvs
-              | otherwise                     = rule_fvs
-        -- See Note [extra_fvs (1): avoid floating into RHS]
-        -- No point in floating in only to float straight out again
-        -- Ditto ok-for-speculation unlifted RHSs
+    (before, new_float, after) = fiBind dflags to_drop bind body_fvs
+    body_fvs    = freeVarsOf body
 
-    [shared_binds, extra_binds, rhs_binds, body_binds]
-        = sepBindsByDropPoint dflags False [extra_fvs, rhs_fvs, body_fvs] to_drop
-
-    new_to_drop = body_binds ++                         -- the bindings used only in the body
-                  [FB (unitVarSet id) rhs_fvs'
-                      (FloatLet (NonRec id rhs'))] ++   -- the new binding itself
-                  extra_binds ++                        -- bindings from extra_fvs
-                  shared_binds                          -- the bindings used both in rhs and body
+{- 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.
 
-        -- Push rhs_binds into the right hand side of the binding
-    rhs'     = fiExpr dflags rhs_binds rhs
-    rhs_fvs' = rhs_fvs `unionVarSet` floatedBindsFVs rhs_binds `unionVarSet` rule_fvs
-                        -- Don't forget the rule_fvs; the binding mentions them!
-
-fiExpr dflags to_drop (_,AnnLet (AnnRec bindings) body)
-  = fiExpr dflags new_to_drop body
-  where
-    (ids, rhss) = unzip bindings
-    rhss_fvs = map (udfmToUfm . freeVarsOf) rhss
-    body_fvs = udfmToUfm $ freeVarsOf body
-
-        -- See Note [extra_fvs (1,2)]
-    rule_fvs = mapUnionVarSet idRuleAndUnfoldingVars ids
-    extra_fvs = rule_fvs `unionVarSet`
-                unionVarSets [ udfmToUfm fvs | (fvs, rhs) <- rhss
-                             , noFloatIntoExpr rhs ]
-
-    (shared_binds:extra_binds:body_binds:rhss_binds)
-        = sepBindsByDropPoint dflags False (extra_fvs:body_fvs:rhss_fvs) to_drop
-
-    new_to_drop = body_binds ++         -- the bindings used only in the body
-                  [FB (mkVarSet ids) rhs_fvs'
-                      (FloatLet (Rec (fi_bind rhss_binds bindings)))] ++
-                                        -- The new binding itself
-                  extra_binds ++        -- Note [extra_fvs (1,2)]
-                  shared_binds          -- Used in more than one place
-
-    rhs_fvs' = unionVarSets rhss_fvs `unionVarSet`
-               unionVarSets (map floatedBindsFVs rhss_binds) `unionVarSet`
-               rule_fvs         -- Don't forget the rule variables!
-
-    -- Push rhs_binds into the right hand side of the binding
-    fi_bind :: [FloatInBinds]       -- one per "drop pt" conjured w/ fvs_of_rhss
-            -> [(Id, CoreExprWithFVs)]
-            -> [(Id, CoreExpr)]
-
-    fi_bind to_drops pairs
-      = [ (binder, fiExpr dflags to_drop rhs)
-        | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
-
-{-
-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
+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 still hits the
+  cases like Trac #5658.   This is implemented in sepBindsByJoinPoint;
+  if is_case is False we dump all floating cases right here.
+
+* Trac #14511 is another example of why we want to restrict float-in
+  of case-expressions.  Consider
+     case indexArray# a n of (# r #) -> writeArray# ma i (f r)
+  Now, floating that indexing operation into the (f r) thunk will
+  not create any new thunks, but it will keep the array 'a' alive
+  for much longer than the programmer expected.
+
+  So again, not floating a case into a let or argument seems like
+  the Right Thing
+
+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)
+  | 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 (mkVarSet (case_bndr : alt_bndrs)) scrut_fvs
+    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] to_drop
-    rhs_fvs   = udfmToUfm (freeVarsOf rhs) `delVarSetList` (case_bndr : alt_bndrs)
-    scrut_fvs = udfmToUfm $ freeVarsOf scrut
+       = sepBindsByDropPoint dflags False
+           [scrut_fvs, rhs_fvs]
+           to_drop
 
 fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
   = wrapFloats drop_here1 $
@@ -404,53 +467,177 @@ fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
   where
         -- Float into the scrut and alts-considered-together just like App
     [drop_here1, scrut_drops, alts_drops]
-       = sepBindsByDropPoint dflags False [scrut_fvs, all_alts_fvs] to_drop
+       = sepBindsByDropPoint dflags False
+           [scrut_fvs, all_alts_fvs]
+           to_drop
 
         -- Float into the alts with the is_case flag set
-    (drop_here2 : alts_drops_s) = sepBindsByDropPoint dflags True alts_fvs alts_drops
+    (drop_here2 : alts_drops_s)
+      | [ _ ] <- alts = [] : [alts_drops]
+      | otherwise     = sepBindsByDropPoint dflags True alts_fvs alts_drops
 
-    scrut_fvs    = udfmToUfm $ freeVarsOf scrut
+    scrut_fvs    = freeVarsOf scrut
     alts_fvs     = map alt_fvs alts
-    all_alts_fvs = unionVarSets alts_fvs
-    alt_fvs (_con, args, rhs) = foldl delVarSet (udfmToUfm $ freeVarsOf rhs) (case_bndr:args)
-                                -- Delete case_bndr and args from free vars of rhs
-                                -- to get free vars of alt
+    all_alts_fvs = unionDVarSets alts_fvs
+    alt_fvs (_con, args, rhs)
+      = 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)
 
-okToFloatInside :: [Var] -> Bool
-okToFloatInside bndrs = all ok bndrs
+------------------
+fiBind :: DynFlags
+       -> FloatInBinds      -- Binds we're trying to drop
+                            -- as far "inwards" as possible
+       -> CoreBindWithFVs   -- Input binding
+       -> DVarSet           -- Free in scope of binding
+       -> ( FloatInBinds    -- Land these before
+          , FloatInBind     -- The binding itself
+          , FloatInBinds)   -- Land these after
+
+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
+          (FloatLet (NonRec id rhs'))
+    , body_binds )                         -- Land these after
+
+  where
+    body_fvs2 = body_fvs `delDVarSet` id
+
+    rule_fvs = bndrRuleAndUnfoldingVarsDSet id        -- See Note [extra_fvs (2): free variables of rules]
+    extra_fvs | noFloatIntoRhs NonRecursive id rhs
+              = rule_fvs `unionDVarSet` rhs_fvs
+              | otherwise
+              = rule_fvs
+        -- See Note [extra_fvs (1): avoid floating into RHS]
+        -- No point in floating in only to float straight out again
+        -- We *can't* float into ok-for-speculation unlifted RHSs
+        -- But do float into join points
+
+    [shared_binds, extra_binds, rhs_binds, body_binds]
+        = sepBindsByDropPoint dflags False
+            [extra_fvs, rhs_fvs, body_fvs2]
+            to_drop
+
+        -- Push rhs_binds into the right hand side of the binding
+    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
+  = ( extra_binds ++ shared_binds
+    , FB (mkDVarSet ids) rhs_fvs'
+         (FloatLet (Rec (fi_bind rhss_binds bindings)))
+    , body_binds )
   where
-    ok b = not (isId b) || isOneShotBndr b
-    -- Push the floats inside there are no non-one-shot value binders
+    (ids, rhss) = unzip bindings
+    rhss_fvs = map freeVarsOf rhss
+
+        -- See Note [extra_fvs (1,2)]
+    rule_fvs = mapUnionDVarSet bndrRuleAndUnfoldingVarsDSet ids
+    extra_fvs = rule_fvs `unionDVarSet`
+                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)
+            to_drop
+
+    rhs_fvs' = unionDVarSets rhss_fvs `unionDVarSet`
+               unionDVarSets (map floatedBindsFVs rhss_binds) `unionDVarSet`
+               rule_fvs         -- Don't forget the rule variables!
 
-noFloatIntoRhs :: AnnExpr' Var (UniqDFM Var) -> Type -> Bool
+    -- Push rhs_binds into the right hand side of the binding
+    fi_bind :: [FloatInBinds]       -- one per "drop pt" conjured w/ fvs_of_rhss
+            -> [(Id, CoreExprWithFVs)]
+            -> [(Id, CoreExpr)]
+
+    fi_bind to_drops pairs
+      = [ (binder, fiRhs dflags to_drop binder rhs)
+        | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
+
+------------------
+fiRhs :: DynFlags -> FloatInBinds -> CoreBndr -> CoreExprWithFVs -> CoreExpr
+fiRhs dflags to_drop bndr rhs
+  | Just join_arity <- isJoinId_maybe bndr
+  , let (bndrs, body) = collectNAnnBndrs join_arity rhs
+  = mkLams bndrs (fiExpr dflags to_drop body)
+  | otherwise
+  = fiExpr dflags to_drop rhs
+
+------------------
+noFloatIntoLam :: [Var] -> Bool
+noFloatIntoLam bndrs = any bad bndrs
+  where
+    bad b = isId b && not (isOneShotBndr b)
+    -- Don't float inside a non-one-shot lambda
+
+noFloatIntoRhs :: RecFlag -> Id -> CoreExprWithFVs' -> Bool
 -- ^ True if it's a bad idea to float bindings into this RHS
--- Preconditio:  rhs :: rhs_ty
-noFloatIntoRhs rhs rhs_ty
-  =  isUnLiftedType rhs_ty   -- See Note [Do not destroy the let/app invariant]
-  || noFloatIntoExpr rhs
-
-noFloatIntoExpr :: AnnExpr' Var (UniqDFM Var) -> 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
+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
+    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@}
@@ -473,14 +660,18 @@ point.
 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
-    -> 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
@@ -490,12 +681,16 @@ sepBindsByDropPoint
 
 type DropBox = (FreeVarSet, FloatInBinds)
 
-sepBindsByDropPoint _ _is_case drop_pts []
-  = [] : [[] | _ <- drop_pts]   -- cut to the chase scene; it happens
-
 sepBindsByDropPoint dflags is_case drop_pts floaters
-  = go floaters (map (\fvs -> (fvs, [])) (emptyVarSet : drop_pts))
+  | null floaters  -- Shortcut common case
+  = [] : [[] | _ <- drop_pts]
+
+  | otherwise
+  = ASSERT( drop_pts `lengthAtLeast` 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
@@ -508,36 +703,30 @@ sepBindsByDropPoint dflags is_case drop_pts floaters
         where
           -- "here" means the group of bindings dropped at the top of the fork
 
-          (used_here : used_in_flags) = [ fvs `intersectsVarSet` bndrs
+          (used_here : used_in_flags) = [ fvs `intersectsDVarSet` bndrs
                                         | (fvs, _) <- drop_boxes]
 
-          drop_here = used_here || not can_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...
+          drop_here = used_here || cant_push
 
-          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 `unionVarSet` bind_fvs, bind_w_fvs:drops)
+          insert (fvs,drops) = (fvs `unionDVarSet` bind_fvs, bind_w_fvs:drops)
 
           insert_maybe box True  = insert box
           insert_maybe box False = box
@@ -545,10 +734,26 @@ sepBindsByDropPoint dflags is_case drop_pts 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 = mapUnionVarSet fbFVs binds
+floatedBindsFVs binds = mapUnionDVarSet fbFVs binds
 
-fbFVs :: FloatInBind -> VarSet
+fbFVs :: FloatInBind -> DVarSet
 fbFVs (FB _ fvs _) = fvs
 
 wrapFloats :: FloatInBinds -> CoreExpr -> CoreExpr
@@ -560,3 +765,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