[project @ 2001-04-13 13:37:24 by panne]
[ghc-hetmet.git] / ghc / compiler / simplCore / SimplUtils.lhs
index 90a759d..501dd60 100644 (file)
@@ -5,13 +5,13 @@
 
 \begin{code}
 module SimplUtils (
-       simplBinder, simplBinders, simplIds,
-       transformRhs,
-       mkCase, findAlt, findDefault,
+       simplBinder, simplBinders, simplRecIds, simplLetId,
+       tryRhsTyLam, tryEtaExpansion,
+       mkCase,
 
        -- The continuation type
        SimplCont(..), DupFlag(..), contIsDupable, contResultType,
-       countValArgs, countArgs,
+       countValArgs, countArgs, mkRhsStop, mkStop,
        getContArgs, interestingCallContext, interestingArg, isStrictType, discardInline
 
     ) where
@@ -19,31 +19,34 @@ module SimplUtils (
 #include "HsVersions.h"
 
 import CmdLineOpts     ( switchIsOn, SimplifierSwitch(..),
-                         opt_SimplDoLambdaEtaExpansion, opt_SimplCaseMerge, opt_DictsStrict
+                         opt_SimplDoLambdaEtaExpansion, opt_SimplCaseMerge, opt_DictsStrict,
+                         opt_UF_UpdateInPlace
                        )
 import CoreSyn
-import CoreUnfold      ( isValueUnfolding )
-import CoreUtils       ( exprIsTrivial, cheapEqExpr, exprType, exprIsCheap, exprEtaExpandArity, bindNonRec )
-import Subst           ( InScopeSet, mkSubst, substBndrs, substBndr, substIds, lookupIdSubst )
-import Id              ( Id, idType, isId, idName, 
-                         idOccInfo, idUnfolding, idStrictness,
-                         mkId, idInfo
+import CoreUtils       ( exprIsTrivial, cheapEqExpr, exprType, exprIsCheap, 
+                         etaExpand, exprEtaExpandArity, bindNonRec, mkCoerce,
+                         findDefault
                        )
-import IdInfo          ( StrictnessInfo(..), arityLowerBound, setOccInfo, vanillaIdInfo )
+import Subst           ( InScopeSet, mkSubst, substExpr )
+import qualified Subst ( simplBndrs, simplBndr, simplLetId )
+import Id              ( idType, idName, 
+                         idUnfolding, idStrictness,
+                         mkLocalId, idInfo
+                       )
+import IdInfo          ( StrictnessInfo(..) )
 import Maybes          ( maybeToBool, catMaybes )
-import Name            ( isLocalName, setNameUnique )
-import Demand          ( Demand, isStrict, wwLazy, wwLazy )
+import Name            ( setNameUnique )
+import Demand          ( isStrict )
 import SimplMonad
-import Type            ( Type, tyVarsOfType, tyVarsOfTypes, mkForAllTys, seqType, repType,
-                         splitTyConApp_maybe, mkTyVarTys, applyTys, splitFunTys, mkFunTys,
-                         isDictTy, isDataType, applyTy, splitFunTy, isUnLiftedType,
+import Type            ( Type, mkForAllTys, seqType, repType,
+                         splitTyConApp_maybe, tyConAppArgs, mkTyVarTys,
+                         isDictTy, isDataType, isUnLiftedType,
                          splitRepFunTys
                        )
 import TyCon           ( tyConDataConsIfAvailable )
 import DataCon         ( dataConRepArity )
-import VarSet
-import VarEnv          ( SubstEnv, SubstResult(..) )
-import Util            ( lengthExceeds )
+import VarEnv          ( SubstEnv )
+import Util            ( lengthExceeds, mapAccumL )
 import Outputable
 \end{code}
 
@@ -56,7 +59,10 @@ import Outputable
 
 \begin{code}
 data SimplCont         -- Strict contexts
-  = Stop OutType               -- Type of the result
+  = Stop     OutType           -- Type of the result
+            Bool               -- True => This is the RHS of a thunk whose type suggests
+                               --         that update-in-place would be possible
+                               --         (This makes the inliner a little keener.)
 
   | CoerceIt OutType                   -- The To-type, simplified
             SimplCont
@@ -83,7 +89,7 @@ data SimplCont                -- Strict contexts
                                -- The result expression in the OutExprStuff has type cont_ty
 
 instance Outputable SimplCont where
-  ppr (Stop _)                      = ptext SLIT("Stop")
+  ppr (Stop _ _)                            = ptext SLIT("Stop")
   ppr (ApplyTo dup arg se cont)      = (ptext SLIT("ApplyTo") <+> ppr dup <+> ppr arg) $$ ppr cont
   ppr (ArgOf   dup _ _)             = ptext SLIT("ArgOf...") <+> ppr dup
   ppr (Select dup bndr alts se cont) = (ptext SLIT("Select") <+> ppr dup <+> ppr bndr) $$ 
@@ -97,9 +103,16 @@ instance Outputable DupFlag where
   ppr OkToDup = ptext SLIT("ok")
   ppr NoDup   = ptext SLIT("nodup")
 
+
+-------------------
+mkRhsStop, mkStop :: OutType -> SimplCont
+mkStop    ty = Stop ty False
+mkRhsStop ty = Stop ty (canUpdateInPlace ty)
+
+
 -------------------
 contIsDupable :: SimplCont -> Bool
-contIsDupable (Stop _)                  = True
+contIsDupable (Stop _ _)                        = True
 contIsDupable (ApplyTo  OkToDup _ _ _)   = True
 contIsDupable (ArgOf    OkToDup _ _)     = True
 contIsDupable (Select   OkToDup _ _ _ _) = True
@@ -115,21 +128,22 @@ discardInline cont                   = cont
 
 -------------------
 discardableCont :: SimplCont -> Bool
-discardableCont (Stop _)           = False
+discardableCont (Stop _ _)         = False
 discardableCont (CoerceIt _ cont)   = discardableCont cont
 discardableCont (InlinePlease cont) = discardableCont cont
 discardableCont other              = True
 
 discardCont :: SimplCont       -- A continuation, expecting
            -> SimplCont        -- Replace the continuation with a suitable coerce
-discardCont (Stop to_ty) = Stop to_ty
-discardCont cont        = CoerceIt to_ty (Stop to_ty)
-                        where
-                          to_ty = contResultType cont
+discardCont cont = case cont of
+                    Stop to_ty _ -> cont
+                    other        -> CoerceIt to_ty (mkStop to_ty)
+                where
+                  to_ty = contResultType cont
 
 -------------------
 contResultType :: SimplCont -> OutType
-contResultType (Stop to_ty)         = to_ty
+contResultType (Stop to_ty _)       = to_ty
 contResultType (ArgOf _ to_ty _)     = to_ty
 contResultType (ApplyTo _ _ _ cont)  = contResultType cont
 contResultType (CoerceIt _ cont)     = contResultType cont
@@ -253,23 +267,24 @@ interestingArg :: InScopeSet -> InExpr -> SubstEnv -> Bool
        -- (i.e. they are probably lambda bound): f x y z
        -- There is little point in inlining f here.
 interestingArg in_scope arg subst
-  = analyse arg
+  = analyse (substExpr (mkSubst in_scope subst) arg)
+       -- 'analyse' only looks at the top part of the result
+       -- and substExpr is lazy, so this isn't nearly as brutal
+       -- as it looks.
   where
-    analyse (Var v)
-       = case lookupIdSubst (mkSubst in_scope subst) v of
-           DoneId v' _ -> hasSomeUnfolding (idUnfolding v')
-                                       -- was: isValueUnfolding (idUnfolding v')
-                                       -- But that seems over-pessimistic
-
-           other       -> True         -- was: False
-                                       -- But that is *definitely* too pessimistic.
-                                       -- E.g.         let x = 3 in f 
-                                       -- Here, x will be unconditionally substituted, via
-                                       -- the substitution!
+    analyse (Var v)          = hasSomeUnfolding (idUnfolding v)
+                               -- Was: isValueUnfolding (idUnfolding v')
+                               -- But that seems over-pessimistic
     analyse (Type _)         = False
     analyse (App fn (Type _)) = analyse fn
     analyse (Note _ a)       = analyse a
     analyse other            = True
+       -- Consider     let x = 3 in f x
+       -- The substitution will contain (x -> ContEx 3), and we want to
+       -- to say that x is an interesting argument.
+       -- But consider also (\x. f x y) y
+       -- The substitution will contain (x -> ContEx y), and we want to say
+       -- that x is not interesting (assuming y has no unfolding)
 \end{code}
 
 Comment about interestingCallContext
@@ -316,11 +331,15 @@ interestingCallContext :: Bool            -- False <=> no args at all
        --      as scrutinee of a case          Select
        --      as arg of a strict fn           ArgOf
        -- then we should not inline it (unless there is some other reason,
-       -- e.g. is is the sole occurrence).  
-       -- Why not?  At least in the case-scrutinee situation, turning
-       --      case x of y -> ...
+       -- e.g. is is the sole occurrence).  We achieve this by making
+       -- interestingCallContext return False for a lone variable.
+       --
+       -- Why?  At least in the case-scrutinee situation, turning
+       --      let x = (a,b) in case x of y -> ...
        -- into
-       --      let y = (a,b) in ...
+       --      let x = (a,b) in case (a,b) of y -> ...
+       -- and thence to 
+       --      let x = (a,b) in let y = (a,b) in ...
        -- is bad if the binding for x will remain.
        --
        -- Another example: I discovered that strings
@@ -333,12 +352,13 @@ interestingCallContext :: Bool            -- False <=> no args at all
        -- the context can ``see'' the unfolding of the variable (e.g. case or a RULE)
        -- so there's no gain.
        --
-       -- However, even a type application isn't a lone variable.  Consider
+       -- However, even a type application or coercion isn't a lone variable.
+       -- Consider
        --      case $fMonadST @ RealWorld of { :DMonad a b c -> c }
        -- We had better inline that sucker!  The case won't see through it.
        --
-       -- For now, I'm treating treating a variable applied to types as
-       -- "lone". The motivating example was
+       -- For now, I'm treating treating a variable applied to types 
+       -- in a *lazy* context "lone". The motivating example was
        --      f = /\a. \x. BIG
        --      g = /\a. \y.  h (f a)
        -- There's no advantage in inlining f here, and perhaps
@@ -347,12 +367,15 @@ interestingCallContext :: Bool            -- False <=> no args at all
 interestingCallContext some_args some_val_args cont
   = interesting cont
   where
-    interesting (InlinePlease _)   = True
-    interesting (ApplyTo _ _ _ _)  = some_args -- Can happen if we have (coerce t (f x)) y
-    interesting (Select _ _ _ _ _) = some_args
-    interesting (ArgOf _ _ _)     = some_val_args
-    interesting (Stop ty)         = some_val_args && canUpdateInPlace ty
-    interesting (CoerceIt _ cont)  = interesting cont
+    interesting (InlinePlease _)       = True
+    interesting (Select _ _ _ _ _)     = some_args
+    interesting (ApplyTo _ _ _ _)      = True  -- Can happen if we have (coerce t (f x)) y
+                                               -- Perhaps True is a bit over-keen, but I've
+                                               -- seen (coerce f) x, where f has an INLINE prag,
+                                               -- So we have to give some motivaiton for inlining it
+    interesting (ArgOf _ _ _)         = some_val_args
+    interesting (Stop ty upd_in_place) = some_val_args && upd_in_place
+    interesting (CoerceIt _ cont)      = interesting cont
        -- If this call is the arg of a strict function, the context
        -- is a bit interesting.  If we inline here, we may get useful
        -- evaluation information to avoid repeated evals: e.g.
@@ -381,7 +404,10 @@ canUpdateInPlace :: Type -> Bool
 
 -- Note the repType: we want to look through newtypes for this purpose
 
-canUpdateInPlace ty = case splitTyConApp_maybe (repType ty) of {
+canUpdateInPlace ty 
+  | not opt_UF_UpdateInPlace = False
+  | otherwise
+  = case splitTyConApp_maybe (repType ty) of {
                        Nothing         -> False ;
                        Just (tycon, _) -> 
 
@@ -406,7 +432,7 @@ simplBinders :: [InBinder] -> ([OutBinder] -> SimplM a) -> SimplM a
 simplBinders bndrs thing_inside
   = getSubst           `thenSmpl` \ subst ->
     let
-       (subst', bndrs') = substBndrs subst bndrs
+       (subst', bndrs') = Subst.simplBndrs subst bndrs
     in
     seqBndrs bndrs'    `seq`
     setSubst subst' (thing_inside bndrs')
@@ -415,23 +441,29 @@ simplBinder :: InBinder -> (OutBinder -> SimplM a) -> SimplM a
 simplBinder bndr thing_inside
   = getSubst           `thenSmpl` \ subst ->
     let
-       (subst', bndr') = substBndr subst bndr
+       (subst', bndr') = Subst.simplBndr subst bndr
     in
     seqBndr bndr'      `seq`
     setSubst subst' (thing_inside bndr')
 
 
--- Same semantics as simplBinders, but a little less 
--- plumbing and hence a little more efficient.
--- Maybe not worth the candle?
-simplIds :: [InBinder] -> ([OutBinder] -> SimplM a) -> SimplM a
-simplIds ids thing_inside
+simplRecIds :: [InBinder] -> ([OutBinder] -> SimplM a) -> SimplM a
+simplRecIds ids thing_inside
   = getSubst           `thenSmpl` \ subst ->
     let
-       (subst', bndrs') = substIds subst ids
+       (subst', ids') = mapAccumL Subst.simplLetId subst ids
     in
-    seqBndrs bndrs'    `seq`
-    setSubst subst' (thing_inside bndrs')
+    seqBndrs ids'      `seq`
+    setSubst subst' (thing_inside ids')
+
+simplLetId :: InBinder -> (OutBinder -> SimplM a) -> SimplM a
+simplLetId id thing_inside
+  = getSubst           `thenSmpl` \ subst ->
+    let
+       (subst', id') = Subst.simplLetId subst id
+    in
+    seqBndr id'        `seq`
+    setSubst subst' (thing_inside id')
 
 seqBndrs [] = ()
 seqBndrs (b:bs) = seqBndr b `seq` seqBndrs bs
