[project @ 2001-06-28 08:36:30 by simonpj]
[ghc-hetmet.git] / ghc / compiler / coreSyn / CoreUtils.lhs
index 131bd47..bea8316 100644 (file)
@@ -5,19 +5,25 @@
 
 \begin{code}
 module CoreUtils (
-       exprType, coreAltsType,
-
        -- Construction
        mkNote, mkInlineMe, mkSCC, mkCoerce,
        bindNonRec, mkIfThenElse, mkAltExpr,
+        mkPiType,
+
+       -- Taking expressions apart
+       findDefault, findAlt, hasDefault,
 
+       -- Properties of expressions
+       exprType, coreAltsType, 
        exprIsBottom, exprIsDupable, exprIsTrivial, exprIsCheap, 
        exprIsValue,exprOkForSpeculation, exprIsBig, 
-       exprArity, exprIsConApp_maybe,
-
+       exprIsConApp_maybe, exprIsAtom,
        idAppIsBottom, idAppIsCheap,
+       exprArity, 
 
-       etaReduceExpr, exprEtaExpandArity,
+       -- Expr transformation
+       etaReduce, etaExpand,
+       exprArity, exprEtaExpandArity, 
 
        -- Size
        coreBindsSize,
@@ -40,29 +46,25 @@ import PprCore              ( pprCoreExpr )
 import Var             ( Var, isId, isTyVar )
 import VarSet
 import VarEnv
-import Name            ( isLocallyDefined, hashName )
-import Literal         ( Literal, hashLiteral, literalType )
+import Name            ( hashName )
+import Literal         ( hashLiteral, literalType, litIsDupable )
 import DataCon         ( DataCon, dataConRepArity )
 import PrimOp          ( primOpOkForSpeculation, primOpIsCheap )
-import Id              ( Id, idType, idFlavour, idStrictness, idLBVarInfo, mkWildId,
-                         idArity, idName, idUnfolding, idInfo, isDataConId_maybe
-
+import Id              ( Id, idType, globalIdDetails, idStrictness, idLBVarInfo, 
+                         mkWildId, idArity, idName, idUnfolding, idInfo, isOneShotLambda,
+                         isDataConId_maybe, mkSysLocal, hasNoBinding
                        )
-import IdInfo          ( arityLowerBound, InlinePragInfo(..),
-                         LBVarInfo(..),  
-                         IdFlavour(..),
+import IdInfo          ( LBVarInfo(..),  
+                         GlobalIdDetails(..),
                          megaSeqIdInfo )
 import Demand          ( appIsBottom )
-import Type            ( Type, mkFunTy, mkForAllTy,
-                         splitFunTy_maybe, tyVarsOfType, tyVarsOfTypes,
-                          isNotUsgTy, mkUsgTy, unUsgTy, UsageAnn(..),
-                         applyTys, isUnLiftedType, seqType
+import Type            ( Type, mkFunTy, mkForAllTy, splitFunTy_maybe, 
+                         applyTys, isUnLiftedType, seqType, mkUTy, mkTyVarTy,
+                         splitForAllTy_maybe, isForAllTy, eqType
                        )
-import TysWiredIn      ( boolTy, stringTy, trueDataCon, falseDataCon )
+import TysWiredIn      ( boolTy, trueDataCon, falseDataCon )
 import CostCentre      ( CostCentre )
-import Unique          ( buildIdKey, augmentIdKey )
-import Util            ( zipWithEqual, mapAccumL )
-import Maybes          ( maybeToBool )
+import UniqSupply      ( UniqSupply, splitUniqSupply, uniqFromSupply )
 import Outputable
 import TysPrim         ( alphaTy )     -- Debugging only
 \end{code}
@@ -82,15 +84,8 @@ exprType (Lit lit)           = literalType lit
 exprType (Let _ body)          = exprType body
 exprType (Case _ _ alts)        = coreAltsType alts
 exprType (Note (Coerce ty _) e) = ty  -- **! should take usage from e
-exprType (Note (TermUsg u) e)   = mkUsgTy u (unUsgTy (exprType e))
 exprType (Note other_note e)    = exprType e
-exprType (Lam binder expr)
-  | isId binder    = (case idLBVarInfo binder of
-                       IsOneShotLambda -> mkUsgTy UsOnce
-                       otherwise       -> id) $
-                     idType binder `mkFunTy` exprType expr
-  | isTyVar binder = mkForAllTy binder (exprType expr)
-
+exprType (Lam binder expr)      = mkPiType binder (exprType expr)
 exprType e@(App _ _)
   = case collectArgs e of
        (fun, args) -> applyTypeToArgs e (exprType fun) args
@@ -101,6 +96,20 @@ coreAltsType :: [CoreAlt] -> Type
 coreAltsType ((_,_,rhs) : _) = exprType rhs
 \end{code}
 
+@mkPiType@ makes a (->) type or a forall type, depending on whether
+it is given a type variable or a term variable.  We cleverly use the
+lbvarinfo field to figure out the right annotation for the arrove in
+case of a term variable.
+
+\begin{code}
+mkPiType :: Var -> Type -> Type                -- The more polymorphic version doesn't work...
+mkPiType v ty | isId v    = (case idLBVarInfo v of
+                               LBVarInfo u -> mkUTy u
+                               otherwise   -> id) $
+                            mkFunTy (idType v) ty
+             | isTyVar v = mkForAllTy v ty
+\end{code}
+
 \begin{code}
 -- The first argument is just for debugging
 applyTypeToArgs :: CoreExpr -> Type -> [CoreExpr] -> Type
@@ -108,9 +117,6 @@ applyTypeToArgs e op_ty [] = op_ty
 
 applyTypeToArgs e op_ty (Type ty : args)
   =    -- Accumulate type arguments so we can instantiate all at once
-    ASSERT2( all isNotUsgTy tys, 
-            ppr e <+> text "of" <+> ppr op_ty <+> text "to" <+> 
-            ppr (Type ty : args) <+> text "i.e." <+> ppr tys )
     applyTypeToArgs e (applyTys op_ty tys) rest_args
   where
     (tys, rest_args)        = go [ty] args
@@ -151,35 +157,52 @@ Drop trivial InlineMe's.  This is somewhat important, because if we have an unfo
 that looks like        (Note InlineMe (Var v)), the InlineMe doesn't go away because it may
 not be *applied* to anything.
 
+We don't use exprIsTrivial here, though, because we sometimes generate worker/wrapper
+bindings like
+       fw = ...
+       f  = inline_me (coerce t fw)
+As usual, the inline_me prevents the worker from getting inlined back into the wrapper.
+We want the split, so that the coerces can cancel at the call site.  
+
+However, we can get left with tiresome type applications.  Notably, consider
+       f = /\ a -> let t = e in (t, w)
+Then lifting the let out of the big lambda gives
+       t' = /\a -> e
+       f = /\ a -> let t = inline_me (t' a) in (t, w)
+The inline_me is to stop the simplifier inlining t' right back
+into t's RHS.  In the next phase we'll substitute for t (since
+its rhs is trivial) and *then* we could get rid of the inline_me.
+But it hardly seems worth it, so I don't bother.
+
 \begin{code}
-mkInlineMe e | exprIsTrivial e = e
-            | otherwise       = Note InlineMe e
+mkInlineMe (Var v) = Var v
+mkInlineMe e      = Note InlineMe e
 \end{code}
 
 
 
 \begin{code}
-mkCoerce :: Type -> Type -> Expr b -> Expr b
--- In (mkCoerce to_ty from_ty e), we require that from_ty = exprType e
--- But exprType is defined in CoreUtils, so we don't check the assertion
+mkCoerce :: Type -> Type -> CoreExpr -> CoreExpr
 
 mkCoerce to_ty from_ty (Note (Coerce to_ty2 from_ty2) expr)
-  = ASSERT( from_ty == to_ty2 )
+  = ASSERT( from_ty `eqType` to_ty2 )
     mkCoerce to_ty from_ty2 expr
 
 mkCoerce to_ty from_ty expr
-  | to_ty == from_ty = expr
-  | otherwise       = Note (Coerce to_ty from_ty) expr
+  | to_ty `eqType` from_ty = expr
+  | otherwise             = ASSERT( from_ty `eqType` exprType expr )
+                            Note (Coerce to_ty from_ty) expr
 \end{code}
 
 \begin{code}
 mkSCC :: CostCentre -> Expr b -> Expr b
        -- Note: Nested SCC's *are* preserved for the benefit of
-       --       cost centre stack profiling (Durham)
-
-mkSCC cc (Lit lit) = Lit lit
-mkSCC cc (Lam x e) = Lam x (mkSCC cc e)        -- Move _scc_ inside lambda
-mkSCC cc expr     = Note (SCC cc) expr
+       --       cost centre stack profiling
+mkSCC cc (Lit lit)         = Lit lit
+mkSCC cc (Lam x e)         = Lam x (mkSCC cc e)  -- Move _scc_ inside lambda
+mkSCC cc (Note (SCC cc') e) = Note (SCC cc) (Note (SCC cc') e)
+mkSCC cc (Note n e)        = Note n (mkSCC cc e) -- Move _scc_ inside notes
+mkSCC cc expr              = Note (SCC cc) expr
 \end{code}
 
 
@@ -221,27 +244,85 @@ mkIfThenElse guard then_expr else_expr
           (DataAlt falseDataCon, [], else_expr) ]
 \end{code}
 
+
+%************************************************************************
+%*                                                                     *
+\subsection{Taking expressions apart}
+%*                                                                     *
+%************************************************************************
+
+The default alternative must be first, if it exists at all.
+This makes it easy to find, though it makes matching marginally harder.
+
+\begin{code}
+hasDefault :: [CoreAlt] -> Bool
+hasDefault ((DEFAULT,_,_) : alts) = True
+hasDefault _                     = False
+
+findDefault :: [CoreAlt] -> ([CoreAlt], Maybe CoreExpr)
+findDefault ((DEFAULT,args,rhs) : alts) = ASSERT( null args ) (alts, Just rhs)
+findDefault alts                       =                     (alts, Nothing)
+
+findAlt :: AltCon -> [CoreAlt] -> CoreAlt
+findAlt con alts
+  = case alts of
+       (deflt@(DEFAULT,_,_):alts) -> go alts deflt
+       other                      -> go alts panic_deflt
+
+  where
+    panic_deflt = pprPanic "Missing alternative" (ppr con $$ vcat (map ppr alts))
+
+    go []                     deflt               = deflt
+    go (alt@(con1,_,_) : alts) deflt | con == con1 = alt
+                                    | otherwise   = ASSERT( not (con1 == DEFAULT) )
+                                                    go alts deflt
+\end{code}
+
+
 %************************************************************************
 %*                                                                     *
 \subsection{Figuring out things about expressions}
 %*                                                                     *
 %************************************************************************
 
-@exprIsTrivial@        is true of expressions we are unconditionally 
-               happy to duplicate; simple variables and constants,
-               and type applications.
+@exprIsTrivial@ is true of expressions we are unconditionally happy to
+               duplicate; simple variables and constants, and type
+               applications.  Note that primop Ids aren't considered
+               trivial unless 
 
 @exprIsBottom@ is true of expressions that are guaranteed to diverge
 
 
 \begin{code}
-exprIsTrivial (Type _)      = True
-exprIsTrivial (Lit lit)      = True
-exprIsTrivial (Var v)       = True
-exprIsTrivial (App e arg)    = isTypeArg arg && exprIsTrivial e
-exprIsTrivial (Note _ e)     = exprIsTrivial e
-exprIsTrivial (Lam b body)   | isTyVar b = exprIsTrivial body
-exprIsTrivial other         = False
+exprIsTrivial (Var v)
+  | hasNoBinding v                    = idArity v == 0
+       -- WAS: | Just op <- isPrimOpId_maybe v      = primOpIsDupable op
+       -- The idea here is that a constructor worker, like $wJust, is
+       -- really short for (\x -> $wJust x), becuase $wJust has no binding.
+       -- So it should be treated like a lambda.
+       -- Ditto unsaturated primops.
+       -- This came up when dealing with eta expansion/reduction for
+       --      x = $wJust
+       -- Here we want to eta-expand.  This looks like an optimisation,
+       -- but it's important (albeit tiresome) that CoreSat doesn't increase 
+       -- anything's arity
+  | otherwise                          = True
+exprIsTrivial (Type _)                = True
+exprIsTrivial (Lit lit)               = True
+exprIsTrivial (App e arg)             = not (isRuntimeArg arg) && exprIsTrivial e
+exprIsTrivial (Note _ e)              = exprIsTrivial e
+exprIsTrivial (Lam b body)             = not (isRuntimeVar b) && exprIsTrivial body
+exprIsTrivial other                   = False
+
+exprIsAtom :: CoreExpr -> Bool
+-- Used to decide whether to let-binding an STG argument
+-- when compiling to ILX => type applications are not allowed
+exprIsAtom (Var v)    = True   -- primOpIsDupable?
+exprIsAtom (Lit lit)  = True
+exprIsAtom (Type ty)  = True
+exprIsAtom (Note (SCC _) e) = False
+exprIsAtom (Note _ e) = exprIsAtom e
+exprIsAtom other      = False
 \end{code}
 
 
@@ -257,10 +338,11 @@ exprIsTrivial other            = False
 
 
 \begin{code}
-exprIsDupable (Type _)      = True
-exprIsDupable (Var v)       = True
-exprIsDupable (Lit lit)      = True
-exprIsDupable (Note _ e)     = exprIsDupable e
+exprIsDupable (Type _)         = True
+exprIsDupable (Var v)          = True
+exprIsDupable (Lit lit)        = litIsDupable lit
+exprIsDupable (Note InlineMe e) = True
+exprIsDupable (Note _ e)        = exprIsDupable e
 exprIsDupable expr          
   = go expr 0
   where
@@ -288,19 +370,16 @@ shared.  The main examples of things which aren't WHNF but are
 
   *    case e of
          pi -> ei
+       (where e, and all the ei are cheap)
 
-       where e, and all the ei are cheap; and
-
-  *    let x = e
-       in b
-
-       where e and b are cheap; and
+  *    let x = e in b
+       (where e and b are cheap)
 
   *    op x1 ... xn
-
-       where op is a cheap primitive operator
+       (where op is a cheap primitive operator)
 
   *    error "foo"
+       (because we are happy to substitute it inside a lambda)
 
 Notice that a variable is considered 'cheap': we can push it inside a lambda,
 because sharing will make sure it is only evaluated once.
@@ -310,12 +389,21 @@ exprIsCheap :: CoreExpr -> Bool
 exprIsCheap (Lit lit)            = True
 exprIsCheap (Type _)             = True
 exprIsCheap (Var _)              = True
+exprIsCheap (Note InlineMe e)            = True
 exprIsCheap (Note _ e)           = exprIsCheap e
-exprIsCheap (Lam x e)            = if isId x then True else exprIsCheap e
-exprIsCheap (Case (Var v) _ alts) = and [exprIsCheap rhs | (_,_,rhs) <- alts]
+exprIsCheap (Lam x e)            = isRuntimeVar x || exprIsCheap e
+exprIsCheap (Case e _ alts)       = exprIsCheap e && 
+                                   and [exprIsCheap rhs | (_,_,rhs) <- alts]
        -- Experimentally, treat (case x of ...) as cheap
+       -- (and case __coerce x etc.)
        -- This improves arities of overloaded functions where
        -- there is only dictionary selection (no construction) involved
+exprIsCheap (Let (NonRec x _) e)  
+      | isUnLiftedType (idType x) = exprIsCheap e
+      | otherwise                = False
+       -- strict lets always have cheap right hand sides, and
+       -- do no allocation.
+
 exprIsCheap other_expr 
   = go other_expr 0 True
   where
@@ -325,13 +413,12 @@ exprIsCheap other_expr
 
          || idAppIsBottom f n_args 
                        -- Application of a function which
-                       -- always gives bottom; we treat this as
-                       -- a WHNF, because it certainly doesn't
-                       -- need to be shared!
+                       -- always gives bottom; we treat this as cheap
+                       -- because it certainly doesn't need to be shared!
        
     go (App f a) n_args args_cheap 
-       | isTypeArg a = go f n_args       args_cheap
-       | otherwise   = go f (n_args + 1) (exprIsCheap a && args_cheap)
+       | not (isRuntimeArg a) = go f n_args      args_cheap
+       | otherwise            = go f (n_args + 1) (exprIsCheap a && args_cheap)
 
     go other   n_args args_cheap = False
 
@@ -340,7 +427,7 @@ idAppIsCheap id n_val_args
   | n_val_args == 0 = True     -- Just a type application of
                                -- a variable (f t1 t2 t3)
                                -- counts as WHNF
-  | otherwise = case idFlavour id of
+  | otherwise = case globalIdDetails id of
                  DataConId _   -> True                 
                  RecordSelId _ -> True                 -- I'm experimenting with making record selection
                                                        -- look cheap, so we will substitute it inside a
@@ -388,7 +475,7 @@ exprOkForSpeculation other_expr
   = go other_expr 0 True
   where
     go (Var f) n_args args_ok 
-      = case idFlavour f of
+      = case globalIdDetails f of
          DataConId _ -> True   -- The strictness of the constructor has already
                                -- been expressed by its "wrapper", so we don't need
                                -- to take the arguments into account
@@ -400,8 +487,8 @@ exprOkForSpeculation other_expr
          other -> False
        
     go (App f a) n_args args_ok 
-       | isTypeArg a = go f n_args       args_ok
-       | otherwise   = go f (n_args + 1) (exprOkForSpeculation a && args_ok)
+       | not (isRuntimeArg a) = go f n_args      args_ok
+       | otherwise            = go f (n_args + 1) (exprOkForSpeculation a && args_ok)
 
     go other n_args args_ok = False
 \end{code}
@@ -425,19 +512,31 @@ idAppIsBottom id n_val_args = appIsBottom (idStrictness id) n_val_args
 \end{code}
 
 @exprIsValue@ returns true for expressions that are certainly *already* 
-evaluated to WHNF.  This is used to decide wether it's ok to change
+evaluated to WHNF.  This is used to decide whether it's ok to change
        case x of _ -> e   ===>   e
 
 and to decide whether it's safe to discard a `seq`
 
-So, it does *not* treat variables as evaluated, unless they say they are
+So, it does *not* treat variables as evaluated, unless they say they are.
+
+But it *does* treat partial applications and constructor applications
+as values, even if their arguments are non-trivial; 
+       e.g.  (:) (f x) (map f xs)      is a value
+             map (...redex...)         is a value
+Because `seq` on such things completes immediately
+
+A possible worry: constructors with unboxed args:
+               C (f x :: Int#)
+Suppose (f x) diverges; then C (f x) is not a value.  True, but
+this form is illegal (see the invariants in CoreSyn).  Args of unboxed
+type must be ok-for-speculation (or trivial).
 
 \begin{code}
 exprIsValue :: CoreExpr -> Bool                -- True => Value-lambda, constructor, PAP
 exprIsValue (Type ty)    = True        -- Types are honorary Values; we don't mind
                                        -- copying them
 exprIsValue (Lit l)      = True
-exprIsValue (Lam b e)            = isId b || exprIsValue e
+exprIsValue (Lam b e)            = isRuntimeVar b || exprIsValue e
 exprIsValue (Note _ e)           = exprIsValue e
 exprIsValue other_expr
   = go other_expr 0
@@ -445,8 +544,8 @@ exprIsValue other_expr
     go (Var f) n_args = idAppIsValue f n_args
        
     go (App f a) n_args
-       | isTypeArg a = go f n_args
-       | otherwise   = go f (n_args + 1) 
+       | not (isRuntimeArg a) = go f n_args
+       | otherwise            = go f (n_args + 1) 
 
     go (Note _ f) n_args = go f n_args
 
@@ -454,7 +553,7 @@ exprIsValue other_expr
 
 idAppIsValue :: Id -> Int -> Bool
 idAppIsValue id n_val_args 
-  = case idFlavour id of
+  = case globalIdDetails id of
        DataConId _ -> True
        PrimOpId _  -> n_val_args < idArity id
        other | n_val_args == 0 -> isEvaldUnfolding (idUnfolding id)
@@ -464,45 +563,31 @@ idAppIsValue id n_val_args
 \end{code}
 
 \begin{code}
-exprArity :: CoreExpr -> Int   -- How many value lambdas are at the top
-exprArity (Lam b e)     | isTyVar b    = exprArity e
-                       | otherwise     = 1 + exprArity e
-
-exprArity (Note note e) | ok_note note = exprArity e
-                       where
-                         ok_note (Coerce _ _) = True
-                               -- We *do* look through coerces when getting arities.
-                               -- Reason: arities are to do with *representation* and
-                               -- work duplication. 
-                         ok_note InlineMe     = True
-                         ok_note InlineCall   = True
-                         ok_note other        = False
-                               -- SCC and TermUsg might be over-conservative?
-
-exprArity other        = 0
-\end{code}
-
-\begin{code}
 exprIsConApp_maybe :: CoreExpr -> Maybe (DataCon, [CoreExpr])
-exprIsConApp_maybe expr
-  = analyse (collectArgs expr)
+exprIsConApp_maybe (Note InlineMe expr) = exprIsConApp_maybe expr
+    -- We ignore InlineMe notes in case we have
+    -- x = __inline_me__ (a,b)
+    -- All part of making sure that INLINE pragmas never hurt
+    -- Marcin tripped on this one when making dictionaries more inlinable
+
+exprIsConApp_maybe expr = analyse (collectArgs expr)
   where
     analyse (Var fun, args)
-       | maybeToBool maybe_con_app = maybe_con_app
-       where
-         maybe_con_app = case isDataConId_maybe fun of
-                               Just con | length args >= dataConRepArity con 
-                                       -- Might be > because the arity excludes type args
-                                        -> Just (con, args)
-                               other    -> Nothing
+       | Just con <- isDataConId_maybe fun,
+         length args >= dataConRepArity con
+               -- Might be > because the arity excludes type args
+       = Just (con,args)
 
+       -- Look through unfoldings, but only cheap ones, because
+       -- we are effectively duplicating the unfolding
     analyse (Var fun, [])
-       = case maybeUnfoldingTemplate (idUnfolding fun) of
-               Nothing  -> Nothing
-               Just unf -> exprIsConApp_maybe unf
+       | let unf = idUnfolding fun,
+         isCheapUnfolding unf
+       = exprIsConApp_maybe (unfoldingTemplate unf)
 
     analyse other = Nothing
-\end{code} 
+\end{code}
+
 
 
 %************************************************************************
@@ -511,7 +596,7 @@ exprIsConApp_maybe expr
 %*                                                                     *
 %************************************************************************
 
-@etaReduceExpr@ trys an eta reduction at the top level of a Core Expr.
+@etaReduce@ trys an eta reduction at the top level of a Core Expr.
 
 e.g.   \ x y -> f x y  ===>  f
 
@@ -520,11 +605,11 @@ The idea is that lambdas are often quite helpful: they indicate
 head normal forms, so we don't want to chuck them away lightly.
 
 \begin{code}
-etaReduceExpr :: CoreExpr -> CoreExpr
+etaReduce :: CoreExpr -> CoreExpr
                -- ToDo: we should really check that we don't turn a non-bottom
                -- lambda into a bottom variable.  Sigh
 
-etaReduceExpr expr@(Lam bndr body)
+etaReduce expr@(Lam bndr body)
   = check (reverse binders) body
   where
     (binders, body) = collectBinders expr
@@ -541,13 +626,18 @@ etaReduceExpr expr@(Lam bndr body)
 
     check _ _ = expr   -- Bale out
 
-etaReduceExpr expr = expr              -- The common case
+etaReduce expr = expr          -- The common case
 \end{code}
        
 
 \begin{code}
-exprEtaExpandArity :: CoreExpr -> Int  -- The number of args the thing can be applied to
-                                       -- without doing much work
+exprEtaExpandArity :: CoreExpr -> (Int, Bool)  
+-- The Int is number of value args the thing can be 
+--     applied to without doing much work
+-- The Bool is True iff there are enough explicit value lambdas
+--     at the top to make this arity apparent
+--     (but ignore it when arity==0)
+
 -- This is used when eta expanding
 --     e  ==>  \xy -> e x y
 --
@@ -555,29 +645,66 @@ exprEtaExpandArity :: CoreExpr -> Int     -- The number of args the thing can be ap
 --     case x of p -> \s -> ...
 -- because for I/O ish things we really want to get that \s to the top.
 -- We are prepared to evaluate x each time round the loop in order to get that
--- Hence "generous" arity
+--
+-- Consider    let x = expensive in \y z -> E
+-- We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
+-- 
+-- Hence the list of Bools returned by go1
+--     NB: this is particularly important/useful for IO state 
+--     transformers, where we often get
+--             let x = E in \ s -> ...
+--     and the \s is a real-world state token abstraction.  Such 
+--     abstractions are almost invariably 1-shot, so we want to
+--     pull the \s out, past the let x=E.  
+--     The hack is in Id.isOneShotLambda
 
 exprEtaExpandArity e
-  = go e
+  = go 0 e
   where
-    go (Var v)                                 = idArity v
-    go (App f (Type _))                        = go f
-    go (App f a)  | exprIsCheap a      = (go f - 1) `max` 0    -- Never go -ve!
-    go (Lam x e)  | isId x             = go e + 1
-                 | otherwise           = go e
-    go (Note n e) | ok_note n          = go e
-    go (Case scrut _ alts)
-      | exprIsCheap scrut              = min_zero [go rhs | (_,_,rhs) <- alts]
-    go (Let b e)       
-      | all exprIsCheap (rhssOfBind b) = go e
-    
-    go other                           = 0
+    go :: Int -> CoreExpr -> (Int,Bool)
+    go ar (Lam x e)  | isId x          = go (ar+1) e
+                    | otherwise        = go ar e
+    go ar (Note n e) | ok_note n       = go ar e
+    go ar other                        = (ar + ar', ar' == 0)
+                                       where
+                                         ar' = length (go1 other)
+
+    go1 :: CoreExpr -> [Bool]
+       -- (go1 e) = [b1,..,bn]
+       -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
+       -- where bi is True <=> the lambda is one-shot
+
+    go1 (Note n e) | ok_note n = go1 e
+    go1 (Var v)                        = replicate (idArity v) False   -- When the type of the Id
+                                                               -- encodes one-shot-ness, use
+                                                               -- the idinfo here
+
+       -- Lambdas; increase arity
+    go1 (Lam x e)  | isId x     = isOneShotLambda x : go1 e
+                  | otherwise  = go1 e
+
+       -- Applications; decrease arity
+    go1 (App f (Type _))       = go1 f
+    go1 (App f a)              = case go1 f of
+                                   (one_shot : xs) | one_shot || exprIsCheap a -> xs
+                                   other                                       -> []
+                                                          
+       -- Case/Let; keep arity if either the expression is cheap
+       -- or it's a 1-shot lambda
+    go1 (Case scrut _ alts) = case foldr1 (zipWith (&&)) [go1 rhs | (_,_,rhs) <- alts] of
+                               xs@(one_shot : _) | one_shot || exprIsCheap scrut -> xs
+                               other                                             -> []
+    go1 (Let b e) = case go1 e of
+                     xs@(one_shot : _) | one_shot || all exprIsCheap (rhssOfBind b) -> xs
+                     other                                                          -> []
+
+    go1 other = []
     
     ok_note (Coerce _ _) = True
     ok_note InlineCall   = True
     ok_note other        = False
            -- Notice that we do not look through __inline_me__
-           -- This one is a bit more surprising, but consider
+           -- This may seem surprising, but consider
            --  f = _inline_me (\x -> e)
            -- We DO NOT want to eta expand this to
            --  f = \x -> (_inline_me (\x -> e)) x
@@ -585,15 +712,95 @@ exprEtaExpandArity e
            -- giving just
            --  f = \x -> e
            -- A Bad Idea
+\end{code}
 
-min_zero :: [Int] -> Int       -- Find the minimum, but zero is the smallest
-min_zero (x:xs) = go x xs
-               where
-                 go 0   xs                 = 0         -- Nothing beats zero
-                 go min []                 = min
-                 go min (x:xs) | x < min   = go x xs
-                               | otherwise = go min xs 
 
+\begin{code}
+etaExpand :: Int               -- Add this number of value args
+         -> UniqSupply
+         -> CoreExpr -> Type   -- Expression and its type
+         -> CoreExpr
+-- (etaExpand n us e ty) returns an expression with 
+-- the same meaning as 'e', but with arity 'n'.  
+
+-- Given e' = etaExpand n us e ty
+-- We should have
+--     ty = exprType e = exprType e'
+--
+-- etaExpand deals with for-alls and coerces. For example:
+--             etaExpand 1 E
+-- where  E :: forall a. T
+--       newtype T = MkT (A -> B)
+--
+-- would return
+--     (/\b. coerce T (\y::A -> (coerce (A->B) (E b) y)
+
+etaExpand n us expr ty
+  | n == 0 && 
+    -- The ILX code generator requires eta expansion for type arguments
+    -- too, but alas the 'n' doesn't tell us how many of them there 
+    -- may be.  So we eagerly eta expand any big lambdas, and just
+    -- cross our fingers about possible loss of sharing in the
+    -- ILX case. 
+    -- The Right Thing is probably to make 'arity' include
+    -- type variables throughout the compiler.  (ToDo.)
+    not (isForAllTy ty)        
+    -- Saturated, so nothing to do
+  = expr
+
+  | otherwise  -- An unsaturated constructor or primop; eta expand it
+  = case splitForAllTy_maybe ty of { 
+         Just (tv,ty') -> Lam tv (etaExpand n us (App expr (Type (mkTyVarTy tv))) ty')
+
+       ; Nothing ->
+  
+       case splitFunTy_maybe ty of {
+         Just (arg_ty, res_ty) -> Lam arg1 (etaExpand (n-1) us2 (App expr (Var arg1)) res_ty)
+                               where
+                                  arg1       = mkSysLocal SLIT("eta") uniq arg_ty
+                                  (us1, us2) = splitUniqSupply us
+                                  uniq       = uniqFromSupply us1 
+                                  
+       ; Nothing -> pprTrace "Bad eta expand" (ppr expr $$ ppr ty) expr
+       }}
+\end{code}
+
+
+exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
+It tells how many things the expression can be applied to before doing
+any work.  It doesn't look inside cases, lets, etc.  The idea is that
+exprEtaExpandArity will do the hard work, leaving something that's easy
+for exprArity to grapple with.  In particular, Simplify uses exprArity to
+compute the ArityInfo for the Id. 
+
+Originally I thought that it was enough just to look for top-level lambdas, but
+it isn't.  I've seen this
+
+       foo = PrelBase.timesInt
+
+We want foo to get arity 2 even though the eta-expander will leave it
+unchanged, in the expectation that it'll be inlined.  But occasionally it
+isn't, because foo is blacklisted (used in a rule).  
+
+Similarly, see the ok_note check in exprEtaExpandArity.  So 
+       f = __inline_me (\x -> e)
+won't be eta-expanded.
+
+And in any case it seems more robust to have exprArity be a bit more intelligent.
+
+\begin{code}
+exprArity :: CoreExpr -> Int
+exprArity e = go e `max` 0
+           where
+             go (Lam x e) | isId x        = go e + 1
+                          | otherwise     = go e
+             go (Note _ e)                = go e
+             go (App e (Type t))          = go e
+             go (App f a) | exprIsCheap a = go f - 1
+               -- Important!  f (fac x) does not have arity 2, 
+               --             even if f does!
+             go (Var v)                   = idArity v
+             go _                         = 0
 \end{code}
 
 
@@ -612,7 +819,7 @@ cheapEqExpr :: Expr b -> Expr b -> Bool
 
 cheapEqExpr (Var v1)   (Var v2)   = v1==v2
 cheapEqExpr (Lit lit1) (Lit lit2) = lit1 == lit2
-cheapEqExpr (Type t1)  (Type t2)  = t1 == t2
+cheapEqExpr (Type t1)  (Type t2)  = t1 `eqType` t2
 
 cheapEqExpr (App f1 a1) (App f2 a2)
   = f1 `cheapEqExpr` f2 && a1 `cheapEqExpr` a2
@@ -632,6 +839,9 @@ exprIsBig other            = True
 \begin{code}
 eqExpr :: CoreExpr -> CoreExpr -> Bool
        -- Works ok at more general type, but only needed at CoreExpr
+       -- Used in rule matching, so when we find a type we use
+       -- eqTcType, which doesn't look through newtypes
+       -- [And it doesn't risk falling into a black hole either.]
 eqExpr e1 e2
   = eq emptyVarEnv e1 e2
   where
@@ -662,7 +872,7 @@ eqExpr e1 e2
                                       env' = extendVarEnv env v1 v2
 
     eq env (Note n1 e1) (Note n2 e2) = eq_note env n1 n2 && eq env e1 e2
-    eq env (Type t1)    (Type t2)    = t1 == t2
+    eq env (Type t1)    (Type t2)    = t1 `eqType` t2
     eq env e1          e2           = False
                                         
     eq_list env []      []       = True
@@ -673,7 +883,7 @@ eqExpr e1 e2
                                         eq (extendVarEnvList env (vs1 `zip` vs2)) r1 r2
 
     eq_note env (SCC cc1)      (SCC cc2)      = cc1 == cc2
-    eq_note env (Coerce t1 f1) (Coerce t2 f2) = t1==t2 && f1==f2
+    eq_note env (Coerce t1 f1) (Coerce t2 f2) = t1 `eqType` t2 && f1 `eqType` f2
     eq_note env InlineCall     InlineCall     = True
     eq_note env other1        other2         = False
 \end{code}
@@ -693,22 +903,24 @@ exprSize :: CoreExpr -> Int
        -- A measure of the size of the expressions
        -- It also forces the expression pretty drastically as a side effect
 exprSize (Var v)       = varSize v 
-exprSize (Lit lit)     = 1
+exprSize (Lit lit)     = lit `seq` 1
 exprSize (App f a)     = exprSize f + exprSize a
 exprSize (Lam b e)     = varSize b + exprSize e
 exprSize (Let b e)     = bindSize b + exprSize e
-exprSize (Case e b as) = exprSize e + varSize b + foldr ((+) . altSize) 0  as
-exprSize (Note n e)    = exprSize e
-exprSize (Type t)      = seqType t `seq`
-                        1
+exprSize (Case e b as) = exprSize e + varSize b + foldr ((+) . altSize) 0 as
+exprSize (Note n e)    = noteSize n + exprSize e
+exprSize (Type t)      = seqType t `seq` 1
 
-exprsSize = foldr ((+) . exprSize) 0 
+noteSize (SCC cc)       = cc `seq` 1
+noteSize (Coerce t1 t2) = seqType t1 `seq` seqType t2 `seq` 1
+noteSize InlineCall     = 1
+noteSize InlineMe       = 1
 
 varSize :: Var -> Int
-varSize b | isTyVar b = 1
-         | otherwise = seqType (idType b)              `seq`
-                       megaSeqIdInfo (idInfo b)        `seq`
-                       1
+varSize b  | isTyVar b = 1
+          | otherwise = seqType (idType b)             `seq`
+                        megaSeqIdInfo (idInfo b)       `seq`
+                        1
 
 varsSize = foldr ((+) . varSize) 0