never executed always true always false
    1 {-
    2 (c) The University of Glasgow 2006
    3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
    4 
    5 
    6         Arity and eta expansion
    7 -}
    8 
    9 {-# LANGUAGE CPP #-}
   10 {-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
   11 
   12 -- | Arity and eta expansion
   13 module GHC.Core.Opt.Arity
   14    ( manifestArity, joinRhsArity, exprArity, typeArity
   15    , exprEtaExpandArity, findRhsArity
   16    , etaExpand, etaExpandAT
   17    , exprBotStrictness_maybe
   18 
   19    -- ** ArityType
   20    , ArityType(..), mkBotArityType, mkTopArityType, expandableArityType
   21    , arityTypeArity, maxWithArity, idArityType
   22 
   23    -- ** Join points
   24    , etaExpandToJoinPoint, etaExpandToJoinPointRule
   25 
   26    -- ** Coercions and casts
   27    , pushCoArg, pushCoArgs, pushCoValArg, pushCoTyArg
   28    , pushCoercionIntoLambda, pushCoDataCon, collectBindersPushingCo
   29    )
   30 where
   31 
   32 import GHC.Prelude
   33 
   34 import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
   35 
   36 import GHC.Core
   37 import GHC.Core.FVs
   38 import GHC.Core.Utils
   39 import GHC.Core.DataCon
   40 import GHC.Core.TyCon     ( tyConArity )
   41 import GHC.Core.TyCon.RecWalk     ( initRecTc, checkRecTc )
   42 import GHC.Core.Predicate ( isDictTy, isCallStackPredTy )
   43 import GHC.Core.Multiplicity
   44 
   45 -- We have two sorts of substitution:
   46 --   GHC.Core.Subst.Subst, and GHC.Core.TyCo.TCvSubst
   47 -- Both have substTy, substCo  Hence need for qualification
   48 import GHC.Core.Subst    as Core
   49 import GHC.Core.Type     as Type
   50 import GHC.Core.Coercion as Type
   51 
   52 import GHC.Types.Demand
   53 import GHC.Types.Var
   54 import GHC.Types.Var.Env
   55 import GHC.Types.Id
   56 import GHC.Types.Var.Set
   57 import GHC.Types.Basic
   58 import GHC.Types.Tickish
   59 
   60 import GHC.Builtin.Uniques
   61 import GHC.Data.FastString
   62 import GHC.Data.Pair
   63 
   64 import GHC.Utils.Constants (debugIsOn)
   65 import GHC.Utils.Outputable
   66 import GHC.Utils.Panic
   67 import GHC.Utils.Panic.Plain
   68 import GHC.Utils.Trace
   69 import GHC.Utils.Misc
   70 
   71 {-
   72 ************************************************************************
   73 *                                                                      *
   74               manifestArity and exprArity
   75 *                                                                      *
   76 ************************************************************************
   77 
   78 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
   79 It tells how many things the expression can be applied to before doing
   80 any work.  It doesn't look inside cases, lets, etc.  The idea is that
   81 exprEtaExpandArity will do the hard work, leaving something that's easy
   82 for exprArity to grapple with.  In particular, Simplify uses exprArity to
   83 compute the ArityInfo for the Id.
   84 
   85 Originally I thought that it was enough just to look for top-level lambdas, but
   86 it isn't.  I've seen this
   87 
   88         foo = PrelBase.timesInt
   89 
   90 We want foo to get arity 2 even though the eta-expander will leave it
   91 unchanged, in the expectation that it'll be inlined.  But occasionally it
   92 isn't, because foo is blacklisted (used in a rule).
   93 
   94 Similarly, see the ok_note check in exprEtaExpandArity.  So
   95         f = __inline_me (\x -> e)
   96 won't be eta-expanded.
   97 
   98 And in any case it seems more robust to have exprArity be a bit more intelligent.
   99 But note that   (\x y z -> f x y z)
  100 should have arity 3, regardless of f's arity.
  101 -}
  102 
  103 manifestArity :: CoreExpr -> Arity
  104 -- ^ manifestArity sees how many leading value lambdas there are,
  105 --   after looking through casts
  106 manifestArity (Lam v e) | isId v        = 1 + manifestArity e
  107                         | otherwise     = manifestArity e
  108 manifestArity (Tick t e) | not (tickishIsCode t) =  manifestArity e
  109 manifestArity (Cast e _)                = manifestArity e
  110 manifestArity _                         = 0
  111 
  112 joinRhsArity :: CoreExpr -> JoinArity
  113 -- Join points are supposed to have manifestly-visible
  114 -- lambdas at the top: no ticks, no casts, nothing
  115 -- Moreover, type lambdas count in JoinArity
  116 joinRhsArity (Lam _ e) = 1 + joinRhsArity e
  117 joinRhsArity _         = 0
  118 
  119 
  120 ---------------
  121 exprArity :: CoreExpr -> Arity
  122 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
  123 exprArity e = go e
  124   where
  125     go (Var v)                     = idArity v
  126     go (Lam x e) | isId x          = go e + 1
  127                  | otherwise       = go e
  128     go (Tick t e) | not (tickishIsCode t) = go e
  129     go (Cast e co)                 = trim_arity (go e) (coercionRKind co)
  130                                         -- Note [exprArity invariant]
  131     go (App e (Type _))            = go e
  132     go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
  133         -- See Note [exprArity for applications]
  134         -- NB: coercions count as a value argument
  135 
  136     go _                           = 0
  137 
  138     trim_arity :: Arity -> Type -> Arity
  139     trim_arity arity ty = arity `min` length (typeArity ty)
  140 
  141 ---------------
  142 typeArity :: Type -> [OneShotInfo]
  143 -- How many value arrows are visible in the type?
  144 -- We look through foralls, and newtypes
  145 -- See Note [exprArity invariant]
  146 typeArity ty
  147   = go initRecTc ty
  148   where
  149     go rec_nts ty
  150       | Just (_, ty')  <- splitForAllTyCoVar_maybe ty
  151       = go rec_nts ty'
  152 
  153       | Just (_,arg,res) <- splitFunTy_maybe ty
  154       = typeOneShot arg : go rec_nts res
  155 
  156       | Just (tc,tys) <- splitTyConApp_maybe ty
  157       , Just (ty', _) <- instNewTyCon_maybe tc tys
  158       , Just rec_nts' <- checkRecTc rec_nts tc  -- See Note [Expanding newtypes]
  159                                                 -- in GHC.Core.TyCon
  160 --   , not (isClassTyCon tc)    -- Do not eta-expand through newtype classes
  161 --                              -- See Note [Newtype classes and eta expansion]
  162 --                              (no longer required)
  163       = go rec_nts' ty'
  164         -- Important to look through non-recursive newtypes, so that, eg
  165         --      (f x)   where f has arity 2, f :: Int -> IO ()
  166         -- Here we want to get arity 1 for the result!
  167         --
  168         -- AND through a layer of recursive newtypes
  169         -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))
  170 
  171       | otherwise
  172       = []
  173 
  174 ---------------
  175 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, DmdSig)
  176 -- A cheap and cheerful function that identifies bottoming functions
  177 -- and gives them a suitable strictness signatures.  It's used during
  178 -- float-out
  179 exprBotStrictness_maybe e
  180   = case getBotArity (arityType botStrictnessArityEnv e) of
  181         Nothing -> Nothing
  182         Just ar -> Just (ar, sig ar)
  183   where
  184     sig ar = mkClosedDmdSig (replicate ar topDmd) botDiv
  185 
  186 {-
  187 Note [exprArity invariant]
  188 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  189 exprArity has the following invariants:
  190 
  191   (1) If typeArity (exprType e) = n,
  192       then manifestArity (etaExpand e n) = n
  193 
  194       That is, etaExpand can always expand as much as typeArity says
  195       So the case analysis in etaExpand and in typeArity must match
  196 
  197   (2) exprArity e <= typeArity (exprType e)
  198 
  199   (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n
  200 
  201       That is, if exprArity says "the arity is n" then etaExpand really
  202       can get "n" manifest lambdas to the top.
  203 
  204 Why is this important?  Because
  205   - In GHC.Iface.Tidy we use exprArity to fix the *final arity* of
  206     each top-level Id, and in
  207   - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
  208     actually match that arity, which in turn means
  209     that the StgRhs has the right number of lambdas
  210 
  211 An alternative would be to do the eta-expansion in GHC.Iface.Tidy, at least
  212 for top-level bindings, in which case we would not need the trim_arity
  213 in exprArity.  That is a less local change, so I'm going to leave it for today!
  214 
  215 Note [Newtype classes and eta expansion]
  216 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  217     NB: this nasty special case is no longer required, because
  218     for newtype classes we don't use the class-op rule mechanism
  219     at all.  See Note [Single-method classes] in GHC.Tc.TyCl.Instance. SLPJ May 2013
  220 
  221 -------- Old out of date comments, just for interest -----------
  222 We have to be careful when eta-expanding through newtypes.  In general
  223 it's a good idea, but annoyingly it interacts badly with the class-op
  224 rule mechanism.  Consider
  225 
  226    class C a where { op :: a -> a }
  227    instance C b => C [b] where
  228      op x = ...
  229 
  230 These translate to
  231 
  232    co :: forall a. (a->a) ~ C a
  233 
  234    $copList :: C b -> [b] -> [b]
  235    $copList d x = ...
  236 
  237    $dfList :: C b -> C [b]
  238    {-# DFunUnfolding = [$copList] #-}
  239    $dfList d = $copList d |> co@[b]
  240 
  241 Now suppose we have:
  242 
  243    dCInt :: C Int
  244 
  245    blah :: [Int] -> [Int]
  246    blah = op ($dfList dCInt)
  247 
  248 Now we want the built-in op/$dfList rule will fire to give
  249    blah = $copList dCInt
  250 
  251 But with eta-expansion 'blah' might (and in #3772, which is
  252 slightly more complicated, does) turn into
  253 
  254    blah = op (\eta. ($dfList dCInt |> sym co) eta)
  255 
  256 and now it is *much* harder for the op/$dfList rule to fire, because
  257 exprIsConApp_maybe won't hold of the argument to op.  I considered
  258 trying to *make* it hold, but it's tricky and I gave up.
  259 
  260 The test simplCore/should_compile/T3722 is an excellent example.
  261 -------- End of old out of date comments, just for interest -----------
  262 
  263 
  264 Note [exprArity for applications]
  265 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  266 When we come to an application we check that the arg is trivial.
  267    eg  f (fac x) does not have arity 2,
  268                  even if f has arity 3!
  269 
  270 * We require that is trivial rather merely cheap.  Suppose f has arity 2.
  271   Then    f (Just y)
  272   has arity 0, because if we gave it arity 1 and then inlined f we'd get
  273           let v = Just y in \w. <f-body>
  274   which has arity 0.  And we try to maintain the invariant that we don't
  275   have arity decreases.
  276 
  277 *  The `max 0` is important!  (\x y -> f x) has arity 2, even if f is
  278    unknown, hence arity 0
  279 
  280 
  281 ************************************************************************
  282 *                                                                      *
  283            Computing the "arity" of an expression
  284 *                                                                      *
  285 ************************************************************************
  286 
  287 Note [Definition of arity]
  288 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  289 The "arity" of an expression 'e' is n if
  290    applying 'e' to *fewer* than n *value* arguments
  291    converges rapidly
  292 
  293 Or, to put it another way
  294 
  295    there is no work lost in duplicating the partial
  296    application (e x1 .. x(n-1))
  297 
  298 In the divergent case, no work is lost by duplicating because if the thing
  299 is evaluated once, that's the end of the program.
  300 
  301 Or, to put it another way, in any context C
  302 
  303    C[ (\x1 .. xn. e x1 .. xn) ]
  304          is as efficient as
  305    C[ e ]
  306 
  307 It's all a bit more subtle than it looks:
  308 
  309 Note [One-shot lambdas]
  310 ~~~~~~~~~~~~~~~~~~~~~~~
  311 Consider one-shot lambdas
  312                 let x = expensive in \y z -> E
  313 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
  314 
  315 Note [Dealing with bottom]
  316 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  317 A Big Deal with computing arities is expressions like
  318 
  319    f = \x -> case x of
  320                True  -> \s -> e1
  321                False -> \s -> e2
  322 
  323 This happens all the time when f :: Bool -> IO ()
  324 In this case we do eta-expand, in order to get that \s to the
  325 top, and give f arity 2.
  326 
  327 This isn't really right in the presence of seq.  Consider
  328         (f bot) `seq` 1
  329 
  330 This should diverge!  But if we eta-expand, it won't.  We ignore this
  331 "problem" (unless -fpedantic-bottoms is on), because being scrupulous
  332 would lose an important transformation for many programs. (See
  333 #5587 for an example.)
  334 
  335 Consider also
  336         f = \x -> error "foo"
  337 Here, arity 1 is fine.  But if it is
  338         f = \x -> case x of
  339                         True  -> error "foo"
  340                         False -> \y -> x+y
  341 then we want to get arity 2.  Technically, this isn't quite right, because
  342         (f True) `seq` 1
  343 should diverge, but it'll converge if we eta-expand f.  Nevertheless, we
  344 do so; it improves some programs significantly, and increasing convergence
  345 isn't a bad thing.  Hence the ABot/ATop in ArityType.
  346 
  347 So these two transformations aren't always the Right Thing, and we
  348 have several tickets reporting unexpected behaviour resulting from
  349 this transformation.  So we try to limit it as much as possible:
  350 
  351  (1) Do NOT move a lambda outside a known-bottom case expression
  352        case undefined of { (a,b) -> \y -> e }
  353      This showed up in #5557
  354 
  355  (2) Do NOT move a lambda outside a case unless
  356      (a) The scrutinee is ok-for-speculation, or
  357      (b) more liberally: the scrutinee is cheap (e.g. a variable), and
  358          -fpedantic-bottoms is not enforced (see #2915 for an example)
  359 
  360 Of course both (1) and (2) are readily defeated by disguising the bottoms.
  361 
  362 4. Note [Newtype arity]
  363 ~~~~~~~~~~~~~~~~~~~~~~~~
  364 Non-recursive newtypes are transparent, and should not get in the way.
  365 We do (currently) eta-expand recursive newtypes too.  So if we have, say
  366 
  367         newtype T = MkT ([T] -> Int)
  368 
  369 Suppose we have
  370         e = coerce T f
  371 where f has arity 1.  Then: etaExpandArity e = 1;
  372 that is, etaExpandArity looks through the coerce.
  373 
  374 When we eta-expand e to arity 1: eta_expand 1 e T
  375 we want to get:                  coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
  376 
  377   HOWEVER, note that if you use coerce bogusly you can ge
  378         coerce Int negate
  379   And since negate has arity 2, you might try to eta expand.  But you can't
  380   decompose Int to a function type.   Hence the final case in eta_expand.
  381 
  382 Note [The state-transformer hack]
  383 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  384 Suppose we have
  385         f = e
  386 where e has arity n.  Then, if we know from the context that f has
  387 a usage type like
  388         t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
  389 then we can expand the arity to m.  This usage type says that
  390 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
  391 Consider f = \x. let y = <expensive>
  392                  in case x of
  393                       True  -> foo
  394                       False -> \(s:RealWorld) -> e
  395 where foo has arity 1.  Then we want the state hack to
  396 apply to foo too, so we can eta expand the case.
  397 
  398 Then we expect that if f is applied to one arg, it'll be applied to two
  399 (that's the hack -- we don't really know, and sometimes it's false)
  400 See also Id.isOneShotBndr.
  401 
  402 Note [State hack and bottoming functions]
  403 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  404 It's a terrible idea to use the state hack on a bottoming function.
  405 Here's what happens (#2861):
  406 
  407   f :: String -> IO T
  408   f = \p. error "..."
  409 
  410 Eta-expand, using the state hack:
  411 
  412   f = \p. (\s. ((error "...") |> g1) s) |> g2
  413   g1 :: IO T ~ (S -> (S,T))
  414   g2 :: (S -> (S,T)) ~ IO T
  415 
  416 Extrude the g2
  417 
  418   f' = \p. \s. ((error "...") |> g1) s
  419   f = f' |> (String -> g2)
  420 
  421 Discard args for bottomming function
  422 
  423   f' = \p. \s. ((error "...") |> g1 |> g3
  424   g3 :: (S -> (S,T)) ~ (S,T)
  425 
  426 Extrude g1.g3
  427 
  428   f'' = \p. \s. (error "...")
  429   f' = f'' |> (String -> S -> g1.g3)
  430 
  431 And now we can repeat the whole loop.  Aargh!  The bug is in applying the
  432 state hack to a function which then swallows the argument.
  433 
  434 This arose in another guise in #3959.  Here we had
  435 
  436      catch# (throw exn >> return ())
  437 
  438 Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()].
  439 After inlining (>>) we get
  440 
  441      catch# (\_. throw {IO ()} exn)
  442 
  443 We must *not* eta-expand to
  444 
  445      catch# (\_ _. throw {...} exn)
  446 
  447 because 'catch#' expects to get a (# _,_ #) after applying its argument to
  448 a State#, not another function!
  449 
  450 In short, we use the state hack to allow us to push let inside a lambda,
  451 but not to introduce a new lambda.
  452 
  453 
  454 Note [ArityType]
  455 ~~~~~~~~~~~~~~~~
  456 ArityType is the result of a compositional analysis on expressions,
  457 from which we can decide the real arity of the expression (extracted
  458 with function exprEtaExpandArity).
  459 
  460 We use the following notation:
  461   at  ::= \o1..on.div
  462   div ::= T | x | ⊥
  463   o   ::= ? | 1
  464 And omit the \. if n = 0. Examples:
  465   \?11.T stands for @AT [NoOneShotInfo,OneShotLam,OneShotLam] topDiv@
  466   ⊥      stands for @AT [] botDiv@
  467 See the 'Outputable' instance for more information. It's pretty simple.
  468 
  469 Here is what the fields mean. If an arbitrary expression 'f' has
  470 ArityType 'at', then
  471 
  472  * If @at = AT [o1,..,on] botDiv@ (notation: \o1..on.⊥), then @f x1..xn@
  473    definitely diverges. Partial applications to fewer than n args may *or
  474    may not* diverge.
  475 
  476    We allow ourselves to eta-expand bottoming functions, even
  477    if doing so may lose some `seq` sharing,
  478        let x = <expensive> in \y. error (g x y)
  479        ==> \y. let x = <expensive> in error (g x y)
  480 
  481  * If @at = AT [o1,..,on] topDiv@ (notation: \o1..on.T), then expanding 'f'
  482    to @\x1..xn. f x1..xn@ loses no sharing, assuming the calls of f respect
  483    the one-shot-ness o1..on of its definition.
  484 
  485    NB 'f' is an arbitrary expression, eg @f = g e1 e2@.  This 'f' can have
  486    arity type @AT oss _@, with @length oss > 0@, only if e1 e2 are themselves
  487    cheap.
  488 
  489  * In both cases, @f@, @f x1@, ... @f x1 ... x(n-1)@ are definitely
  490    really functions, or bottom, but *not* casts from a data type, in
  491    at least one case branch.  (If it's a function in one case branch but
  492    an unsafe cast from a data type in another, the program is bogus.)
  493    So eta expansion is dynamically ok; see Note [State hack and
  494    bottoming functions], the part about catch#
  495 
  496 Example:
  497       f = \x\y. let v = <expensive> in
  498           \s(one-shot) \t(one-shot). blah
  499       'f' has arity type \??11.T
  500       The one-shot-ness means we can, in effect, push that
  501       'let' inside the \st.
  502 
  503 
  504 Suppose f = \xy. x+y
  505 Then  f             :: \??.T
  506       f v           :: \?.T
  507       f <expensive> :: T
  508 -}
  509 
  510 
  511 -- | The analysis lattice of arity analysis. It is isomorphic to
  512 --
  513 -- @
  514 --    data ArityType'
  515 --      = AEnd Divergence
  516 --      | ALam OneShotInfo ArityType'
  517 -- @
  518 --
  519 -- Which is easier to display the Hasse diagram for:
  520 --
  521 -- @
  522 --  ALam OneShotLam at
  523 --          |
  524 --      AEnd topDiv
  525 --          |
  526 --  ALam NoOneShotInfo at
  527 --          |
  528 --      AEnd exnDiv
  529 --          |
  530 --      AEnd botDiv
  531 -- @
  532 --
  533 -- where the @at@ fields of @ALam@ are inductively subject to the same order.
  534 -- That is, @ALam os at1 < ALam os at2@ iff @at1 < at2@.
  535 --
  536 -- Why the strange Top element? See Note [Combining case branches].
  537 --
  538 -- We rely on this lattice structure for fixed-point iteration in
  539 -- 'findRhsArity'. For the semantics of 'ArityType', see Note [ArityType].
  540 data ArityType
  541   = AT ![OneShotInfo] !Divergence
  542   -- ^ @AT oss div@ means this value can safely be eta-expanded @length oss@
  543   -- times, provided use sites respect the 'OneShotInfo's in @oss@.
  544   -- A 'OneShotLam' annotation can come from two sources:
  545   --     * The user annotated a lambda as one-shot with 'GHC.Exts.oneShot'
  546   --     * It's from a lambda binder of a type affected by `-fstate-hack`.
  547   --       See 'idStateHackOneShotInfo'.
  548   -- In both cases, 'OneShotLam' should win over 'NoOneShotInfo', see
  549   -- Note [Combining case branches].
  550   --
  551   -- If @div@ is dead-ending ('isDeadEndDiv'), then application to
  552   -- @length os@ arguments will surely diverge, similar to the situation
  553   -- with 'DmdType'.
  554   deriving Eq
  555 
  556 -- | This is the BNF of the generated output:
  557 --
  558 -- @
  559 -- @
  560 --
  561 -- We format
  562 -- @AT [o1,..,on] topDiv@ as @\o1..on.T@ and
  563 -- @AT [o1,..,on] botDiv@ as @\o1..on.⊥@, respectively.
  564 -- More concretely, @AT [NOI,OS,OS] topDiv@ is formatted as @\?11.T@.
  565 -- If the one-shot info is empty, we omit the leading @\.@.
  566 instance Outputable ArityType where
  567   ppr (AT oss div)
  568     | null oss  = pp_div div
  569     | otherwise = char '\\' <> hcat (map pp_os oss) <> dot <> pp_div div
  570     where
  571       pp_div Diverges = char '⊥'
  572       pp_div ExnOrDiv = char 'x'
  573       pp_div Dunno    = char 'T'
  574       pp_os OneShotLam    = char '1'
  575       pp_os NoOneShotInfo = char '?'
  576 
  577 mkBotArityType :: [OneShotInfo] -> ArityType
  578 mkBotArityType oss = AT oss botDiv
  579 
  580 botArityType :: ArityType
  581 botArityType = mkBotArityType []
  582 
  583 mkTopArityType :: [OneShotInfo] -> ArityType
  584 mkTopArityType oss = AT oss topDiv
  585 
  586 topArityType :: ArityType
  587 topArityType = mkTopArityType []
  588 
  589 -- | The number of value args for the arity type
  590 arityTypeArity :: ArityType -> Arity
  591 arityTypeArity (AT oss _) = length oss
  592 
  593 -- | True <=> eta-expansion will add at least one lambda
  594 expandableArityType :: ArityType -> Bool
  595 expandableArityType at = arityTypeArity at > 0
  596 
  597 -- | See Note [Dead ends] in "GHC.Types.Demand".
  598 -- Bottom implies a dead end.
  599 isDeadEndArityType :: ArityType -> Bool
  600 isDeadEndArityType (AT _ div) = isDeadEndDiv div
  601 
  602 -- | Expand a non-bottoming arity type so that it has at least the given arity.
  603 maxWithArity :: ArityType -> Arity -> ArityType
  604 maxWithArity at@(AT oss div) !ar
  605   | isDeadEndArityType at    = at
  606   | oss `lengthAtLeast` ar   = at
  607   | otherwise                = AT (take ar $ oss ++ repeat NoOneShotInfo) div
  608 
  609 -- | Trim an arity type so that it has at most the given arity.
  610 -- Any excess 'OneShotInfo's are truncated to 'topDiv', even if they end in
  611 -- 'ABot'.
  612 minWithArity :: ArityType -> Arity -> ArityType
  613 minWithArity at@(AT oss _) ar
  614   | oss `lengthAtMost` ar = at
  615   | otherwise             = AT (take ar oss) topDiv
  616 
  617 takeWhileOneShot :: ArityType -> ArityType
  618 takeWhileOneShot (AT oss div)
  619   | isDeadEndDiv div = AT (takeWhile isOneShotInfo oss) topDiv
  620   | otherwise        = AT (takeWhile isOneShotInfo oss) div
  621 
  622 -- | The Arity returned is the number of value args the
  623 -- expression can be applied to without doing much work
  624 exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
  625 -- exprEtaExpandArity is used when eta expanding
  626 --      e  ==>  \xy -> e x y
  627 exprEtaExpandArity dflags e = arityType (etaExpandArityEnv dflags) e
  628 
  629 getBotArity :: ArityType -> Maybe Arity
  630 -- Arity of a divergent function
  631 getBotArity (AT oss div)
  632   | isDeadEndDiv div = Just $ length oss
  633   | otherwise        = Nothing
  634 
  635 ----------------------
  636 findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
  637 -- This implements the fixpoint loop for arity analysis
  638 -- See Note [Arity analysis]
  639 -- If findRhsArity e = (n, is_bot) then
  640 --  (a) any application of e to <n arguments will not do much work,
  641 --      so it is safe to expand e  ==>  (\x1..xn. e x1 .. xn)
  642 --  (b) if is_bot=True, then e applied to n args is guaranteed bottom
  643 findRhsArity dflags bndr rhs old_arity
  644   = go 0 botArityType
  645       -- We always do one step, but usually that produces a result equal to
  646       -- old_arity, and then we stop right away, because old_arity is assumed
  647       -- to be sound. In other words, arities should never decrease.
  648       -- Result: the common case is that there is just one iteration
  649   where
  650     go :: Int -> ArityType -> ArityType
  651     go !n cur_at@(AT oss div)
  652       | not (isDeadEndDiv div)           -- the "stop right away" case
  653       , length oss <= old_arity = cur_at -- from above
  654       | next_at == cur_at       = cur_at
  655       | otherwise               =
  656          -- Warn if more than 2 iterations. Why 2? See Note [Exciting arity]
  657          warnPprTrace (debugIsOn && n > 2)
  658             (text "Exciting arity" $$ nest 2
  659               ( ppr bndr <+> ppr cur_at <+> ppr next_at $$ ppr rhs)) $
  660             go (n+1) next_at
  661       where
  662         next_at = step cur_at
  663 
  664     step :: ArityType -> ArityType
  665     step at = -- pprTrace "step" (ppr bndr <+> ppr at <+> ppr (arityType env rhs)) $
  666               arityType env rhs
  667       where
  668         env = extendSigEnv (findRhsArityEnv dflags) bndr at
  669 
  670 
  671 {-
  672 Note [Arity analysis]
  673 ~~~~~~~~~~~~~~~~~~~~~
  674 The motivating example for arity analysis is this:
  675 
  676   f = \x. let g = f (x+1)
  677           in \y. ...g...
  678 
  679 What arity does f have?  Really it should have arity 2, but a naive
  680 look at the RHS won't see that.  You need a fixpoint analysis which
  681 says it has arity "infinity" the first time round.
  682 
  683 This example happens a lot; it first showed up in Andy Gill's thesis,
  684 fifteen years ago!  It also shows up in the code for 'rnf' on lists
  685 in #4138.
  686 
  687 We do the necessary, quite simple fixed-point iteration in 'findRhsArity',
  688 which assumes for a single binding 'ABot' on the first run and iterates
  689 until it finds a stable arity type. Two wrinkles
  690 
  691 * We often have to ask (see the Case or Let case of 'arityType') whether some
  692   expression is cheap. In the case of an application, that depends on the arity
  693   of the application head! That's why we have our own version of 'exprIsCheap',
  694   'myExprIsCheap', that will integrate the optimistic arity types we have on
  695   f and g into the cheapness check.
  696 
  697 * Consider this (#18793)
  698 
  699     go = \ds. case ds of
  700            []     -> id
  701            (x:ys) -> let acc = go ys in
  702                      case blah of
  703                        True  -> acc
  704                        False -> \ x1 -> acc (negate x1)
  705 
  706   We must propagate go's optimistically large arity to @acc@, so that the
  707   tail call to @acc@ in the True branch has sufficient arity.  This is done
  708   by the 'am_sigs' field in 'FindRhsArity', and 'lookupSigEnv' in the Var case
  709   of 'arityType'.
  710 
  711 Note [Exciting Arity]
  712 ~~~~~~~~~~~~~~~~~~~~~
  713 The fixed-point iteration in 'findRhsArity' stabilises very quickly in almost
  714 all cases. To get notified of cases where we need an usual number of iterations,
  715 we emit a warning in debug mode, so that we can investigate and make sure that
  716 we really can't do better. It's a gross hack, but catches real bugs (#18870).
  717 
  718 Now, which number is "unusual"? We pick n > 2. Here's a pretty common and
  719 expected example that takes two iterations and would ruin the specificity
  720 of the warning (from T18937):
  721 
  722   f :: [Int] -> Int -> Int
  723   f []     = id
  724   f (x:xs) = let y = sum [0..x]
  725              in \z -> f xs (y + z)
  726 
  727 Fixed-point iteration starts with arity type ⊥ for f. After the first
  728 iteration, we get arity type \??.T, e.g. arity 2, because we unconditionally
  729 'floatIn' the let-binding (see its bottom case).  After the second iteration,
  730 we get arity type \?.T, e.g. arity 1, because now we are no longer allowed
  731 to floatIn the non-cheap let-binding.  Which is all perfectly benign, but
  732 means we do two iterations (well, actually 3 'step's to detect we are stable)
  733 and don't want to emit the warning.
  734 
  735 Note [Eta expanding through dictionaries]
  736 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  737 If the experimental -fdicts-cheap flag is on, we eta-expand through
  738 dictionary bindings.  This improves arities. Thereby, it also
  739 means that full laziness is less prone to floating out the
  740 application of a function to its dictionary arguments, which
  741 can thereby lose opportunities for fusion.  Example:
  742         foo :: Ord a => a -> ...
  743      foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
  744         -- So foo has arity 1
  745 
  746      f = \x. foo dInt $ bar x
  747 
  748 The (foo DInt) is floated out, and makes ineffective a RULE
  749      foo (bar x) = ...
  750 
  751 One could go further and make exprIsCheap reply True to any
  752 dictionary-typed expression, but that's more work.
  753 -}
  754 
  755 arityLam :: Id -> ArityType -> ArityType
  756 arityLam id (AT oss div) = AT (idStateHackOneShotInfo id : oss) div
  757 
  758 floatIn :: Bool -> ArityType -> ArityType
  759 -- We have something like (let x = E in b),
  760 -- where b has the given arity type.
  761 floatIn cheap at
  762   | isDeadEndArityType at || cheap = at
  763   -- If E is not cheap, keep arity only for one-shots
  764   | otherwise                      = takeWhileOneShot at
  765 
  766 arityApp :: ArityType -> Bool -> ArityType
  767 -- Processing (fun arg) where at is the ArityType of fun,
  768 -- Knock off an argument and behave like 'let'
  769 arityApp (AT (_:oss) div) cheap = floatIn cheap (AT oss div)
  770 arityApp at               _     = at
  771 
  772 -- | Least upper bound in the 'ArityType' lattice.
  773 -- See the haddocks on 'ArityType' for the lattice.
  774 --
  775 -- Used for branches of a @case@.
  776 andArityType :: ArityType -> ArityType -> ArityType
  777 andArityType (AT (os1:oss1) div1) (AT (os2:oss2) div2)
  778   | AT oss' div' <- andArityType (AT oss1 div1) (AT oss2 div2)
  779   = AT ((os1 `bestOneShot` os2) : oss') div' -- See Note [Combining case branches]
  780 andArityType (AT []         div1) at2
  781   | isDeadEndDiv div1 = at2                  -- Note [ABot branches: max arity wins]
  782   | otherwise         = takeWhileOneShot at2 -- See Note [Combining case branches]
  783 andArityType at1                  (AT []         div2)
  784   | isDeadEndDiv div2 = at1                  -- Note [ABot branches: max arity wins]
  785   | otherwise         = takeWhileOneShot at1 -- See Note [Combining case branches]
  786 
  787 {- Note [ABot branches: max arity wins]
  788 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  789 Consider   case x of
  790              True  -> \x.  error "urk"
  791              False -> \xy. error "urk2"
  792 
  793 Remember: \o1..on.⊥ means "if you apply to n args, it'll definitely diverge".
  794 So we need \??.⊥ for the whole thing, the /max/ of both arities.
  795 
  796 Note [Combining case branches]
  797 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  798 Consider
  799   go = \x. let z = go e0
  800                go2 = \x. case x of
  801                            True  -> z
  802                            False -> \s(one-shot). e1
  803            in go2 x
  804 We *really* want to respect the one-shot annotation provided by the
  805 user and eta-expand go and go2.
  806 When combining the branches of the case we have
  807      T `andAT` \1.T
  808 and we want to get \1.T.
  809 But if the inner lambda wasn't one-shot (\?.T) we don't want to do this.
  810 (We need a usage analysis to justify that.)
  811 
  812 So we combine the best of the two branches, on the (slightly dodgy)
  813 basis that if we know one branch is one-shot, then they all must be.
  814 Surprisingly, this means that the one-shot arity type is effectively the top
  815 element of the lattice.
  816 
  817 Note [Arity trimming]
  818 ~~~~~~~~~~~~~~~~~~~~~
  819 Consider ((\x y. blah) |> co), where co :: (Int->Int->Int) ~ (Int -> F a) , and
  820 F is some type family.
  821 
  822 Because of Note [exprArity invariant], item (2), we must return with arity at
  823 most 1, because typeArity (Int -> F a) = 1.  So we have to trim the result of
  824 calling arityType on (\x y. blah).  Failing to do so, and hence breaking the
  825 exprArity invariant, led to #5441.
  826 
  827 How to trim?  If we end in topDiv, it's easy.  But we must take great care with
  828 dead ends (i.e. botDiv). Suppose the expression was (\x y. error "urk"),
  829 we'll get \??.⊥.  We absolutely must not trim that to \?.⊥, because that
  830 claims that ((\x y. error "urk") |> co) diverges when given one argument,
  831 which it absolutely does not. And Bad Things happen if we think something
  832 returns bottom when it doesn't (#16066).
  833 
  834 So, if we need to trim a dead-ending arity type, switch (conservatively) to
  835 topDiv.
  836 
  837 Historical note: long ago, we unconditionally switched to topDiv when we
  838 encountered a cast, but that is far too conservative: see #5475
  839 
  840 Note [Eta expanding through CallStacks]
  841 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  842 Just as it's good to eta-expand through dictionaries, so it is good to
  843 do so through CallStacks.  #20103 is a case in point, where we got
  844   foo :: HasCallStack => Int -> Int
  845   foo = \(d::CallStack). let d2 = pushCallStack blah d in
  846         \(x:Int). blah
  847 
  848 We really want to eta-expand this!  #20103 is quite convincing!
  849 We do this regardless of -fdicts-cheap; it's not really a dictionary.
  850 -}
  851 
  852 ---------------------------
  853 
  854 -- | Each of the entry-points of the analyser ('arityType') has different
  855 -- requirements. The entry-points are
  856 --
  857 --   1. 'exprBotStrictness_maybe'
  858 --   2. 'exprEtaExpandArity'
  859 --   3. 'findRhsArity'
  860 --
  861 -- For each of the entry-points, there is a separate mode that governs
  862 --
  863 --   1. How pedantic we are wrt. ⊥, in 'pedanticBottoms'.
  864 --   2. Whether we store arity signatures for non-recursive let-bindings,
  865 --      accessed in 'extendSigEnv'/'lookupSigEnv'.
  866 --      See Note [Arity analysis] why that's important.
  867 --   3. Which expressions we consider cheap to float inside a lambda,
  868 --      in 'myExprIsCheap'.
  869 data AnalysisMode
  870   = BotStrictness
  871   -- ^ Used during 'exprBotStrictness_maybe'.
  872   | EtaExpandArity { am_ped_bot :: !Bool
  873                    , am_dicts_cheap :: !Bool }
  874   -- ^ Used for finding an expression's eta-expanding arity quickly, without
  875   -- fixed-point iteration ('exprEtaExpandArity').
  876   | FindRhsArity { am_ped_bot :: !Bool
  877                  , am_dicts_cheap :: !Bool
  878                  , am_sigs :: !(IdEnv ArityType) }
  879   -- ^ Used for regular, fixed-point arity analysis ('findRhsArity').
  880   --   See Note [Arity analysis] for details about fixed-point iteration.
  881   --   INVARIANT: Disjoint with 'ae_joins'.
  882 
  883 data ArityEnv
  884   = AE
  885   { ae_mode   :: !AnalysisMode
  886   -- ^ The analysis mode. See 'AnalysisMode'.
  887   , ae_joins  :: !IdSet
  888   -- ^ In-scope join points. See Note [Eta-expansion and join points]
  889   --   INVARIANT: Disjoint with the domain of 'am_sigs' (if present).
  890   }
  891 
  892 -- | The @ArityEnv@ used by 'exprBotStrictness_maybe'. Pedantic about bottoms
  893 -- and no application is ever considered cheap.
  894 botStrictnessArityEnv :: ArityEnv
  895 botStrictnessArityEnv = AE { ae_mode = BotStrictness, ae_joins = emptyVarSet }
  896 
  897 -- | The @ArityEnv@ used by 'exprEtaExpandArity'.
  898 etaExpandArityEnv :: DynFlags -> ArityEnv
  899 etaExpandArityEnv dflags
  900   = AE { ae_mode  = EtaExpandArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
  901                                    , am_dicts_cheap = gopt Opt_DictsCheap dflags }
  902        , ae_joins = emptyVarSet }
  903 
  904 -- | The @ArityEnv@ used by 'findRhsArity'.
  905 findRhsArityEnv :: DynFlags -> ArityEnv
  906 findRhsArityEnv dflags
  907   = AE { ae_mode  = FindRhsArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
  908                                  , am_dicts_cheap = gopt Opt_DictsCheap dflags
  909                                  , am_sigs = emptyVarEnv }
  910        , ae_joins = emptyVarSet }
  911 
  912 -- First some internal functions in snake_case for deleting in certain VarEnvs
  913 -- of the ArityType. Don't call these; call delInScope* instead!
  914 
  915 modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
  916 modifySigEnv f env@AE { ae_mode = am@FindRhsArity{am_sigs = sigs} } =
  917   env { ae_mode = am { am_sigs = f sigs } }
  918 modifySigEnv _ env = env
  919 {-# INLINE modifySigEnv #-}
  920 
  921 del_sig_env :: Id -> ArityEnv -> ArityEnv -- internal!
  922 del_sig_env id = modifySigEnv (\sigs -> delVarEnv sigs id)
  923 {-# INLINE del_sig_env #-}
  924 
  925 del_sig_env_list :: [Id] -> ArityEnv -> ArityEnv -- internal!
  926 del_sig_env_list ids = modifySigEnv (\sigs -> delVarEnvList sigs ids)
  927 {-# INLINE del_sig_env_list #-}
  928 
  929 del_join_env :: JoinId -> ArityEnv -> ArityEnv -- internal!
  930 del_join_env id env@(AE { ae_joins = joins })
  931   = env { ae_joins = delVarSet joins id }
  932 {-# INLINE del_join_env #-}
  933 
  934 del_join_env_list :: [JoinId] -> ArityEnv -> ArityEnv -- internal!
  935 del_join_env_list ids env@(AE { ae_joins = joins })
  936   = env { ae_joins = delVarSetList joins ids }
  937 {-# INLINE del_join_env_list #-}
  938 
  939 -- end of internal deletion functions
  940 
  941 extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
  942 extendJoinEnv env@(AE { ae_joins = joins }) join_ids
  943   = del_sig_env_list join_ids
  944   $ env { ae_joins = joins `extendVarSetList` join_ids }
  945 
  946 extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
  947 extendSigEnv env id ar_ty
  948   = del_join_env id (modifySigEnv (\sigs -> extendVarEnv sigs id ar_ty) env)
  949 
  950 delInScope :: ArityEnv -> Id -> ArityEnv
  951 delInScope env id = del_join_env id $ del_sig_env id env
  952 
  953 delInScopeList :: ArityEnv -> [Id] -> ArityEnv
  954 delInScopeList env ids = del_join_env_list ids $ del_sig_env_list ids env
  955 
  956 lookupSigEnv :: ArityEnv -> Id -> Maybe ArityType
  957 lookupSigEnv AE{ ae_mode = mode } id = case mode of
  958   BotStrictness                  -> Nothing
  959   EtaExpandArity{}               -> Nothing
  960   FindRhsArity{ am_sigs = sigs } -> lookupVarEnv sigs id
  961 
  962 -- | Whether the analysis should be pedantic about bottoms.
  963 -- 'exprBotStrictness_maybe' always is.
  964 pedanticBottoms :: ArityEnv -> Bool
  965 pedanticBottoms AE{ ae_mode = mode } = case mode of
  966   BotStrictness                          -> True
  967   EtaExpandArity{ am_ped_bot = ped_bot } -> ped_bot
  968   FindRhsArity{ am_ped_bot = ped_bot }   -> ped_bot
  969 
  970 -- | A version of 'exprIsCheap' that considers results from arity analysis
  971 -- and optionally the expression's type.
  972 -- Under 'exprBotStrictness_maybe', no expressions are cheap.
  973 myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
  974 myExprIsCheap AE{ae_mode = mode} e mb_ty = case mode of
  975   BotStrictness -> False
  976   _             -> cheap_dict || cheap_fun e
  977     where
  978       cheap_dict = case mb_ty of
  979                      Nothing -> False
  980                      Just ty -> (am_dicts_cheap mode && isDictTy ty)
  981                                 || isCallStackPredTy ty
  982         -- See Note [Eta expanding through dictionaries]
  983         -- See Note [Eta expanding through CallStacks]
  984 
  985       cheap_fun e = case mode of
  986 #if __GLASGOW_HASKELL__ <= 900
  987         BotStrictness                -> panic "impossible"
  988 #endif
  989         EtaExpandArity{}             -> exprIsCheap e
  990         FindRhsArity{am_sigs = sigs} -> exprIsCheapX (myIsCheapApp sigs) e
  991 
  992 -- | A version of 'isCheapApp' that considers results from arity analysis.
  993 -- See Note [Arity analysis] for what's in the signature environment and why
  994 -- it's important.
  995 myIsCheapApp :: IdEnv ArityType -> CheapAppFun
  996 myIsCheapApp sigs fn n_val_args = case lookupVarEnv sigs fn of
  997   -- Nothing means not a local function, fall back to regular
  998   -- 'GHC.Core.Utils.isCheapApp'
  999   Nothing         -> isCheapApp fn n_val_args
 1000   -- @Just at@ means local function with @at@ as current ArityType.
 1001   -- Roughly approximate what 'isCheapApp' is doing.
 1002   Just (AT oss div)
 1003     | isDeadEndDiv div -> True -- See Note [isCheapApp: bottoming functions] in GHC.Core.Utils
 1004     | n_val_args < length oss -> True -- Essentially isWorkFreeApp
 1005     | otherwise -> False
 1006 
 1007 ----------------
 1008 arityType :: ArityEnv -> CoreExpr -> ArityType
 1009 
 1010 arityType env (Cast e co)
 1011   = minWithArity (arityType env e) co_arity -- See Note [Arity trimming]
 1012   where
 1013     co_arity = length (typeArity (coercionRKind co))
 1014     -- See Note [exprArity invariant] (2); must be true of
 1015     -- arityType too, since that is how we compute the arity
 1016     -- of variables, and they in turn affect result of exprArity
 1017     -- #5441 is a nice demo
 1018 
 1019 arityType env (Var v)
 1020   | v `elemVarSet` ae_joins env
 1021   = botArityType  -- See Note [Eta-expansion and join points]
 1022   | Just at <- lookupSigEnv env v -- Local binding
 1023   = at
 1024   | otherwise
 1025   = idArityType v
 1026 
 1027         -- Lambdas; increase arity
 1028 arityType env (Lam x e)
 1029   | isId x    = arityLam x (arityType env' e)
 1030   | otherwise = arityType env' e
 1031   where
 1032     env' = delInScope env x
 1033 
 1034         -- Applications; decrease arity, except for types
 1035 arityType env (App fun (Type _))
 1036    = arityType env fun
 1037 arityType env (App fun arg )
 1038    = arityApp (arityType env fun) (myExprIsCheap env arg Nothing)
 1039 
 1040         -- Case/Let; keep arity if either the expression is cheap
 1041         -- or it's a 1-shot lambda
 1042         -- The former is not really right for Haskell
 1043         --      f x = case x of { (a,b) -> \y. e }
 1044         --  ===>
 1045         --      f x y = case x of { (a,b) -> e }
 1046         -- The difference is observable using 'seq'
 1047         --
 1048 arityType env (Case scrut bndr _ alts)
 1049   | exprIsDeadEnd scrut || null alts
 1050   = botArityType    -- Do not eta expand. See Note [Dealing with bottom (1)]
 1051   | not (pedanticBottoms env)  -- See Note [Dealing with bottom (2)]
 1052   , myExprIsCheap env scrut (Just (idType bndr))
 1053   = alts_type
 1054   | exprOkForSpeculation scrut
 1055   = alts_type
 1056 
 1057   | otherwise                  -- In the remaining cases we may not push
 1058   = takeWhileOneShot alts_type -- evaluation of the scrutinee in
 1059   where
 1060     env' = delInScope env bndr
 1061     arity_type_alt (Alt _con bndrs rhs) = arityType (delInScopeList env' bndrs) rhs
 1062     alts_type = foldr1 andArityType (map arity_type_alt alts)
 1063 
 1064 arityType env (Let (NonRec j rhs) body)
 1065   | Just join_arity <- isJoinId_maybe j
 1066   , (_, rhs_body)   <- collectNBinders join_arity rhs
 1067   = -- See Note [Eta-expansion and join points]
 1068     andArityType (arityType env rhs_body)
 1069                  (arityType env' body)
 1070   where
 1071      env' = extendJoinEnv env [j]
 1072 
 1073 arityType env (Let (Rec pairs) body)
 1074   | ((j,_):_) <- pairs
 1075   , isJoinId j
 1076   = -- See Note [Eta-expansion and join points]
 1077     foldr (andArityType . do_one) (arityType env' body) pairs
 1078   where
 1079     env' = extendJoinEnv env (map fst pairs)
 1080     do_one (j,rhs)
 1081       | Just arity <- isJoinId_maybe j
 1082       = arityType env' $ snd $ collectNBinders arity rhs
 1083       | otherwise
 1084       = pprPanic "arityType:joinrec" (ppr pairs)
 1085 
 1086 arityType env (Let (NonRec b r) e)
 1087   = floatIn cheap_rhs (arityType env' e)
 1088   where
 1089     cheap_rhs = myExprIsCheap env r (Just (idType b))
 1090     env'      = extendSigEnv env b (arityType env r)
 1091 
 1092 arityType env (Let (Rec prs) e)
 1093   = floatIn (all is_cheap prs) (arityType env' e)
 1094   where
 1095     env'           = delInScopeList env (map fst prs)
 1096     is_cheap (b,e) = myExprIsCheap env' e (Just (idType b))
 1097 
 1098 arityType env (Tick t e)
 1099   | not (tickishIsCode t)     = arityType env e
 1100 
 1101 arityType _ _ = topArityType
 1102 
 1103 {- Note [Eta-expansion and join points]
 1104 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1105 Consider this (#18328)
 1106 
 1107   f x = join j y = case y of
 1108                       True -> \a. blah
 1109                       False -> \b. blah
 1110         in case x of
 1111               A -> j True
 1112               B -> \c. blah
 1113               C -> j False
 1114 
 1115 and suppose the join point is too big to inline.  Now, what is the
 1116 arity of f?  If we inlined the join point, we'd definitely say "arity
 1117 2" because we are prepared to push case-scrutinisation inside a
 1118 lambda.  But currently the join point totally messes all that up,
 1119 because (thought of as a vanilla let-binding) the arity pinned on 'j'
 1120 is just 1.
 1121 
 1122 Why don't we eta-expand j?  Because of
 1123 Note [Do not eta-expand join points] in GHC.Core.Opt.Simplify.Utils
 1124 
 1125 Even if we don't eta-expand j, why is its arity only 1?
 1126 See invariant 2b in Note [Invariants on join points] in GHC.Core.
 1127 
 1128 So we do this:
 1129 
 1130 * Treat the RHS of a join-point binding, /after/ stripping off
 1131   join-arity lambda-binders, as very like the body of the let.
 1132   More precisely, do andArityType with the arityType from the
 1133   body of the let.
 1134 
 1135 * Dually, when we come to a /call/ of a join point, just no-op
 1136   by returning ABot, the bottom element of ArityType,
 1137   which so that: bot `andArityType` x = x
 1138 
 1139 * This works if the join point is bound in the expression we are
 1140   taking the arityType of.  But if it's bound further out, it makes
 1141   no sense to say that (say) the arityType of (j False) is ABot.
 1142   Bad things happen.  So we keep track of the in-scope join-point Ids
 1143   in ae_join.
 1144 
 1145 This will make f, above, have arity 2. Then, we'll eta-expand it thus:
 1146 
 1147   f x eta = (join j y = ... in case x of ...) eta
 1148 
 1149 and the Simplify will automatically push that application of eta into
 1150 the join points.
 1151 
 1152 An alternative (roughly equivalent) idea would be to carry an
 1153 environment mapping let-bound Ids to their ArityType.
 1154 -}
 1155 
 1156 idArityType :: Id -> ArityType
 1157 idArityType v
 1158   | strict_sig <- idDmdSig v
 1159   , not $ isTopSig strict_sig
 1160   , (ds, div) <- splitDmdSig strict_sig
 1161   , let arity = length ds
 1162   -- Every strictness signature admits an arity signature!
 1163   = AT (take arity one_shots) div
 1164   | otherwise
 1165   = AT (take (idArity v) one_shots) topDiv
 1166   where
 1167     one_shots :: [OneShotInfo]  -- One-shot-ness derived from the type
 1168     one_shots = typeArity (idType v)
 1169 
 1170 {-
 1171 %************************************************************************
 1172 %*                                                                      *
 1173               The main eta-expander
 1174 %*                                                                      *
 1175 %************************************************************************
 1176 
 1177 We go for:
 1178    f = \x1..xn -> N  ==>   f = \x1..xn y1..ym -> N y1..ym
 1179                                  (n >= 0)
 1180 
 1181 where (in both cases)
 1182 
 1183         * The xi can include type variables
 1184 
 1185         * The yi are all value variables
 1186 
 1187         * N is a NORMAL FORM (i.e. no redexes anywhere)
 1188           wanting a suitable number of extra args.
 1189 
 1190 The biggest reason for doing this is for cases like
 1191 
 1192         f = \x -> case x of
 1193                     True  -> \y -> e1
 1194                     False -> \y -> e2
 1195 
 1196 Here we want to get the lambdas together.  A good example is the nofib
 1197 program fibheaps, which gets 25% more allocation if you don't do this
 1198 eta-expansion.
 1199 
 1200 We may have to sandwich some coerces between the lambdas
 1201 to make the types work.   exprEtaExpandArity looks through coerces
 1202 when computing arity; and etaExpand adds the coerces as necessary when
 1203 actually computing the expansion.
 1204 
 1205 Note [No crap in eta-expanded code]
 1206 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1207 The eta expander is careful not to introduce "crap".  In particular,
 1208 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
 1209 returns a CoreExpr satisfying the same invariant. See Note [Eta
 1210 expansion and the CorePrep invariants] in CorePrep.
 1211 
 1212 This means the eta-expander has to do a bit of on-the-fly
 1213 simplification but it's not too hard.  The alternative, of relying on
 1214 a subsequent clean-up phase of the Simplifier to de-crapify the result,
 1215 means you can't really use it in CorePrep, which is painful.
 1216 
 1217 Note [Eta expansion for join points]
 1218 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1219 The no-crap rule is very tiresome to guarantee when
 1220 we have join points. Consider eta-expanding
 1221    let j :: Int -> Int -> Bool
 1222        j x = e
 1223    in b
 1224 
 1225 The simple way is
 1226   \(y::Int). (let j x = e in b) y
 1227 
 1228 The no-crap way is
 1229   \(y::Int). let j' :: Int -> Bool
 1230                  j' x = e y
 1231              in b[j'/j] y
 1232 where I have written to stress that j's type has
 1233 changed.  Note that (of course!) we have to push the application
 1234 inside the RHS of the join as well as into the body.  AND if j
 1235 has an unfolding we have to push it into there too.  AND j might
 1236 be recursive...
 1237 
 1238 So for now I'm abandoning the no-crap rule in this case. I think
 1239 that for the use in CorePrep it really doesn't matter; and if
 1240 it does, then CoreToStg.myCollectArgs will fall over.
 1241 
 1242 (Moreover, I think that casts can make the no-crap rule fail too.)
 1243 
 1244 Note [Eta expansion and SCCs]
 1245 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1246 Note that SCCs are not treated specially by etaExpand.  If we have
 1247         etaExpand 2 (\x -> scc "foo" e)
 1248         = (\xy -> (scc "foo" e) y)
 1249 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
 1250 
 1251 Note [Eta expansion and source notes]
 1252 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1253 CorePrep puts floatable ticks outside of value applications, but not
 1254 type applications. As a result we might be trying to eta-expand an
 1255 expression like
 1256 
 1257   (src<...> v) @a
 1258 
 1259 which we want to lead to code like
 1260 
 1261   \x -> src<...> v @a x
 1262 
 1263 This means that we need to look through type applications and be ready
 1264 to re-add floats on the top.
 1265 
 1266 Note [Eta expansion with ArityType]
 1267 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1268 The etaExpandAT function takes an ArityType (not just an Arity) to
 1269 guide eta-expansion.  Why? Because we want to preserve one-shot info.
 1270 Consider
 1271   foo = \x. case x of
 1272               True  -> (\s{os}. blah) |> co
 1273               False -> wubble
 1274 We'll get an ArityType for foo of \?1.T.
 1275 
 1276 Then we want to eta-expand to
 1277   foo = \x. (\eta{os}. (case x of ...as before...) eta) |> some_co
 1278 
 1279 That 'eta' binder is fresh, and we really want it to have the
 1280 one-shot flag from the inner \s{os}.  By expanding with the
 1281 ArityType gotten from analysing the RHS, we achieve this neatly.
 1282 
 1283 This makes a big difference to the one-shot monad trick;
 1284 see Note [The one-shot state monad trick] in GHC.Utils.Monad.
 1285 -}
 1286 
 1287 -- | @etaExpand n e@ returns an expression with
 1288 -- the same meaning as @e@, but with arity @n@.
 1289 --
 1290 -- Given:
 1291 --
 1292 -- > e' = etaExpand n e
 1293 --
 1294 -- We should have that:
 1295 --
 1296 -- > ty = exprType e = exprType e'
 1297 
 1298 etaExpand :: Arity -> CoreExpr -> CoreExpr
 1299 etaExpand n orig_expr
 1300   = eta_expand in_scope (replicate n NoOneShotInfo) orig_expr
 1301   where
 1302     in_scope = {-#SCC "eta_expand:in-scopeX" #-}
 1303                mkInScopeSet (exprFreeVars orig_expr)
 1304 
 1305 etaExpandAT :: InScopeSet -> ArityType -> CoreExpr -> CoreExpr
 1306 -- See Note [Eta expansion with ArityType]
 1307 --
 1308 -- We pass in the InScopeSet from the simplifier to avoid recomputing
 1309 -- it here, which can be jolly expensive if the casts are big
 1310 -- In #18223 it took 10% of compile time just to do the exprFreeVars!
 1311 etaExpandAT in_scope (AT oss _) orig_expr
 1312   = eta_expand in_scope oss orig_expr
 1313 
 1314 -- etaExpand arity e = res
 1315 -- Then 'res' has at least 'arity' lambdas at the top
 1316 --    possibly with a cast wrapped around the outside
 1317 -- See Note [Eta expansion with ArityType]
 1318 --
 1319 -- etaExpand deals with for-alls. For example:
 1320 --              etaExpand 1 E
 1321 -- where  E :: forall a. a -> a
 1322 -- would return
 1323 --      (/\b. \y::a -> E b y)
 1324 
 1325 eta_expand :: InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
 1326 eta_expand in_scope one_shots (Cast expr co)
 1327   = Cast (eta_expand in_scope one_shots expr) co
 1328 
 1329 eta_expand in_scope one_shots orig_expr
 1330   = go in_scope one_shots [] orig_expr
 1331   where
 1332       -- Strip off existing lambdas and casts before handing off to mkEtaWW
 1333       -- This is mainly to avoid spending time cloning binders and substituting
 1334       -- when there is actually nothing to do.  It's slightly awkward to deal
 1335       -- with casts here, apart from the topmost one, and they are rare, so
 1336       -- if we find one we just hand off to mkEtaWW anyway
 1337       -- Note [Eta expansion and SCCs]
 1338     go _ [] _ _ = orig_expr  -- Already has the specified arity; no-op
 1339 
 1340     go in_scope oss@(_:oss1) vs (Lam v body)
 1341       | isTyVar v = go (in_scope `extendInScopeSet` v) oss  (v:vs) body
 1342       | otherwise = go (in_scope `extendInScopeSet` v) oss1 (v:vs) body
 1343 
 1344     go in_scope oss rev_vs expr
 1345       = -- pprTrace "ee" (vcat [ppr in_scope', ppr top_bndrs, ppr eis]) $
 1346         retick $
 1347         etaInfoAbs top_eis $
 1348         etaInfoApp in_scope' sexpr eis
 1349       where
 1350           (in_scope', eis@(EI eta_bndrs mco))
 1351                     = mkEtaWW oss (ppr orig_expr) in_scope (exprType expr)
 1352           top_bndrs = reverse rev_vs
 1353           top_eis   = EI (top_bndrs ++ eta_bndrs) (mkPiMCos top_bndrs mco)
 1354 
 1355           -- Find ticks behind type apps.
 1356           -- See Note [Eta expansion and source notes]
 1357           -- I don't really understand this code SLPJ May 21
 1358           (expr', args) = collectArgs expr
 1359           (ticks, expr'') = stripTicksTop tickishFloatable expr'
 1360           sexpr = mkApps expr'' args
 1361           retick expr = foldr mkTick expr ticks
 1362 
 1363 {- *********************************************************************
 1364 *                                                                      *
 1365               The EtaInfo mechanism
 1366           mkEtaWW, etaInfoAbs, etaInfoApp
 1367 *                                                                      *
 1368 ********************************************************************* -}
 1369 
 1370 {- Note [The EtaInfo mechanism]
 1371 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1372 Suppose we have (e :: ty) and we want to eta-expand it to arity N.
 1373 This what eta_expand does.  We do it in two steps:
 1374 
 1375 1.  mkEtaWW: from 'ty' and 'N' build a EtaInfo which describes
 1376     the shape of the expansion necessary to expand to arity N.
 1377 
 1378 2.  Build the term
 1379        \ v1..vn.  e v1 .. vn
 1380     where those abstractions and applications are described by
 1381     the same EtaInfo.  Specifically we build the term
 1382 
 1383        etaInfoAbs etas (etaInfoApp in_scope e etas)
 1384 
 1385    where etas :: EtaInfo
 1386          etaInfoAbs builds the lambdas
 1387          etaInfoApp builds the applictions
 1388 
 1389    Note that the /same/ EtaInfo drives both etaInfoAbs and etaInfoApp
 1390 
 1391 To a first approximation EtaInfo is just [Var].  But
 1392 casts complicate the question.  If we have
 1393    newtype N a = MkN (S -> a)
 1394      axN :: N a  ~  S -> a
 1395 and
 1396    e :: N (N Int)
 1397 then the eta-expansion should look like
 1398    (\(x::S) (y::S) -> e |> co x y) |> sym co
 1399 where
 1400   co :: N (N Int) ~ S -> S -> Int
 1401   co = axN @(N Int) ; (S -> axN @Int)
 1402 
 1403 We want to get one cast, at the top, to account for all those
 1404 nested newtypes. This is expressed by the EtaInfo type:
 1405 
 1406    data EtaInfo = EI [Var] MCoercionR
 1407 
 1408 Note [Check for reflexive casts in eta expansion]
 1409 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1410 It turns out that the casts created by teh above mechanism are often Refl.
 1411 When casts are very deeply nested (as happens in #18223), the repetition
 1412 of types can make the overall term very large.  So there is a big
 1413 payoff in cancelling out casts aggressively wherever possible.
 1414 (See also Note [No crap in eta-expanded code].)
 1415 
 1416 This matters particularly in etaInfoApp, where we
 1417 * Do beta-reduction on the fly
 1418 * Use getArg_maybe to get a cast out of the way,
 1419   so that we can do beta reduction
 1420 Together this makes a big difference.  Consider when e is
 1421    case x of
 1422       True  -> (\x -> e1) |> c1
 1423       False -> (\p -> e2) |> c2
 1424 
 1425 When we eta-expand this to arity 1, say, etaInfoAbs will wrap
 1426 a (\eta) around the outside and use etaInfoApp to apply each
 1427 alternative to 'eta'.  We want to beta-reduce all that junk
 1428 away.
 1429 
 1430 #18223 was a dramatic example in which the intermediate term was
 1431 grotesquely huge, even though the next Simplifier iteration squashed
 1432 it.  Better to kill it at birth.
 1433 
 1434 The crucial spots in etaInfoApp are:
 1435 * `checkReflexiveMCo` in the (Cast e co) case of `go`
 1436 * `checkReflexiveMCo` in `pushCoArg`
 1437 * Less important: checkReflexiveMCo in the final case of `go`
 1438 Collectively these make a factor-of-5 difference to the total
 1439 allocation of T18223, so take care if you change this stuff!
 1440 
 1441 Example:
 1442    newtype N = MkN (Y->Z)
 1443    f :: X -> N
 1444    f = \(x::X). ((\(y::Y). blah) |> fco)
 1445 
 1446 where fco :: (Y->Z) ~ N
 1447 
 1448 mkEtaWW makes an EtaInfo of (EI [(eta1:X), (eta2:Y)] eta_co
 1449   where
 1450     eta_co :: (X->N) ~ (X->Y->Z)
 1451     eta_co =  (<X> -> nco)
 1452     nco :: N ~ (Y->Z)  -- Comes from topNormaliseNewType_maybe
 1453 
 1454 Now, when we push that eta_co inward in etaInfoApp:
 1455 * In the (Cast e co) case, the 'fco' and 'nco' will meet, and
 1456   should cancel.
 1457 * When we meet the (\y.e) we want no cast on the y.
 1458 
 1459 -}
 1460 
 1461 --------------
 1462 data EtaInfo = EI [Var] MCoercionR
 1463 
 1464 -- (EI bs co) describes a particular eta-expansion, as follows:
 1465 --  Abstraction:  (\b1 b2 .. bn. []) |> sym co
 1466 --  Application:  ([] |> co) b1 b2 .. bn
 1467 --
 1468 --    e :: T    co :: T ~ (t1 -> t2 -> .. -> tn -> tr)
 1469 --    e = (\b1 b2 ... bn. (e |> co) b1 b2 .. bn) |> sym co
 1470 
 1471 instance Outputable EtaInfo where
 1472   ppr (EI vs mco) = text "EI" <+> ppr vs <+> parens (ppr mco)
 1473 
 1474 
 1475 etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
 1476 -- (etaInfoApp s e (EI bs mco) returns something equivalent to
 1477 --             ((substExpr s e) |> mco b1 .. bn)
 1478 -- See Note [The EtaInfo mechanism]
 1479 --
 1480 -- NB: With very deeply nested casts, this function can be expensive
 1481 --     In T18223, this function alone costs 15% of allocation, all
 1482 --     spent in the calls to substExprSC and substBindSC
 1483 
 1484 etaInfoApp in_scope expr eis
 1485   = go (mkEmptySubst in_scope) expr eis
 1486   where
 1487     go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
 1488     -- 'go' pushed down the eta-infos into the branch of a case
 1489     -- and the body of a let; and does beta-reduction if possible
 1490     --   go subst fun co [b1,..,bn]  returns  (subst(fun) |> co) b1 .. bn
 1491     go subst (Tick t e) eis
 1492       = Tick (substTickish subst t) (go subst e eis)
 1493 
 1494     go subst (Cast e co) (EI bs mco)
 1495       = go subst e (EI bs mco')
 1496       where
 1497         mco' = checkReflexiveMCo (Core.substCo subst co `mkTransMCoR` mco)
 1498                -- See Note [Check for reflexive casts in eta-expansion]
 1499 
 1500     go subst (Case e b ty alts) eis
 1501       = Case (Core.substExprSC subst e) b1 ty' alts'
 1502       where
 1503         (subst1, b1) = Core.substBndr subst b
 1504         alts' = map subst_alt alts
 1505         ty'   = etaInfoAppTy (Core.substTy subst ty) eis
 1506         subst_alt (Alt con bs rhs) = Alt con bs' (go subst2 rhs eis)
 1507                  where
 1508                   (subst2,bs') = Core.substBndrs subst1 bs
 1509 
 1510     go subst (Let b e) eis
 1511       | not (isJoinBind b) -- See Note [Eta expansion for join points]
 1512       = Let b' (go subst' e eis)
 1513       where
 1514         (subst', b') = Core.substBindSC subst b
 1515 
 1516     -- Beta-reduction if possible, pushing any intervening casts past
 1517     -- the argument. See Note [The EtaInfo mechansim]
 1518     go subst (Lam v e) (EI (b:bs) mco)
 1519       | Just (arg,mco') <- pushMCoArg mco (varToCoreExpr b)
 1520       = go (Core.extendSubst subst v arg) e (EI bs mco')
 1521 
 1522     -- Stop pushing down; just wrap the expression up
 1523     -- See Note [Check for reflexive casts in eta expansion]
 1524     go subst e (EI bs mco) = Core.substExprSC subst e
 1525                              `mkCastMCo` checkReflexiveMCo mco
 1526                              `mkVarApps` bs
 1527 
 1528 --------------
 1529 etaInfoAppTy :: Type -> EtaInfo -> Type
 1530 -- If                    e :: ty
 1531 -- then   etaInfoApp e eis :: etaInfoApp ty eis
 1532 etaInfoAppTy ty (EI bs mco)
 1533   = applyTypeToArgs (text "etaInfoAppTy") ty1 (map varToCoreExpr bs)
 1534   where
 1535     ty1 = case mco of
 1536              MRefl  -> ty
 1537              MCo co -> coercionRKind co
 1538 
 1539 --------------
 1540 etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
 1541 -- See Note [The EtaInfo mechanism]
 1542 etaInfoAbs (EI bs mco) expr = (mkLams bs expr) `mkCastMCo` mkSymMCo mco
 1543 
 1544 --------------
 1545 -- | @mkEtaWW n _ fvs ty@ will compute the 'EtaInfo' necessary for eta-expanding
 1546 -- an expression @e :: ty@ to take @n@ value arguments, where @fvs@ are the
 1547 -- free variables of @e@.
 1548 --
 1549 -- Note that this function is entirely unconcerned about cost centres and other
 1550 -- semantically-irrelevant source annotations, so call sites must take care to
 1551 -- preserve that info. See Note [Eta expansion and SCCs].
 1552 mkEtaWW
 1553   :: [OneShotInfo]
 1554   -- ^ How many value arguments to eta-expand
 1555   -> SDoc
 1556   -- ^ The pretty-printed original expression, for warnings.
 1557   -> InScopeSet
 1558   -- ^ A super-set of the free vars of the expression to eta-expand.
 1559   -> Type
 1560   -> (InScopeSet, EtaInfo)
 1561   -- ^ The variables in 'EtaInfo' are fresh wrt. to the incoming 'InScopeSet'.
 1562   -- The outgoing 'InScopeSet' extends the incoming 'InScopeSet' with the
 1563   -- fresh variables in 'EtaInfo'.
 1564 
 1565 mkEtaWW orig_oss ppr_orig_expr in_scope orig_ty
 1566   = go 0 orig_oss empty_subst orig_ty
 1567   where
 1568     empty_subst = mkEmptyTCvSubst in_scope
 1569 
 1570     go :: Int                -- For fresh names
 1571        -> [OneShotInfo]      -- Number of value args to expand to
 1572        -> TCvSubst -> Type   -- We are really looking at subst(ty)
 1573        -> (InScopeSet, EtaInfo)
 1574     -- (go [o1,..,on] subst ty) = (in_scope, EI [b1,..,bn] co)
 1575     --    co :: subst(ty) ~ b1_ty -> ... -> bn_ty -> tr
 1576 
 1577     go _ [] subst _       -- See Note [exprArity invariant]
 1578        ----------- Done!  No more expansion needed
 1579        = (getTCvInScope subst, EI [] MRefl)
 1580 
 1581     go n oss@(one_shot:oss1) subst ty       -- See Note [exprArity invariant]
 1582        ----------- Forall types  (forall a. ty)
 1583        | Just (tcv,ty') <- splitForAllTyCoVar_maybe ty
 1584        , (subst', tcv') <- Type.substVarBndr subst tcv
 1585        , let oss' | isTyVar tcv = oss
 1586                   | otherwise   = oss1
 1587          -- A forall can bind a CoVar, in which case
 1588          -- we consume one of the [OneShotInfo]
 1589        , (in_scope, EI bs mco) <- go n oss' subst' ty'
 1590        = (in_scope, EI (tcv' : bs) (mkHomoForAllMCo tcv' mco))
 1591 
 1592        ----------- Function types  (t1 -> t2)
 1593        | Just (mult, arg_ty, res_ty) <- splitFunTy_maybe ty
 1594        , typeHasFixedRuntimeRep arg_ty
 1595           -- See Note [Representation polymorphism invariants] in GHC.Core
 1596           -- See also test case typecheck/should_run/EtaExpandLevPoly
 1597 
 1598        , (subst', eta_id) <- freshEtaId n subst (Scaled mult arg_ty)
 1599           -- Avoid free vars of the original expression
 1600 
 1601        , let eta_id' = eta_id `setIdOneShotInfo` one_shot
 1602        , (in_scope, EI bs mco) <- go (n+1) oss1 subst' res_ty
 1603        = (in_scope, EI (eta_id' : bs) (mkFunResMCo (idScaledType eta_id') mco))
 1604 
 1605        ----------- Newtypes
 1606        -- Given this:
 1607        --      newtype T = MkT ([T] -> Int)
 1608        -- Consider eta-expanding this
 1609        --      eta_expand 1 e T
 1610        -- We want to get
 1611        --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 1612        | Just (co, ty') <- topNormaliseNewType_maybe ty
 1613        , -- co :: ty ~ ty'
 1614          let co' = Type.substCo subst co
 1615              -- Remember to apply the substitution to co (#16979)
 1616              -- (or we could have applied to ty, but then
 1617              --  we'd have had to zap it for the recursive call)
 1618        , (in_scope, EI bs mco) <- go n oss subst ty'
 1619          -- mco :: subst(ty') ~ b1_ty -> ... -> bn_ty -> tr
 1620        = (in_scope, EI bs (mkTransMCoR co' mco))
 1621 
 1622        | otherwise       -- We have an expression of arity > 0,
 1623                          -- but its type isn't a function, or a binder
 1624                          -- does not have a fixed runtime representation
 1625        = warnPprTrace True ((ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr)
 1626          (getTCvInScope subst, EI [] MRefl)
 1627         -- This *can* legitimately happen:
 1628         -- e.g.  coerce Int (\x. x) Essentially the programmer is
 1629         -- playing fast and loose with types (Happy does this a lot).
 1630         -- So we simply decline to eta-expand.  Otherwise we'd end up
 1631         -- with an explicit lambda having a non-function type
 1632 
 1633 
 1634 {- *********************************************************************
 1635 *                                                                      *
 1636               The "push rules"
 1637 *                                                                      *
 1638 ************************************************************************
 1639 
 1640 Here we implement the "push rules" from FC papers:
 1641 
 1642 * The push-argument rules, where we can move a coercion past an argument.
 1643   We have
 1644       (fun |> co) arg
 1645   and we want to transform it to
 1646     (fun arg') |> co'
 1647   for some suitable co' and transformed arg'.
 1648 
 1649 * The PushK rule for data constructors.  We have
 1650        (K e1 .. en) |> co
 1651   and we want to transform to
 1652        (K e1' .. en')
 1653   by pushing the coercion into the arguments
 1654 -}
 1655 
 1656 pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], MCoercion)
 1657 pushCoArgs co []         = return ([], MCo co)
 1658 pushCoArgs co (arg:args) = do { (arg',  m_co1) <- pushCoArg  co  arg
 1659                               ; case m_co1 of
 1660                                   MCo co1 -> do { (args', m_co2) <- pushCoArgs co1 args
 1661                                                  ; return (arg':args', m_co2) }
 1662                                   MRefl  -> return (arg':args, MRefl) }
 1663 
 1664 pushMCoArg :: MCoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
 1665 pushMCoArg MRefl    arg = Just (arg, MRefl)
 1666 pushMCoArg (MCo co) arg = pushCoArg co arg
 1667 
 1668 pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
 1669 -- We have (fun |> co) arg, and we want to transform it to
 1670 --         (fun arg) |> co
 1671 -- This may fail, e.g. if (fun :: N) where N is a newtype
 1672 -- C.f. simplCast in GHC.Core.Opt.Simplify
 1673 -- 'co' is always Representational
 1674 pushCoArg co arg
 1675   | Type ty <- arg
 1676   = do { (ty', m_co') <- pushCoTyArg co ty
 1677        ; return (Type ty', m_co') }
 1678   | otherwise
 1679   = do { (arg_mco, m_co') <- pushCoValArg co
 1680        ; let arg_mco' = checkReflexiveMCo arg_mco
 1681              -- checkReflexiveMCo: see Note [Check for reflexive casts in eta expansion]
 1682              -- The coercion is very often (arg_co -> res_co), but without
 1683              -- the argument coercion actually being ReflCo
 1684        ; return (arg `mkCastMCo` arg_mco', m_co') }
 1685 
 1686 pushCoTyArg :: CoercionR -> Type -> Maybe (Type, MCoercionR)
 1687 -- We have (fun |> co) @ty
 1688 -- Push the coercion through to return
 1689 --         (fun @ty') |> co'
 1690 -- 'co' is always Representational
 1691 -- If the returned coercion is Nothing, then it would have been reflexive;
 1692 -- it's faster not to compute it, though.
 1693 pushCoTyArg co ty
 1694   -- The following is inefficient - don't do `eqType` here, the coercion
 1695   -- optimizer will take care of it. See #14737.
 1696   -- -- | tyL `eqType` tyR
 1697   -- -- = Just (ty, Nothing)
 1698 
 1699   | isReflCo co
 1700   = Just (ty, MRefl)
 1701 
 1702   | isForAllTy_ty tyL
 1703   = assertPpr (isForAllTy_ty tyR) (ppr co $$ ppr ty) $
 1704     Just (ty `mkCastTy` co1, MCo co2)
 1705 
 1706   | otherwise
 1707   = Nothing
 1708   where
 1709     Pair tyL tyR = coercionKind co
 1710        -- co :: tyL ~ tyR
 1711        -- tyL = forall (a1 :: k1). ty1
 1712        -- tyR = forall (a2 :: k2). ty2
 1713 
 1714     co1 = mkSymCo (mkNthCo Nominal 0 co)
 1715        -- co1 :: k2 ~N k1
 1716        -- Note that NthCo can extract a Nominal equality between the
 1717        -- kinds of the types related by a coercion between forall-types.
 1718        -- See the NthCo case in GHC.Core.Lint.
 1719 
 1720     co2 = mkInstCo co (mkGReflLeftCo Nominal ty co1)
 1721         -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
 1722         -- Arg of mkInstCo is always nominal, hence mkNomReflCo
 1723 
 1724 pushCoValArg :: CoercionR -> Maybe (MCoercionR, MCoercionR)
 1725 -- We have (fun |> co) arg
 1726 -- Push the coercion through to return
 1727 --         (fun (arg |> co_arg)) |> co_res
 1728 -- 'co' is always Representational
 1729 -- If the second returned Coercion is actually Nothing, then no cast is necessary;
 1730 -- the returned coercion would have been reflexive.
 1731 pushCoValArg co
 1732   -- The following is inefficient - don't do `eqType` here, the coercion
 1733   -- optimizer will take care of it. See #14737.
 1734   -- -- | tyL `eqType` tyR
 1735   -- -- = Just (mkRepReflCo arg, Nothing)
 1736 
 1737   | isReflCo co
 1738   = Just (MRefl, MRefl)
 1739 
 1740   | isFunTy tyL
 1741   , (co_mult, co1, co2) <- decomposeFunCo Representational co
 1742   , isReflexiveCo co_mult
 1743     -- We can't push the coercion in the case where co_mult isn't reflexivity:
 1744     -- it could be an unsafe axiom, and losing this information could yield
 1745     -- ill-typed terms. For instance (fun x ::(1) Int -> (fun _ -> () |> co) x)
 1746     -- with co :: (Int -> ()) ~ (Int %1 -> ()), would reduce to (fun x ::(1) Int
 1747     -- -> (fun _ ::(Many) Int -> ()) x) which is ill-typed
 1748 
 1749               -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
 1750               -- then co1 :: tyL1 ~ tyR1
 1751               --      co2 :: tyL2 ~ tyR2
 1752   = assertPpr (isFunTy tyR) (ppr co $$ ppr arg) $
 1753     Just (coToMCo (mkSymCo co1), coToMCo co2)
 1754     -- Critically, coToMCo to checks for ReflCo; the whole coercion may not
 1755     -- be reflexive, but either of its components might be
 1756     -- We could use isReflexiveCo, but it's not clear if the benefit
 1757     -- is worth the cost, and it makes no difference in #18223
 1758 
 1759   | otherwise
 1760   = Nothing
 1761   where
 1762     arg = funArgTy tyR
 1763     Pair tyL tyR = coercionKind co
 1764 
 1765 pushCoercionIntoLambda
 1766     :: HasDebugCallStack => InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
 1767 -- This implements the Push rule from the paper on coercions
 1768 --    (\x. e) |> co
 1769 -- ===>
 1770 --    (\x'. e |> co')
 1771 pushCoercionIntoLambda in_scope x e co
 1772     | assert (not (isTyVar x) && not (isCoVar x)) True
 1773     , Pair s1s2 t1t2 <- coercionKind co
 1774     , Just (_, _s1,_s2) <- splitFunTy_maybe s1s2
 1775     , Just (w1, t1,_t2) <- splitFunTy_maybe t1t2
 1776     , (co_mult, co1, co2) <- decomposeFunCo Representational co
 1777     , isReflexiveCo co_mult
 1778       -- We can't push the coercion in the case where co_mult isn't
 1779       -- reflexivity. See pushCoValArg for more details.
 1780     = let
 1781           -- Should we optimize the coercions here?
 1782           -- Otherwise they might not match too well
 1783           x' = x `setIdType` t1 `setIdMult` w1
 1784           in_scope' = in_scope `extendInScopeSet` x'
 1785           subst = extendIdSubst (mkEmptySubst in_scope')
 1786                                 x
 1787                                 (mkCast (Var x') co1)
 1788       in Just (x', substExpr subst e `mkCast` co2)
 1789     | otherwise
 1790     = pprTrace "exprIsLambda_maybe: Unexpected lambda in case" (ppr (Lam x e))
 1791       Nothing
 1792 
 1793 pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
 1794               -> Maybe (DataCon
 1795                        , [Type]      -- Universal type args
 1796                        , [CoreExpr]) -- All other args incl existentials
 1797 -- Implement the KPush reduction rule as described in "Down with kinds"
 1798 -- The transformation applies iff we have
 1799 --      (C e1 ... en) `cast` co
 1800 -- where co :: (T t1 .. tn) ~ to_ty
 1801 -- The left-hand one must be a T, because exprIsConApp returned True
 1802 -- but the right-hand one might not be.  (Though it usually will.)
 1803 pushCoDataCon dc dc_args co
 1804   | isReflCo co || from_ty `eqType` to_ty  -- try cheap test first
 1805   , let (univ_ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) dc_args
 1806   = Just (dc, map exprToType univ_ty_args, rest_args)
 1807 
 1808   | Just (to_tc, to_tc_arg_tys) <- splitTyConApp_maybe to_ty
 1809   , to_tc == dataConTyCon dc
 1810         -- These two tests can fail; we might see
 1811         --      (C x y) `cast` (g :: T a ~ S [a]),
 1812         -- where S is a type function.  In fact, exprIsConApp
 1813         -- will probably not be called in such circumstances,
 1814         -- but there's nothing wrong with it
 1815 
 1816   = let
 1817         tc_arity       = tyConArity to_tc
 1818         dc_univ_tyvars = dataConUnivTyVars dc
 1819         dc_ex_tcvars   = dataConExTyCoVars dc
 1820         arg_tys        = dataConRepArgTys dc
 1821 
 1822         non_univ_args  = dropList dc_univ_tyvars dc_args
 1823         (ex_args, val_args) = splitAtList dc_ex_tcvars non_univ_args
 1824 
 1825         -- Make the "Psi" from the paper
 1826         omegas = decomposeCo tc_arity co (tyConRolesRepresentational to_tc)
 1827         (psi_subst, to_ex_arg_tys)
 1828           = liftCoSubstWithEx Representational
 1829                               dc_univ_tyvars
 1830                               omegas
 1831                               dc_ex_tcvars
 1832                               (map exprToType ex_args)
 1833 
 1834           -- Cast the value arguments (which include dictionaries)
 1835         new_val_args = zipWith cast_arg (map scaledThing arg_tys) val_args
 1836         cast_arg arg_ty arg = mkCast arg (psi_subst arg_ty)
 1837 
 1838         to_ex_args = map Type to_ex_arg_tys
 1839 
 1840         dump_doc = vcat [ppr dc,      ppr dc_univ_tyvars, ppr dc_ex_tcvars,
 1841                          ppr arg_tys, ppr dc_args,
 1842                          ppr ex_args, ppr val_args, ppr co, ppr from_ty, ppr to_ty, ppr to_tc
 1843                          , ppr $ mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args) ]
 1844     in
 1845     assertPpr (eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args))) dump_doc $
 1846     assertPpr (equalLength val_args arg_tys) dump_doc $
 1847     Just (dc, to_tc_arg_tys, to_ex_args ++ new_val_args)
 1848 
 1849   | otherwise
 1850   = Nothing
 1851 
 1852   where
 1853     Pair from_ty to_ty = coercionKind co
 1854 
 1855 collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
 1856 -- Collect lambda binders, pushing coercions inside if possible
 1857 -- E.g.   (\x.e) |> g         g :: <Int> -> blah
 1858 --        = (\x. e |> Nth 1 g)
 1859 --
 1860 -- That is,
 1861 --
 1862 -- collectBindersPushingCo ((\x.e) |> g) === ([x], e |> Nth 1 g)
 1863 collectBindersPushingCo e
 1864   = go [] e
 1865   where
 1866     -- Peel off lambdas until we hit a cast.
 1867     go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
 1868     -- The accumulator is in reverse order
 1869     go bs (Lam b e)   = go (b:bs) e
 1870     go bs (Cast e co) = go_c bs e co
 1871     go bs e           = (reverse bs, e)
 1872 
 1873     -- We are in a cast; peel off casts until we hit a lambda.
 1874     go_c :: [Var] -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
 1875     -- (go_c bs e c) is same as (go bs e (e |> c))
 1876     go_c bs (Cast e co1) co2 = go_c bs e (co1 `mkTransCo` co2)
 1877     go_c bs (Lam b e)    co  = go_lam bs b e co
 1878     go_c bs e            co  = (reverse bs, mkCast e co)
 1879 
 1880     -- We are in a lambda under a cast; peel off lambdas and build a
 1881     -- new coercion for the body.
 1882     go_lam :: [Var] -> Var -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
 1883     -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
 1884     go_lam bs b e co
 1885       | isTyVar b
 1886       , let Pair tyL tyR = coercionKind co
 1887       , assert (isForAllTy_ty tyL) $
 1888         isForAllTy_ty tyR
 1889       , isReflCo (mkNthCo Nominal 0 co)  -- See Note [collectBindersPushingCo]
 1890       = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkTyVarTy b)))
 1891 
 1892       | isCoVar b
 1893       , let Pair tyL tyR = coercionKind co
 1894       , assert (isForAllTy_co tyL) $
 1895         isForAllTy_co tyR
 1896       , isReflCo (mkNthCo Nominal 0 co)  -- See Note [collectBindersPushingCo]
 1897       , let cov = mkCoVarCo b
 1898       = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkCoercionTy cov)))
 1899 
 1900       | isId b
 1901       , let Pair tyL tyR = coercionKind co
 1902       , assert (isFunTy tyL) $ isFunTy tyR
 1903       , (co_mult, co_arg, co_res) <- decomposeFunCo Representational co
 1904       , isReflCo co_mult -- See Note [collectBindersPushingCo]
 1905       , isReflCo co_arg  -- See Note [collectBindersPushingCo]
 1906       = go_c (b:bs) e co_res
 1907 
 1908       | otherwise = (reverse bs, mkCast (Lam b e) co)
 1909 
 1910 {-
 1911 
 1912 Note [collectBindersPushingCo]
 1913 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 1914 We just look for coercions of form
 1915    <type> # w -> blah
 1916 (and similarly for foralls) to keep this function simple.  We could do
 1917 more elaborate stuff, but it'd involve substitution etc.
 1918 
 1919 -}
 1920 
 1921 {- *********************************************************************
 1922 *                                                                      *
 1923                 Join points
 1924 *                                                                      *
 1925 ********************************************************************* -}
 1926 
 1927 -------------------
 1928 -- | Split an expression into the given number of binders and a body,
 1929 -- eta-expanding if necessary. Counts value *and* type binders.
 1930 etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr)
 1931 etaExpandToJoinPoint join_arity expr
 1932   = go join_arity [] expr
 1933   where
 1934     go 0 rev_bs e         = (reverse rev_bs, e)
 1935     go n rev_bs (Lam b e) = go (n-1) (b : rev_bs) e
 1936     go n rev_bs e         = case etaBodyForJoinPoint n e of
 1937                               (bs, e') -> (reverse rev_bs ++ bs, e')
 1938 
 1939 etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
 1940 etaExpandToJoinPointRule _ rule@(BuiltinRule {})
 1941   = warnPprTrace True (sep [text "Can't eta-expand built-in rule:", ppr rule])
 1942       -- How did a local binding get a built-in rule anyway? Probably a plugin.
 1943     rule
 1944 etaExpandToJoinPointRule join_arity rule@(Rule { ru_bndrs = bndrs, ru_rhs = rhs
 1945                                                , ru_args  = args })
 1946   | need_args == 0
 1947   = rule
 1948   | need_args < 0
 1949   = pprPanic "etaExpandToJoinPointRule" (ppr join_arity $$ ppr rule)
 1950   | otherwise
 1951   = rule { ru_bndrs = bndrs ++ new_bndrs, ru_args = args ++ new_args
 1952          , ru_rhs = new_rhs }
 1953   where
 1954     need_args = join_arity - length args
 1955     (new_bndrs, new_rhs) = etaBodyForJoinPoint need_args rhs
 1956     new_args = varsToCoreExprs new_bndrs
 1957 
 1958 -- Adds as many binders as asked for; assumes expr is not a lambda
 1959 etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr)
 1960 etaBodyForJoinPoint need_args body
 1961   = go need_args (exprType body) (init_subst body) [] body
 1962   where
 1963     go 0 _  _     rev_bs e
 1964       = (reverse rev_bs, e)
 1965     go n ty subst rev_bs e
 1966       | Just (tv, res_ty) <- splitForAllTyCoVar_maybe ty
 1967       , let (subst', tv') = substVarBndr subst tv
 1968       = go (n-1) res_ty subst' (tv' : rev_bs) (e `App` varToCoreExpr tv')
 1969       | Just (mult, arg_ty, res_ty) <- splitFunTy_maybe ty
 1970       , let (subst', b) = freshEtaId n subst (Scaled mult arg_ty)
 1971       = go (n-1) res_ty subst' (b : rev_bs) (e `App` Var b)
 1972       | otherwise
 1973       = pprPanic "etaBodyForJoinPoint" $ int need_args $$
 1974                                          ppr body $$ ppr (exprType body)
 1975 
 1976     init_subst e = mkEmptyTCvSubst (mkInScopeSet (exprFreeVars e))
 1977 
 1978 
 1979 
 1980 --------------
 1981 freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
 1982 -- Make a fresh Id, with specified type (after applying substitution)
 1983 -- It should be "fresh" in the sense that it's not in the in-scope set
 1984 -- of the TvSubstEnv; and it should itself then be added to the in-scope
 1985 -- set of the TvSubstEnv
 1986 --
 1987 -- The Int is just a reasonable starting point for generating a unique;
 1988 -- it does not necessarily have to be unique itself.
 1989 freshEtaId n subst ty
 1990       = (subst', eta_id')
 1991       where
 1992         Scaled mult' ty' = Type.substScaledTyUnchecked subst ty
 1993         eta_id' = uniqAway (getTCvInScope subst) $
 1994                   mkSysLocalOrCoVar (fsLit "eta") (mkBuiltinUnique n) mult' ty'
 1995                   -- "OrCoVar" since this can be used to eta-expand
 1996                   -- coercion abstractions
 1997         subst'  = extendTCvInScope subst eta_id'