@@ -445,25 +477,6 @@ seqBndr b | isTyVar b = b `seq` ()
 
 %************************************************************************
 %*                                                                     *
-\subsection{Transform a RHS}
-%*                                                                     *
-%************************************************************************
-
-Try (a) eta expansion
-    (b) type-lambda swizzling
-
-\begin{code}
-transformRhs :: InExpr -> SimplM InExpr
-transformRhs rhs 
-  = tryEtaExpansion body               `thenSmpl` \ body' ->
-    mkRhsTyLam tyvars body'
-  where
-    (tyvars, body) = collectTyBinders rhs
-\end{code}
-
-
-%************************************************************************
-%*                                                                     *
 \subsection{Local tyvar-lifting}
 %*                                                                     *
 %************************************************************************
@@ -491,7 +504,7 @@ let-floating.
 This optimisation is CRUCIAL in eliminating the junk introduced by
 desugaring mutually recursive definitions.  Don't eliminate it lightly!
 
-So far as the implemtation is concerned:
+So far as the implementation is concerned:
 
        Invariant: go F e = /\tvs -> F e
        
@@ -533,25 +546,35 @@ as we would normally do.
 
 
 \begin{code}
-mkRhsTyLam tyvars body                 -- Only does something if there's a let
-  | null tyvars || not (worth_it body) -- inside a type lambda, and a WHNF inside that
-  = returnSmpl (mkLams tyvars body)
+tryRhsTyLam :: OutExpr -> SimplM ([OutBind], OutExpr)
+
+tryRhsTyLam rhs                        -- Only does something if there's a let
+  | null tyvars || not (worth_it body) -- inside a type lambda, 
+  = returnSmpl ([], rhs)               -- and a WHNF inside that
+
   | otherwise
