+\begin{code}
+exprEtaExpandArity :: CoreExpr -> Arity
+{- The Arity returned is the number of value args the
+ thing can be applied to without doing much work
+
+exprEtaExpandArity 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:
+
+1. One-shot lambdas
+
+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
+
+2. The state-transformer hack
+
+The one-shot lambda special cause 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, even if E is expensive. So we treat state-token lambdas as
+one-shot even if they aren't really. The hack is in Id.isOneShotLambda.
+
+3. Dealing with bottom
+
+Consider also
+ f = \x -> error "foo"
+Here, arity 1 is fine. But if it is
+ f = \x -> case x of
+ True -> error "foo"
+ False -> \y -> x+y
+then we want to get arity 2. Tecnically, this isn't quite right, because
+ (f True) `seq` 1
+should diverge, but it'll converge if we eta-expand f. Nevertheless, we
+do so; it improves some programs significantly, and increasing convergence
+isn't a bad thing. Hence the ABot/ATop in ArityType.
+
+Actually, the situation is worse. Consider
+ f = \x -> case x of
+ True -> \y -> x+y
+ False -> \y -> x-y
+Can we eta-expand here? At first the answer looks like "yes of course", but
+consider
+ (f bot) `seq` 1
+This should diverge! But if we eta-expand, it won't. Again, we ignore this
+"problem", because being scrupulous would lose an important transformation for
+many programs.
+-}
+
+
+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 || isStateHack 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 | 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
+
+isStateHack id = case splitTyConApp_maybe (idType id) of
+ Just (tycon,_) | tycon == statePrimTyCon -> True
+ other -> False
+
+ -- The last clause is a gross hack. It claims that
+ -- every function over realWorldStatePrimTy is a one-shot
+ -- function. This is pretty true in practice, and makes a big
+ -- difference. For example, consider
+ -- a `thenST` \ r -> ...E...
+ -- The early full laziness pass, if it doesn't know that r is one-shot
+ -- will pull out E (let's say it doesn't mention r) to give
+ -- let lvl = E in a `thenST` \ r -> ...lvl...
+ -- When `thenST` gets inlined, we end up with
+ -- let lvl = E in \s -> case a s of (r, s') -> ...lvl...
+ -- and we don't re-inline E.
+ --
+ -- It would be better to spot that r was one-shot to start with, but
+ -- I don't want to rely on that.
+ --
+ -- Another good example is in fill_in in PrelPack.lhs. We should be able to
+ -- spot that fill_in has arity 2 (and when Keith is done, we will) but we can't yet.
+
+{- 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}