+exprEtaExpandArity :: CoreExpr -> Arity
+-- The Int is number of value args the thing can be
+-- applied to without doing much work
+--
+-- This is used when eta expanding
+-- e ==> \xy -> e x y
+--
+-- It returns 1 (or more) to:
+-- 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
+
+-- It's all a bit more subtle than it looks. Consider one-shot lambdas
+-- 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 ArityType returned by arityType
+
+-- 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
+--
+-- Consider also
+-- f = \x -> error "foo"
+-- Here, arity 1 is fine. But if it is
+-- f = \x -> case e of
+-- True -> error "foo"
+-- False -> \y -> x+y
+-- then we want to get arity 2.
+-- Hence the ABot/ATop in ArityType
+
+
+exprEtaExpandArity e = arityDepth (arityType e)
+
+-- A limited sort of function type
+data ArityType = AFun Bool ArityType -- True <=> one-shot
+ | ATop -- Know nothing
+ | ABot -- Diverges
+
+arityDepth :: ArityType -> Arity
+arityDepth (AFun _ ty) = 1 + arityDepth ty
+arityDepth ty = 0
+
+andArityType ABot at2 = at2
+andArityType ATop at2 = ATop
+andArityType (AFun t1 at1) (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
+andArityType at1 at2 = andArityType at2 at1
+
+arityType :: CoreExpr -> ArityType
+ -- (go1 e) = [b1,..,bn]
+ -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
+ -- where bi is True <=> the lambda is one-shot
+
+arityType (Note n e) = arityType e
+-- Not needed any more: etaExpand is cleverer
+-- | ok_note n = arityType e
+-- | otherwise = ATop
+
+arityType (Var v)
+ = mk (idArity v)
+ where
+ mk :: Arity -> ArityType
+ mk 0 | isBottomingId v = ABot
+ | otherwise = ATop
+ mk n = AFun False (mk (n-1))
+
+ -- When the type of the Id encodes one-shot-ness,
+ -- use the idinfo here
+
+ -- Lambdas; increase arity
+arityType (Lam x e) | isId x = AFun (isOneShotLambda x) (arityType e)
+ | otherwise = arityType e
+
+ -- Applications; decrease arity
+arityType (App f (Type _)) = arityType f
+arityType (App f a) = case arityType f of
+ AFun one_shot xs | one_shot -> xs
+ | exprIsCheap a -> xs
+ other -> ATop
+
+ -- Case/Let; keep arity if either the expression is cheap
+ -- or it's a 1-shot lambda
+arityType (Case scrut _ alts) = case foldr1 andArityType [arityType rhs | (_,_,rhs) <- alts] of
+ xs@(AFun one_shot _) | one_shot -> xs
+ xs | exprIsCheap scrut -> xs
+ | otherwise -> ATop
+
+arityType (Let b e) = case arityType e of
+ xs@(AFun one_shot _) | one_shot -> xs
+ xs | all exprIsCheap (rhssOfBind b) -> xs
+ | otherwise -> ATop
+
+arityType other = ATop
+
+{- NOT NEEDED ANY MORE: etaExpand is cleverer
+ok_note InlineMe = False
+ok_note other = True
+ -- Notice that we do not look through __inline_me__
+ -- 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
+ -- because the _inline_me gets dropped now it is applied,
+ -- giving just
+ -- f = \x -> e
+ -- A Bad Idea
+-}
+\end{code}
+
+
+\begin{code}
+etaExpand :: Arity -- Result should have this number of value args
+ -> [Unique]
+ -> 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 n us expr ty
+ | manifestArity expr >= n = expr -- The no-op case
+ | otherwise = eta_expand n us expr ty
+ where
+
+-- manifestArity sees how many leading value lambdas there are
+manifestArity :: CoreExpr -> Arity
+manifestArity (Lam v e) | isId v = 1 + manifestArity e
+ | otherwise = manifestArity e
+manifestArity (Note _ e) = manifestArity e
+manifestArity e = 0
+
+-- etaExpand deals with for-alls. For example:
+-- etaExpand 1 E
+-- where E :: forall a. a -> a
+-- would return
+-- (/\b. \y::a -> E b y)
+--
+-- It deals with coerces too, though they are now rare
+-- so perhaps the extra code isn't worth it
+
+eta_expand 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
+
+eta_expand n us (Note note@(Coerce _ ty) e) _
+ = Note note (eta_expand n us e ty)
+
+ -- Use mkNote so that _scc_s get pushed inside any lambdas that
+ -- are generated as part of the eta expansion. We rely on this
+ -- behaviour in CorePrep, when we eta expand an already-prepped RHS.
+eta_expand n us (Note note e) ty
+ = mkNote note (eta_expand n us e ty)
+
+ -- Short cut for the case where there already
+ -- is a lambda; no point in gratuitously adding more
+eta_expand n us (Lam v body) ty
+ | isTyVar v
+ = Lam v (eta_expand n us body (applyTy ty (mkTyVarTy v)))
+
+ | otherwise
+ = Lam v (eta_expand (n-1) us body (funResultTy ty))
+
+eta_expand n us expr ty
+ = case splitForAllTy_maybe ty of {
+ Just (tv,ty') -> Lam tv (eta_expand n us (App expr (Type (mkTyVarTy tv))) ty')
+
+ ; Nothing ->
+
+ case splitFunTy_maybe ty of {
+ Just (arg_ty, res_ty) -> Lam arg1 (eta_expand (n-1) us2 (App expr (Var arg1)) res_ty)
+ where
+ arg1 = mkSysLocal SLIT("eta") uniq arg_ty
+ (uniq:us2) = us
+
+ ; Nothing ->
+
+ case splitNewType_maybe ty of {
+ Just ty' -> mkCoerce ty ty' (eta_expand n us (mkCoerce ty' ty expr) ty') ;
+ Nothing -> pprTrace "Bad eta expand" (ppr expr $$ ppr ty) expr
+ }}}