-  = go (\x -> x) body
+  = go (\x -> x) body          `thenSmpl` \ (binds, body') ->
+    returnSmpl (binds,  mkLams tyvars body')
+
   where
-    worth_it (Let _ e)      = whnf_in_middle e
-    worth_it other                  = False
+    (tyvars, body) = collectTyBinders rhs
+
+    worth_it e@(Let _ _) = whnf_in_middle e
+    worth_it e          = False
+
+    whnf_in_middle (Let (NonRec x rhs) e) | isUnLiftedType (idType x) = False
     whnf_in_middle (Let _ e) = whnf_in_middle e
     whnf_in_middle e        = exprIsCheap e
 
-
-    go fn (Let bind@(NonRec var rhs) body) | exprIsTrivial rhs
+    go fn (Let bind@(NonRec var rhs) body)
+      | exprIsTrivial rhs
       = go (fn . Let bind) body
 
-    go fn (Let bind@(NonRec var rhs) body)
+    go fn (Let (NonRec var rhs) body)
       = mk_poly tyvars_here var                                `thenSmpl` \ (var', rhs') ->
-       go (fn . Let (mk_silly_bind var rhs')) body     `thenSmpl` \ body' ->
-       returnSmpl (Let (NonRec var' (mkLams tyvars_here (fn rhs))) body')
+       go (fn . Let (mk_silly_bind var rhs')) body     `thenSmpl` \ (binds, body') ->
+       returnSmpl (NonRec var' (mkLams tyvars_here (fn rhs)) : binds, body')
+
       where
        tyvars_here = tyvars
                --      main_tyvar_set = mkVarSet tyvars
@@ -576,10 +599,11 @@ mkRhsTyLam tyvars body                    -- Only does something if there's a let
     go fn (Let (Rec prs) body)
        = mapAndUnzipSmpl (mk_poly tyvars_here) vars    `thenSmpl` \ (vars', rhss') ->
         let
-           gn body = fn $ foldr Let body (zipWith mk_silly_bind vars rhss')
+           gn body  = fn (foldr Let body (zipWith mk_silly_bind vars rhss'))
+           new_bind = Rec (vars' `zip` [mkLams tyvars_here (gn rhs) | rhs <- rhss])
         in
-        go gn body                             `thenSmpl` \ body' ->
-        returnSmpl (Let (Rec (vars' `zip` [mkLams tyvars_here (gn rhs) | rhs <- rhss])) body')
+        go gn body                             `thenSmpl` \ (binds, body') -> 
+        returnSmpl (new_bind : binds, body')
        where
         (vars,rhss) = unzip prs
         tyvars_here = tyvars
@@ -587,18 +611,19 @@ mkRhsTyLam tyvars body                    -- Only does something if there's a let
                --       var_tys     = map idType vars
                -- See notes with tyvars_here above
 
-
-    go fn body = returnSmpl (mkLams tyvars (fn body))
+    go fn body = returnSmpl ([], fn body)
 
     mk_poly tyvars_here var
       = getUniqueSmpl          `thenSmpl` \ uniq ->
        let
            poly_name = setNameUnique (idName var) uniq         -- Keep same name
            poly_ty   = mkForAllTys tyvars_here (idType var)    -- But new type of course
+           poly_id   = mkLocalId poly_name poly_ty 
 
-               -- It's crucial to copy the occInfo of the original var, because
-               -- we're looking at occurrence-analysed but as yet unsimplified code!
-               -- In particular, we mustn't lose the loop breakers.
+               -- In the olden days, it was crucial to copy the occInfo of the original var, 
+               -- because we were looking at occurrence-analysed but as yet unsimplified code!
+               -- In particular, we mustn't lose the loop breakers.  BUT NOW we are looking
+               -- at already simplified code, so it doesn't matter
                -- 
                -- It's even right to retain single-occurrence or dead-var info:
                -- Suppose we started with  /\a -> let x = E in B
@@ -607,28 +632,29 @@ mkRhsTyLam tyvars body                    -- Only does something if there's a let
                -- where x* has an INLINE prag on it.  Now, once x* is inlined,
                -- the occurrences of x' will be just the occurrences originally
                -- pinned on x.
-           poly_info = vanillaIdInfo `setOccInfo` idOccInfo var
-
-           poly_id   = mkId poly_name poly_ty poly_info
        in
        returnSmpl (poly_id, mkTyApps (Var poly_id) (mkTyVarTys tyvars_here))
 
-    mk_silly_bind var rhs = NonRec var rhs
-               -- We need to be careful about inlining.
+    mk_silly_bind var rhs = NonRec var (Note InlineMe rhs)
                -- Suppose we start with:
                --
-               --      x = let g = /\a -> \x -> f x x
-               --          in 
-               --          /\ b -> let g* = g b in E
+               --      x = /\ a -> let g = G in E
+               --
+               -- Then we'll float to get
+               --
+               --      x = let poly_g = /\ a -> G
+               --          in /\ a -> let g = poly_g a in E
                --
-               -- Then:        * the binding for g gets floated out
-               --              * but then it MIGHT get inlined into the rhs of g*
-               --              * then the binding for g* is floated out of the /\b
-               --              * so we're back to square one
-               -- We rely on the simplifier not to inline g into the RHS of g*,
-               -- because it's a "lone" occurrence, and there is no benefit in
-               -- inlining.  But it's a slightly delicate property, and there's
-               -- a danger of making the simplifier loop here.
+               -- But now the occurrence analyser will see just one occurrence
+               -- of poly_g, not inside a lambda, so the simplifier will
+               -- PreInlineUnconditionally poly_g back into g!  Badk to square 1!
+               -- (I used to think that the "don't inline lone occurrences" stuff
+               --  would stop this happening, but since it's the *only* occurrence,
+               --  PreInlineUnconditionally kicks in first!)
+               --
+               -- Solution: put an INLINE note on g's RHS, so that poly_g seems
+               --           to appear many times.  (NB: mkInlineMe eliminates
+               --           such notes on trivial RHSs, so do it manually.)
 \end{code}
 
 
@@ -641,91 +667,68 @@ mkRhsTyLam tyvars body                    -- Only does something if there's a let
        Try eta expansion for RHSs
 
 We go for:
-               \x1..xn -> N    ==>   \x1..xn y1..ym -> N y1..ym
-       AND             
-               N E1..En        ==>   let z1=E1 .. zn=En in \y1..ym -> N z1..zn y1..ym
+   Case 1    f = \x1..xn -> N  ==>   f = \x1..xn y1..ym -> N y1..ym
+                (n >= 0)
+     OR                
+   Case 2    f = N E1..En      ==>   z1=E1
+                (n > 0)                 .. 
+                                    zn=En
+                                    f = \y1..ym -> N z1..zn y1..ym
+
+where (in both cases) 
+
+       * The xi can include type variables
 
-where (in both cases) N is a NORMAL FORM (i.e. no redexes anywhere)
-wanting a suitable number of extra args.
+       * The yi are all value variables
 
-NB: the Ei may have unlifted type, but the simplifier (which is applied
-to the result) deals OK with this.
+       * N is a NORMAL FORM (i.e. no redexes anywhere)
+         wanting a suitable number of extra args.
 
-There is no point in looking for a combination of the two, 
-because that would leave use with some lets sandwiched between lambdas;
-that's what the final test in the first equation is for.
+       * the Ei must not have unlifted type
+
+There is no point in looking for a combination of the two, because
+that would leave use with some lets sandwiched between lambdas; that's
+what the final test in the first equation is for.
+
+In Case 1, we may have to sandwich some coerces between the lambdas
+to make the types work.   exprEtaExpandArity looks through coerces
+when computing arity; and etaExpand adds the coerces as necessary when
+actually computing the expansion.
 
 \begin{code}
-tryEtaExpansion :: InExpr -> SimplM InExpr
-tryEtaExpansion rhs
-  |  not opt_SimplDoLambdaEtaExpansion
-  || exprIsTrivial rhs                         -- Don't eta-expand a trival RHS
-  || null y_tys                                        -- No useful expansion
-  || not (null x_bndrs || and trivial_args)    -- Not (no x-binders or no z-binds)
-  = returnSmpl rhs
-
-  | otherwise  -- Consider eta expansion
-  = newIds SLIT("y") y_tys                                     $ ( \ y_bndrs ->
-    tick (EtaExpansion (head y_bndrs))                         `thenSmpl_`
-    mapAndUnzipSmpl bind_z_arg (args `zip` trivial_args)       `thenSmpl` (\ (maybe_z_binds, z_args) ->
-    returnSmpl (mkLams x_bndrs                         $ 
-               mkLets (catMaybes maybe_z_binds)        $
-               mkLams y_bndrs                          $
-               mkApps (mkApps fun z_args) (map Var y_bndrs))))
-  where
-    (x_bndrs, body) = collectValBinders rhs
-    (fun, args)            = collectArgs body
-    trivial_args    = map exprIsTrivial args
-    fun_arity      = exprEtaExpandArity fun
+tryEtaExpansion :: OutExpr -> OutType -> SimplM ([OutBind], OutExpr)
+tryEtaExpansion rhs rhs_ty
+  |  not opt_SimplDoLambdaEtaExpansion                 -- Not if switched off
+  || exprIsTrivial rhs                         -- Not if RHS is trivial
+  || final_arity == 0                          -- Not if arity is zero
+  = returnSmpl ([], rhs)
+
+  | n_val_args == 0 && not arity_is_manifest
+  =    -- Some lambdas but not enough: case 1
+    getUniqSupplySmpl                          `thenSmpl` \ us ->
+    returnSmpl ([], etaExpand final_arity us rhs rhs_ty)
+
+  | n_val_args > 0 && not (any cant_bind arg_infos)
+  =    -- Partial application: case 2
+    mapAndUnzipSmpl bind_z_arg arg_infos       `thenSmpl` \ (maybe_z_binds, z_args) ->
+    getUniqSupplySmpl                          `thenSmpl` \ us ->
+    returnSmpl (catMaybes maybe_z_binds, 
+               etaExpand final_arity us (mkApps fun z_args) rhs_ty)
 
-    bind_z_arg (arg, trivial_arg) 
+  | otherwise
+  = returnSmpl ([], rhs)
+  where
+    (fun, args)                           = collectArgs rhs
+    n_val_args                    = valArgCount args
+    (fun_arity, arity_is_manifest) = exprEtaExpandArity fun
+    final_arity                           = 0 `max` (fun_arity - n_val_args)
+    arg_infos                     = [(arg, exprType arg, exprIsTrivial arg) | arg <- args]
+    cant_bind (_, ty, triv)       = not triv && isUnLiftedType ty
+
+    bind_z_arg (arg, arg_ty, trivial_arg) 
        | trivial_arg = returnSmpl (Nothing, arg)
-        | otherwise   = newId SLIT("z") (exprType arg) $ \ z ->
+        | otherwise   = newId SLIT("z") arg_ty $ \ z ->
                        returnSmpl (Just (NonRec z arg), Var z)
-
-       -- Note: I used to try to avoid the exprType call by using
-       -- the type of the binder.  But this type doesn't necessarily
-       -- belong to the same substitution environment as this rhs;
-       -- and we are going to make extra term binders (y_bndrs) from the type
-       -- which will be processed with the rhs substitution environment.
-       -- This only went wrong in a mind bendingly complicated case.
-    (potential_extra_arg_tys, _) = splitFunTys (exprType body)
-       
-    y_tys :: [InType]
-    y_tys  = take no_extras_wanted potential_extra_arg_tys
-       
-    no_extras_wanted :: Int
-    no_extras_wanted = 0 `max`
-
-       -- We used to expand the arity to the previous arity fo the
-       -- function; but this is pretty dangerous.  Consdier
-       --      f = \xy -> e
-       -- so that f has arity 2.  Now float something into f's RHS:
-       --      f = let z = BIG in \xy -> e
-       -- The last thing we want to do now is to put some lambdas
-       -- outside, to get
-       --      f = \xy -> let z = BIG in e
-       --
-       -- (bndr_arity - no_of_xs)              `max`
-
-       -- See if the body could obviously do with more args
-       (fun_arity - valArgCount args)
-
--- This case is now deal with by exprEtaExpandArity
-       -- Finally, see if it's a state transformer, and xs is non-null
-       -- (so it's also a function not a thunk) in which
-       -- case we eta-expand on principle! This can waste work,
-       -- but usually doesn't.
-       -- I originally checked for a singleton type [ty] in this case
-       -- but then I found a situation in which I had
-       --      \ x -> let {..} in \ s -> f (...) s
-       -- AND f RETURNED A FUNCTION.  That is, 's' wasn't the only
-       -- potential extra arg.
---     case (x_bndrs, potential_extra_arg_tys) of
---         (_:_, ty:_)  -> case splitTyConApp_maybe ty of
---                               Just (tycon,_) | tycon == statePrimTyCon -> 1
---                               other                                    -> 0
---         other -> 0
 \end{code}
 
 
@@ -804,15 +807,28 @@ and similar friends.
 mkCase scrut case_bndr alts
   | all identity_alt alts
   = tick (CaseIdentity case_bndr)              `thenSmpl_`
-    returnSmpl scrut
+    returnSmpl (re_note scrut)
   where
-    identity_alt (DEFAULT, [], Var v)     = v == case_bndr
-    identity_alt (DataAlt con, args, rhs) = cheapEqExpr rhs
-                                                       (mkConApp con (map Type arg_tys ++ map varToCoreExpr args))
-    identity_alt other                   = False
-
-    arg_tys = case splitTyConApp_maybe (idType case_bndr) of
-               Just (tycon, arg_tys) -> arg_tys
+    identity_alt (con, args, rhs) = de_note rhs `cheapEqExpr` identity_rhs con args
+
+    identity_rhs (DataAlt con) args = mkConApp con (arg_tys ++ map varToCoreExpr args)
+    identity_rhs (LitAlt lit)  _    = Lit lit
+    identity_rhs DEFAULT       _    = Var case_bndr
+
+    arg_tys = map Type (tyConAppArgs (idType case_bndr))
+
+       -- We've seen this:
+       --      case coerce T e of x { _ -> coerce T' x }
+       -- And we definitely want to eliminate this case!
+       -- So we throw away notes from the RHS, and reconstruct
+       -- (at least an approximation) at the other end
+    de_note (Note _ e) = de_note e
+    de_note e         = e
+
+       -- re_note wraps a coerce if it might be necessary
+    re_note scrut = case head alts of
+                       (_,_,rhs1@(Note _ _)) -> mkCoerce (exprType rhs1) (idType case_bndr) scrut
+                       other                 -> scrut
 \end{code}
 
 The catch-all case
@@ -823,22 +839,3 @@ mkCase other_scrut case_bndr other_alts
 \end{code}
 
 
-\begin{code}
-findDefault :: [CoreAlt] -> ([CoreAlt], Maybe CoreExpr)
-findDefault []                         = ([], Nothing)
-findDefault ((DEFAULT,args,rhs) : alts) = ASSERT( null alts && null args ) 
-                                         ([], Just rhs)
-findDefault (alt : alts)               = case findDefault alts of 
-                                           (alts', deflt) -> (alt : alts', deflt)
-
-findAlt :: AltCon -> [CoreAlt] -> CoreAlt
-findAlt con alts
-  = go alts
-  where
-    go []          = pprPanic "Missing alternative" (ppr con $$ vcat (map ppr alts))
-    go (alt : alts) | matches alt = alt
-                   | otherwise   = go alts
-
-    matches (DEFAULT, _, _) = True
-    matches (con1, _, _)    = con == con1
-\end{code}