{- (c) The University of Glasgow 2006 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998 Arity and eta expansion -} {-# LANGUAGE CPP #-} {-# OPTIONS_GHC -Wno-incomplete-record-updates #-} -- | Arity and eta expansion module GHC.Core.Opt.Arity ( manifestArity, joinRhsArity, exprArity, typeArity , exprEtaExpandArity, findRhsArity , etaExpand, etaExpandAT , etaExpandToJoinPoint, etaExpandToJoinPointRule , exprBotStrictness_maybe , ArityType(..), expandableArityType, arityTypeArity , maxWithArity, isBotArityType, idArityType ) where #include "HsVersions.h" import GHC.Prelude import GHC.Core import GHC.Core.FVs import GHC.Core.Utils import GHC.Core.Subst import GHC.Types.Demand import GHC.Types.Var import GHC.Types.Var.Env import GHC.Types.Id import GHC.Core.Type as Type import GHC.Core.TyCon ( initRecTc, checkRecTc ) import GHC.Core.Predicate ( isDictTy ) import GHC.Core.Coercion as Coercion import GHC.Core.Multiplicity import GHC.Types.Var.Set import GHC.Types.Basic import GHC.Types.Unique import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt ) import GHC.Utils.Outputable import GHC.Data.FastString import GHC.Utils.Misc ( lengthAtLeast ) {- ************************************************************************ * * manifestArity and exprArity * * ************************************************************************ exprArity is a cheap-and-cheerful version of exprEtaExpandArity. It tells how many things the expression can be applied to before doing any work. It doesn't look inside cases, lets, etc. The idea is that exprEtaExpandArity will do the hard work, leaving something that's easy for exprArity to grapple with. In particular, Simplify uses exprArity to compute the ArityInfo for the Id. Originally I thought that it was enough just to look for top-level lambdas, but it isn't. I've seen this foo = PrelBase.timesInt We want foo to get arity 2 even though the eta-expander will leave it unchanged, in the expectation that it'll be inlined. But occasionally it isn't, because foo is blacklisted (used in a rule). Similarly, see the ok_note check in exprEtaExpandArity. So f = __inline_me (\x -> e) won't be eta-expanded. And in any case it seems more robust to have exprArity be a bit more intelligent. But note that (\x y z -> f x y z) should have arity 3, regardless of f's arity. -} manifestArity :: CoreExpr -> Arity -- ^ manifestArity sees how many leading value lambdas there are, -- after looking through casts manifestArity :: CoreExpr -> JoinArity manifestArity (Lam Id v CoreExpr e) | Id -> Bool isId Id v = JoinArity 1 JoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a + CoreExpr -> JoinArity manifestArity CoreExpr e | Bool otherwise = CoreExpr -> JoinArity manifestArity CoreExpr e manifestArity (Tick Tickish Id t CoreExpr e) | Bool -> Bool not (Tickish Id -> Bool forall id. Tickish id -> Bool tickishIsCode Tickish Id t) = CoreExpr -> JoinArity manifestArity CoreExpr e manifestArity (Cast CoreExpr e CoercionR _) = CoreExpr -> JoinArity manifestArity CoreExpr e manifestArity CoreExpr _ = JoinArity 0 joinRhsArity :: CoreExpr -> JoinArity -- Join points are supposed to have manifestly-visible -- lambdas at the top: no ticks, no casts, nothing -- Moreover, type lambdas count in JoinArity joinRhsArity :: CoreExpr -> JoinArity joinRhsArity (Lam Id _ CoreExpr e) = JoinArity 1 JoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a + CoreExpr -> JoinArity joinRhsArity CoreExpr e joinRhsArity CoreExpr _ = JoinArity 0 --------------- exprArity :: CoreExpr -> Arity -- ^ An approximate, fast, version of 'exprEtaExpandArity' exprArity :: CoreExpr -> JoinArity exprArity CoreExpr e = CoreExpr -> JoinArity go CoreExpr e where go :: CoreExpr -> JoinArity go (Var Id v) = Id -> JoinArity idArity Id v go (Lam Id x CoreExpr e) | Id -> Bool isId Id x = CoreExpr -> JoinArity go CoreExpr e JoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a + JoinArity 1 | Bool otherwise = CoreExpr -> JoinArity go CoreExpr e go (Tick Tickish Id t CoreExpr e) | Bool -> Bool not (Tickish Id -> Bool forall id. Tickish id -> Bool tickishIsCode Tickish Id t) = CoreExpr -> JoinArity go CoreExpr e go (Cast CoreExpr e CoercionR co) = JoinArity -> Type -> JoinArity trim_arity (CoreExpr -> JoinArity go CoreExpr e) (CoercionR -> Type coercionRKind CoercionR co) -- Note [exprArity invariant] go (App CoreExpr e (Type Type _)) = CoreExpr -> JoinArity go CoreExpr e go (App CoreExpr f CoreExpr a) | CoreExpr -> Bool exprIsTrivial CoreExpr a = (CoreExpr -> JoinArity go CoreExpr f JoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a - JoinArity 1) JoinArity -> JoinArity -> JoinArity forall a. Ord a => a -> a -> a `max` JoinArity 0 -- See Note [exprArity for applications] -- NB: coercions count as a value argument go CoreExpr _ = JoinArity 0 trim_arity :: Arity -> Type -> Arity trim_arity :: JoinArity -> Type -> JoinArity trim_arity JoinArity arity Type ty = JoinArity arity JoinArity -> JoinArity -> JoinArity forall a. Ord a => a -> a -> a `min` [OneShotInfo] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length (Type -> [OneShotInfo] typeArity Type ty) --------------- typeArity :: Type -> [OneShotInfo] -- How many value arrows are visible in the type? -- We look through foralls, and newtypes -- See Note [exprArity invariant] typeArity :: Type -> [OneShotInfo] typeArity Type ty = RecTcChecker -> Type -> [OneShotInfo] go RecTcChecker initRecTc Type ty where go :: RecTcChecker -> Type -> [OneShotInfo] go RecTcChecker rec_nts Type ty | Just (Id _, Type ty') <- Type -> Maybe (Id, Type) splitForAllTy_maybe Type ty = RecTcChecker -> Type -> [OneShotInfo] go RecTcChecker rec_nts Type ty' | Just (Type _,Type arg,Type res) <- Type -> Maybe (Type, Type, Type) splitFunTy_maybe Type ty = Type -> OneShotInfo typeOneShot Type arg OneShotInfo -> [OneShotInfo] -> [OneShotInfo] forall a. a -> [a] -> [a] : RecTcChecker -> Type -> [OneShotInfo] go RecTcChecker rec_nts Type res | Just (TyCon tc,[Type] tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type]) Type -> Maybe (TyCon, [Type]) splitTyConApp_maybe Type ty , Just (Type ty', CoercionR _) <- TyCon -> [Type] -> Maybe (Type, CoercionR) instNewTyCon_maybe TyCon tc [Type] tys , Just RecTcChecker rec_nts' <- RecTcChecker -> TyCon -> Maybe RecTcChecker checkRecTc RecTcChecker rec_nts TyCon tc -- See Note [Expanding newtypes] -- in GHC.Core.TyCon -- , not (isClassTyCon tc) -- Do not eta-expand through newtype classes -- -- See Note [Newtype classes and eta expansion] -- (no longer required) = RecTcChecker -> Type -> [OneShotInfo] go RecTcChecker rec_nts' Type ty' -- Important to look through non-recursive newtypes, so that, eg -- (f x) where f has arity 2, f :: Int -> IO () -- Here we want to get arity 1 for the result! -- -- AND through a layer of recursive newtypes -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b))) | Bool otherwise = [] --------------- exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig) -- A cheap and cheerful function that identifies bottoming functions -- and gives them a suitable strictness signatures. It's used during -- float-out exprBotStrictness_maybe :: CoreExpr -> Maybe (JoinArity, StrictSig) exprBotStrictness_maybe CoreExpr e = case ArityType -> Maybe JoinArity getBotArity (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e) of Maybe JoinArity Nothing -> Maybe (JoinArity, StrictSig) forall a. Maybe a Nothing Just JoinArity ar -> (JoinArity, StrictSig) -> Maybe (JoinArity, StrictSig) forall a. a -> Maybe a Just (JoinArity ar, JoinArity -> StrictSig sig JoinArity ar) where env :: ArityEnv env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv AE { ae_ped_bot :: Bool ae_ped_bot = Bool True , ae_cheap_fn :: CheapFun ae_cheap_fn = \ CoreExpr _ Maybe Type _ -> Bool False , ae_joins :: IdSet ae_joins = IdSet emptyVarSet } sig :: JoinArity -> StrictSig sig JoinArity ar = [Demand] -> Divergence -> StrictSig mkClosedStrictSig (JoinArity -> Demand -> [Demand] forall a. JoinArity -> a -> [a] replicate JoinArity ar Demand topDmd) Divergence botDiv {- Note [exprArity invariant] ~~~~~~~~~~~~~~~~~~~~~~~~~~ exprArity has the following invariants: (1) If typeArity (exprType e) = n, then manifestArity (etaExpand e n) = n That is, etaExpand can always expand as much as typeArity says So the case analysis in etaExpand and in typeArity must match (2) exprArity e <= typeArity (exprType e) (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n That is, if exprArity says "the arity is n" then etaExpand really can get "n" manifest lambdas to the top. Why is this important? Because - In GHC.Iface.Tidy we use exprArity to fix the *final arity* of each top-level Id, and in - In CorePrep we use etaExpand on each rhs, so that the visible lambdas actually match that arity, which in turn means that the StgRhs has the right number of lambdas An alternative would be to do the eta-expansion in GHC.Iface.Tidy, at least for top-level bindings, in which case we would not need the trim_arity in exprArity. That is a less local change, so I'm going to leave it for today! Note [Newtype classes and eta expansion] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ NB: this nasty special case is no longer required, because for newtype classes we don't use the class-op rule mechanism at all. See Note [Single-method classes] in GHC.Tc.TyCl.Instance. SLPJ May 2013 -------- Old out of date comments, just for interest ----------- We have to be careful when eta-expanding through newtypes. In general it's a good idea, but annoyingly it interacts badly with the class-op rule mechanism. Consider class C a where { op :: a -> a } instance C b => C [b] where op x = ... These translate to co :: forall a. (a->a) ~ C a $copList :: C b -> [b] -> [b] $copList d x = ... $dfList :: C b -> C [b] {-# DFunUnfolding = [$copList] #-} $dfList d = $copList d |> co@[b] Now suppose we have: dCInt :: C Int blah :: [Int] -> [Int] blah = op ($dfList dCInt) Now we want the built-in op/$dfList rule will fire to give blah = $copList dCInt But with eta-expansion 'blah' might (and in #3772, which is slightly more complicated, does) turn into blah = op (\eta. ($dfList dCInt |> sym co) eta) and now it is *much* harder for the op/$dfList rule to fire, because exprIsConApp_maybe won't hold of the argument to op. I considered trying to *make* it hold, but it's tricky and I gave up. The test simplCore/should_compile/T3722 is an excellent example. -------- End of old out of date comments, just for interest ----------- Note [exprArity for applications] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ When we come to an application we check that the arg is trivial. eg f (fac x) does not have arity 2, even if f has arity 3! * We require that is trivial rather merely cheap. Suppose f has arity 2. Then f (Just y) has arity 0, because if we gave it arity 1 and then inlined f we'd get let v = Just y in \w. <f-body> which has arity 0. And we try to maintain the invariant that we don't have arity decreases. * The `max 0` is important! (\x y -> f x) has arity 2, even if f is unknown, hence arity 0 ************************************************************************ * * Computing the "arity" of an expression * * ************************************************************************ Note [Definition of arity] ~~~~~~~~~~~~~~~~~~~~~~~~~~ The "arity" of an expression 'e' is n if applying 'e' to *fewer* than n *value* arguments converges rapidly Or, to put it another way there is no work lost in duplicating the partial application (e x1 .. x(n-1)) In the divergent case, no work is lost by duplicating because if the thing is evaluated once, that's the end of the program. Or, to put it another way, in any context C C[ (\x1 .. xn. e x1 .. xn) ] is as efficient as C[ e ] It's all a bit more subtle than it looks: Note [One-shot lambdas] ~~~~~~~~~~~~~~~~~~~~~~~ Consider one-shot lambdas let x = expensive in \y z -> E We want this to have arity 1 if the \y-abstraction is a 1-shot lambda. Note [Dealing with bottom] ~~~~~~~~~~~~~~~~~~~~~~~~~~ A Big Deal with computing arities is expressions like f = \x -> case x of True -> \s -> e1 False -> \s -> e2 This happens all the time when f :: Bool -> IO () In this case we do eta-expand, in order to get that \s to the top, and give f arity 2. This isn't really right in the presence of seq. Consider (f bot) `seq` 1 This should diverge! But if we eta-expand, it won't. We ignore this "problem" (unless -fpedantic-bottoms is on), because being scrupulous would lose an important transformation for many programs. (See #5587 for an example.) 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. Technically, 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. So these two transformations aren't always the Right Thing, and we have several tickets reporting unexpected behaviour resulting from this transformation. So we try to limit it as much as possible: (1) Do NOT move a lambda outside a known-bottom case expression case undefined of { (a,b) -> \y -> e } This showed up in #5557 (2) Do NOT move a lambda outside a case if all the branches of the case are known to return bottom. case x of { (a,b) -> \y -> error "urk" } This case is less important, but the idea is that if the fn is going to diverge eventually anyway then getting the best arity isn't an issue, so we might as well play safe (3) Do NOT move a lambda outside a case unless (a) The scrutinee is ok-for-speculation, or (b) more liberally: the scrutinee is cheap (e.g. a variable), and -fpedantic-bottoms is not enforced (see #2915 for an example) Of course both (1) and (2) are readily defeated by disguising the bottoms. 4. Note [Newtype arity] ~~~~~~~~~~~~~~~~~~~~~~~~ Non-recursive newtypes are transparent, and should not get in the way. We do (currently) eta-expand recursive newtypes too. So if we have, say newtype T = MkT ([T] -> Int) Suppose we have e = coerce T f where f has arity 1. Then: etaExpandArity e = 1; that is, etaExpandArity looks through the coerce. When we eta-expand e to arity 1: eta_expand 1 e T we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x) HOWEVER, note that if you use coerce bogusly you can ge coerce Int negate And since negate has arity 2, you might try to eta expand. But you can't decompose Int to a function type. Hence the final case in eta_expand. Note [The state-transformer hack] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Suppose we have f = e where e has arity n. Then, if we know from the context that f has a usage type like t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ... then we can expand the arity to m. This usage type says that any application (x e1 .. en) will be applied to uniquely to (m-n) more args Consider f = \x. let y = <expensive> in case x of True -> foo False -> \(s:RealWorld) -> e where foo has arity 1. Then we want the state hack to apply to foo too, so we can eta expand the case. Then we expect that if f is applied to one arg, it'll be applied to two (that's the hack -- we don't really know, and sometimes it's false) See also Id.isOneShotBndr. Note [State hack and bottoming functions] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ It's a terrible idea to use the state hack on a bottoming function. Here's what happens (#2861): f :: String -> IO T f = \p. error "..." Eta-expand, using the state hack: f = \p. (\s. ((error "...") |> g1) s) |> g2 g1 :: IO T ~ (S -> (S,T)) g2 :: (S -> (S,T)) ~ IO T Extrude the g2 f' = \p. \s. ((error "...") |> g1) s f = f' |> (String -> g2) Discard args for bottomming function f' = \p. \s. ((error "...") |> g1 |> g3 g3 :: (S -> (S,T)) ~ (S,T) Extrude g1.g3 f'' = \p. \s. (error "...") f' = f'' |> (String -> S -> g1.g3) And now we can repeat the whole loop. Aargh! The bug is in applying the state hack to a function which then swallows the argument. This arose in another guise in #3959. Here we had catch# (throw exn >> return ()) Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()]. After inlining (>>) we get catch# (\_. throw {IO ()} exn) We must *not* eta-expand to catch# (\_ _. throw {...} exn) because 'catch#' expects to get a (# _,_ #) after applying its argument to a State#, not another function! In short, we use the state hack to allow us to push let inside a lambda, but not to introduce a new lambda. Note [ArityType] ~~~~~~~~~~~~~~~~ ArityType is the result of a compositional analysis on expressions, from which we can decide the real arity of the expression (extracted with function exprEtaExpandArity). Here is what the fields mean. If an arbitrary expression 'f' has ArityType 'at', then * If at = ABot n, then (f x1..xn) definitely diverges. Partial applications to fewer than n args may *or may not* diverge. We allow ourselves to eta-expand bottoming functions, even if doing so may lose some `seq` sharing, let x = <expensive> in \y. error (g x y) ==> \y. let x = <expensive> in error (g x y) * If at = ATop as, and n=length as, then expanding 'f' to (\x1..xn. f x1 .. xn) loses no sharing, assuming the calls of f respect the one-shot-ness of its definition. NB 'f' is an arbitrary expression, eg (f = g e1 e2). This 'f' can have ArityType as ATop, with length as > 0, only if e1 e2 are themselves. * In both cases, f, (f x1), ... (f x1 ... f(n-1)) are definitely really functions, or bottom, but *not* casts from a data type, in at least one case branch. (If it's a function in one case branch but an unsafe cast from a data type in another, the program is bogus.) So eta expansion is dynamically ok; see Note [State hack and bottoming functions], the part about catch# Example: f = \x\y. let v = <expensive> in \s(one-shot) \t(one-shot). blah 'f' has ArityType [ManyShot,ManyShot,OneShot,OneShot] The one-shot-ness means we can, in effect, push that 'let' inside the \st. Suppose f = \xy. x+y Then f :: AT [False,False] ATop f v :: AT [False] ATop f <expensive> :: AT [] ATop -------------------- Main arity code ---------------------------- -} data ArityType -- See Note [ArityType] = ATop [OneShotInfo] | ABot Arity deriving( ArityType -> ArityType -> Bool (ArityType -> ArityType -> Bool) -> (ArityType -> ArityType -> Bool) -> Eq ArityType forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a /= :: ArityType -> ArityType -> Bool $c/= :: ArityType -> ArityType -> Bool == :: ArityType -> ArityType -> Bool $c== :: ArityType -> ArityType -> Bool Eq ) -- There is always an explicit lambda -- to justify the [OneShot], or the Arity instance Outputable ArityType where ppr :: ArityType -> SDoc ppr (ATop [OneShotInfo] os) = String -> SDoc text String "ATop" SDoc -> SDoc -> SDoc <> SDoc -> SDoc parens (JoinArity -> SDoc forall a. Outputable a => a -> SDoc ppr ([OneShotInfo] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length [OneShotInfo] os)) ppr (ABot JoinArity n) = String -> SDoc text String "ABot" SDoc -> SDoc -> SDoc <> SDoc -> SDoc parens (JoinArity -> SDoc forall a. Outputable a => a -> SDoc ppr JoinArity n) arityTypeArity :: ArityType -> Arity -- The number of value args for the arity type arityTypeArity :: ArityType -> JoinArity arityTypeArity (ATop [OneShotInfo] oss) = [OneShotInfo] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length [OneShotInfo] oss arityTypeArity (ABot JoinArity ar) = JoinArity ar expandableArityType :: ArityType -> Bool -- True <=> eta-expansion will add at least one lambda expandableArityType :: ArityType -> Bool expandableArityType (ATop [OneShotInfo] oss) = Bool -> Bool not ([OneShotInfo] -> Bool forall (t :: * -> *) a. Foldable t => t a -> Bool null [OneShotInfo] oss) expandableArityType (ABot JoinArity ar) = JoinArity ar JoinArity -> JoinArity -> Bool forall a. Eq a => a -> a -> Bool /= JoinArity 0 isBotArityType :: ArityType -> Bool isBotArityType :: ArityType -> Bool isBotArityType (ABot {}) = Bool True isBotArityType (ATop {}) = Bool False arityTypeOneShots :: ArityType -> [OneShotInfo] arityTypeOneShots :: ArityType -> [OneShotInfo] arityTypeOneShots (ATop [OneShotInfo] oss) = [OneShotInfo] oss arityTypeOneShots (ABot JoinArity ar) = JoinArity -> OneShotInfo -> [OneShotInfo] forall a. JoinArity -> a -> [a] replicate JoinArity ar OneShotInfo OneShotLam -- If we are diveging or throwing an exception anyway -- it's fine to push redexes inside the lambdas botArityType :: ArityType botArityType :: ArityType botArityType = JoinArity -> ArityType ABot JoinArity 0 -- Unit for andArityType maxWithArity :: ArityType -> Arity -> ArityType maxWithArity :: ArityType -> JoinArity -> ArityType maxWithArity at :: ArityType at@(ABot {}) JoinArity _ = ArityType at maxWithArity at :: ArityType at@(ATop [OneShotInfo] oss) JoinArity ar | [OneShotInfo] oss [OneShotInfo] -> JoinArity -> Bool forall a. [a] -> JoinArity -> Bool `lengthAtLeast` JoinArity ar = ArityType at | Bool otherwise = [OneShotInfo] -> ArityType ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo] forall a. JoinArity -> [a] -> [a] take JoinArity ar ([OneShotInfo] oss [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo] forall a. [a] -> [a] -> [a] ++ OneShotInfo -> [OneShotInfo] forall a. a -> [a] repeat OneShotInfo NoOneShotInfo)) vanillaArityType :: ArityType vanillaArityType :: ArityType vanillaArityType = [OneShotInfo] -> ArityType ATop [] -- Totally uninformative -- ^ The Arity returned is the number of value args the -- expression can be applied to without doing much work exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType -- exprEtaExpandArity is used when eta expanding -- e ==> \xy -> e x y exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType exprEtaExpandArity DynFlags dflags CoreExpr e = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e where env :: ArityEnv env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv AE { ae_cheap_fn :: CheapFun ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun mk_cheap_fn DynFlags dflags CheapAppFun isCheapApp , ae_ped_bot :: Bool ae_ped_bot = GeneralFlag -> DynFlags -> Bool gopt GeneralFlag Opt_PedanticBottoms DynFlags dflags , ae_joins :: IdSet ae_joins = IdSet emptyVarSet } getBotArity :: ArityType -> Maybe Arity -- Arity of a divergent function getBotArity :: ArityType -> Maybe JoinArity getBotArity (ABot JoinArity n) = JoinArity -> Maybe JoinArity forall a. a -> Maybe a Just JoinArity n getBotArity ArityType _ = Maybe JoinArity forall a. Maybe a Nothing mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun mk_cheap_fn DynFlags dflags CheapAppFun cheap_app | Bool -> Bool not (GeneralFlag -> DynFlags -> Bool gopt GeneralFlag Opt_DictsCheap DynFlags dflags) = \CoreExpr e Maybe Type _ -> CheapAppFun -> CoreExpr -> Bool exprIsCheapX CheapAppFun cheap_app CoreExpr e | Bool otherwise = \CoreExpr e Maybe Type mb_ty -> CheapAppFun -> CoreExpr -> Bool exprIsCheapX CheapAppFun cheap_app CoreExpr e Bool -> Bool -> Bool || case Maybe Type mb_ty of Maybe Type Nothing -> Bool False Just Type ty -> Type -> Bool isDictTy Type ty ---------------------- findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType -- This implements the fixpoint loop for arity analysis -- See Note [Arity analysis] -- If findRhsArity e = (n, is_bot) then -- (a) any application of e to <n arguments will not do much work, -- so it is safe to expand e ==> (\x1..xn. e x1 .. xn) -- (b) if is_bot=True, then e applied to n args is guaranteed bottom findRhsArity :: DynFlags -> Id -> CoreExpr -> JoinArity -> ArityType findRhsArity DynFlags dflags Id bndr CoreExpr rhs JoinArity old_arity = ArityType -> ArityType go (CheapAppFun -> ArityType get_arity CheapAppFun init_cheap_app) -- We always call exprEtaExpandArity once, but usually -- that produces a result equal to old_arity, and then -- we stop right away (since arities should not decrease) -- Result: the common case is that there is just one iteration where init_cheap_app :: CheapAppFun init_cheap_app :: CheapAppFun init_cheap_app Id fn JoinArity n_val_args | Id fn Id -> Id -> Bool forall a. Eq a => a -> a -> Bool == Id bndr = Bool True -- On the first pass, this binder gets infinite arity | Bool otherwise = CheapAppFun isCheapApp Id fn JoinArity n_val_args go :: ArityType -> ArityType go :: ArityType -> ArityType go ArityType cur_atype | JoinArity cur_arity JoinArity -> JoinArity -> Bool forall a. Ord a => a -> a -> Bool <= JoinArity old_arity = ArityType cur_atype | ArityType new_atype ArityType -> ArityType -> Bool forall a. Eq a => a -> a -> Bool == ArityType cur_atype = ArityType cur_atype | Bool otherwise = #if defined(DEBUG) pprTrace "Exciting arity" (vcat [ ppr bndr <+> ppr cur_atype <+> ppr new_atype , ppr rhs]) #endif ArityType -> ArityType go ArityType new_atype where new_atype :: ArityType new_atype = CheapAppFun -> ArityType get_arity CheapAppFun cheap_app cur_arity :: JoinArity cur_arity = ArityType -> JoinArity arityTypeArity ArityType cur_atype cheap_app :: CheapAppFun cheap_app :: CheapAppFun cheap_app Id fn JoinArity n_val_args | Id fn Id -> Id -> Bool forall a. Eq a => a -> a -> Bool == Id bndr = JoinArity n_val_args JoinArity -> JoinArity -> Bool forall a. Ord a => a -> a -> Bool < JoinArity cur_arity | Bool otherwise = CheapAppFun isCheapApp Id fn JoinArity n_val_args get_arity :: CheapAppFun -> ArityType get_arity :: CheapAppFun -> ArityType get_arity CheapAppFun cheap_app = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr rhs where env :: ArityEnv env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv AE { ae_cheap_fn :: CheapFun ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun mk_cheap_fn DynFlags dflags CheapAppFun cheap_app , ae_ped_bot :: Bool ae_ped_bot = GeneralFlag -> DynFlags -> Bool gopt GeneralFlag Opt_PedanticBottoms DynFlags dflags , ae_joins :: IdSet ae_joins = IdSet emptyVarSet } {- Note [Arity analysis] ~~~~~~~~~~~~~~~~~~~~~ The motivating example for arity analysis is this: f = \x. let g = f (x+1) in \y. ...g... What arity does f have? Really it should have arity 2, but a naive look at the RHS won't see that. You need a fixpoint analysis which says it has arity "infinity" the first time round. This example happens a lot; it first showed up in Andy Gill's thesis, fifteen years ago! It also shows up in the code for 'rnf' on lists in #4138. The analysis is easy to achieve because exprEtaExpandArity takes an argument type CheapFun = CoreExpr -> Maybe Type -> Bool used to decide if an expression is cheap enough to push inside a lambda. And exprIsCheapX in turn takes an argument type CheapAppFun = Id -> Int -> Bool which tells when an application is cheap. This makes it easy to write the analysis loop. The analysis is cheap-and-cheerful because it doesn't deal with mutual recursion. But the self-recursive case is the important one. Note [Eta expanding through dictionaries] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ If the experimental -fdicts-cheap flag is on, we eta-expand through dictionary bindings. This improves arities. Thereby, it also means that full laziness is less prone to floating out the application of a function to its dictionary arguments, which can thereby lose opportunities for fusion. Example: foo :: Ord a => a -> ... foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). .... -- So foo has arity 1 f = \x. foo dInt $ bar x The (foo DInt) is floated out, and makes ineffective a RULE foo (bar x) = ... One could go further and make exprIsCheap reply True to any dictionary-typed expression, but that's more work. -} arityLam :: Id -> ArityType -> ArityType arityLam :: Id -> ArityType -> ArityType arityLam Id id (ATop [OneShotInfo] as) = [OneShotInfo] -> ArityType ATop (Id -> OneShotInfo idStateHackOneShotInfo Id id OneShotInfo -> [OneShotInfo] -> [OneShotInfo] forall a. a -> [a] -> [a] : [OneShotInfo] as) arityLam Id _ (ABot JoinArity n) = JoinArity -> ArityType ABot (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a +JoinArity 1) floatIn :: Bool -> ArityType -> ArityType -- We have something like (let x = E in b), -- where b has the given arity type. floatIn :: Bool -> ArityType -> ArityType floatIn Bool _ (ABot JoinArity n) = JoinArity -> ArityType ABot JoinArity n floatIn Bool True (ATop [OneShotInfo] as) = [OneShotInfo] -> ArityType ATop [OneShotInfo] as floatIn Bool False (ATop [OneShotInfo] as) = [OneShotInfo] -> ArityType ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo] forall a. (a -> Bool) -> [a] -> [a] takeWhile OneShotInfo -> Bool isOneShotInfo [OneShotInfo] as) -- If E is not cheap, keep arity only for one-shots arityApp :: ArityType -> Bool -> ArityType -- Processing (fun arg) where at is the ArityType of fun, -- Knock off an argument and behave like 'let' arityApp :: ArityType -> Bool -> ArityType arityApp (ABot JoinArity 0) Bool _ = JoinArity -> ArityType ABot JoinArity 0 arityApp (ABot JoinArity n) Bool _ = JoinArity -> ArityType ABot (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a -JoinArity 1) arityApp (ATop []) Bool _ = [OneShotInfo] -> ArityType ATop [] arityApp (ATop (OneShotInfo _:[OneShotInfo] as)) Bool cheap = Bool -> ArityType -> ArityType floatIn Bool cheap ([OneShotInfo] -> ArityType ATop [OneShotInfo] as) andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case' -- This is least upper bound in the ArityType lattice andArityType :: ArityType -> ArityType -> ArityType andArityType (ABot JoinArity n1) (ABot JoinArity n2) = JoinArity -> ArityType ABot (JoinArity n1 JoinArity -> JoinArity -> JoinArity forall a. Ord a => a -> a -> a `max` JoinArity n2) -- Note [ABot branches: use max] andArityType (ATop [OneShotInfo] as) (ABot JoinArity _) = [OneShotInfo] -> ArityType ATop [OneShotInfo] as andArityType (ABot JoinArity _) (ATop [OneShotInfo] bs) = [OneShotInfo] -> ArityType ATop [OneShotInfo] bs andArityType (ATop [OneShotInfo] as) (ATop [OneShotInfo] bs) = [OneShotInfo] -> ArityType ATop ([OneShotInfo] as [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo] `combine` [OneShotInfo] bs) where -- See Note [Combining case branches] combine :: [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo] combine (OneShotInfo a:[OneShotInfo] as) (OneShotInfo b:[OneShotInfo] bs) = (OneShotInfo a OneShotInfo -> OneShotInfo -> OneShotInfo `bestOneShot` OneShotInfo b) OneShotInfo -> [OneShotInfo] -> [OneShotInfo] forall a. a -> [a] -> [a] : [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo] combine [OneShotInfo] as [OneShotInfo] bs combine [] [OneShotInfo] bs = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo] forall a. (a -> Bool) -> [a] -> [a] takeWhile OneShotInfo -> Bool isOneShotInfo [OneShotInfo] bs combine [OneShotInfo] as [] = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo] forall a. (a -> Bool) -> [a] -> [a] takeWhile OneShotInfo -> Bool isOneShotInfo [OneShotInfo] as {- Note [ABot branches: use max] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Consider case x of True -> \x. error "urk" False -> \xy. error "urk2" Remember: ABot n means "if you apply to n args, it'll definitely diverge". So we need (ABot 2) for the whole thing, the /max/ of the ABot arities. Note [Combining case branches] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Consider go = \x. let z = go e0 go2 = \x. case x of True -> z False -> \s(one-shot). e1 in go2 x We *really* want to eta-expand go and go2. When combining the branches of the case we have ATop [] `andAT` ATop [OneShotLam] and we want to get ATop [OneShotLam]. But if the inner lambda wasn't one-shot we don't want to do this. (We need a proper arity analysis to justify that.) So we combine the best of the two branches, on the (slightly dodgy) basis that if we know one branch is one-shot, then they all must be. Note [Arity trimming] ~~~~~~~~~~~~~~~~~~~~~ Consider ((\x y. blah) |> co), where co :: (Int->Int->Int) ~ (Int -> F a) , and F is some type family. Because of Note [exprArity invariant], item (2), we must return with arity at most 1, because typeArity (Int -> F a) = 1. So we have to trim the result of calling arityType on (\x y. blah). Failing to do so, and hence breaking the exprArity invariant, led to #5441. How to trim? For ATop, it's easy. But we must take great care with ABot. Suppose the expression was (\x y. error "urk"), we'll get (ABot 2). We absolutely must not trim that to (ABot 1), because that claims that ((\x y. error "urk") |> co) diverges when given one argument, which it absolutely does not. And Bad Things happen if we think something returns bottom when it doesn't (#16066). So, do not reduce the 'n' in (ABot n); rather, switch (conservatively) to ATop. Historical note: long ago, we unconditionally switched to ATop when we encountered a cast, but that is far too conservative: see #5475 -} --------------------------- type CheapFun = CoreExpr -> Maybe Type -> Bool -- How to decide if an expression is cheap -- If the Maybe is Just, the type is the type -- of the expression; Nothing means "don't know" data ArityEnv = AE { ArityEnv -> CheapFun ae_cheap_fn :: CheapFun , ArityEnv -> Bool ae_ped_bot :: Bool -- True <=> be pedantic about bottoms , ArityEnv -> IdSet ae_joins :: IdSet -- In-scope join points -- See Note [Eta-expansion and join points] } extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv extendJoinEnv :: ArityEnv -> [Id] -> ArityEnv extendJoinEnv env :: ArityEnv env@(AE { ae_joins :: ArityEnv -> IdSet ae_joins = IdSet joins }) [Id] join_ids = ArityEnv env { ae_joins :: IdSet ae_joins = IdSet joins IdSet -> [Id] -> IdSet `extendVarSetList` [Id] join_ids } ---------------- arityType :: ArityEnv -> CoreExpr -> ArityType arityType :: ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env (Cast CoreExpr e CoercionR co) = case ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e of ATop [OneShotInfo] os -> [OneShotInfo] -> ArityType ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo] forall a. JoinArity -> [a] -> [a] take JoinArity co_arity [OneShotInfo] os) -- See Note [Arity trimming] ABot JoinArity n | JoinArity co_arity JoinArity -> JoinArity -> Bool forall a. Ord a => a -> a -> Bool < JoinArity n -> [OneShotInfo] -> ArityType ATop (JoinArity -> OneShotInfo -> [OneShotInfo] forall a. JoinArity -> a -> [a] replicate JoinArity co_arity OneShotInfo noOneShotInfo) | Bool otherwise -> JoinArity -> ArityType ABot JoinArity n where co_arity :: JoinArity co_arity = [OneShotInfo] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length (Type -> [OneShotInfo] typeArity (CoercionR -> Type coercionRKind CoercionR co)) -- See Note [exprArity invariant] (2); must be true of -- arityType too, since that is how we compute the arity -- of variables, and they in turn affect result of exprArity -- #5441 is a nice demo -- However, do make sure that ATop -> ATop and ABot -> ABot! -- Casts don't affect that part. Getting this wrong provoked #5475 arityType ArityEnv env (Var Id v) | Id v Id -> IdSet -> Bool `elemVarSet` ArityEnv -> IdSet ae_joins ArityEnv env = ArityType botArityType -- See Note [Eta-expansion and join points] | Bool otherwise = Id -> ArityType idArityType Id v -- Lambdas; increase arity arityType ArityEnv env (Lam Id x CoreExpr e) | Id -> Bool isId Id x = Id -> ArityType -> ArityType arityLam Id x (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e) | Bool otherwise = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e -- Applications; decrease arity, except for types arityType ArityEnv env (App CoreExpr fun (Type Type _)) = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr fun arityType ArityEnv env (App CoreExpr fun CoreExpr arg ) = ArityType -> Bool -> ArityType arityApp (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr fun) (ArityEnv -> CheapFun ae_cheap_fn ArityEnv env CoreExpr arg Maybe Type forall a. Maybe a Nothing) -- Case/Let; keep arity if either the expression is cheap -- or it's a 1-shot lambda -- The former is not really right for Haskell -- f x = case x of { (a,b) -> \y. e } -- ===> -- f x y = case x of { (a,b) -> e } -- The difference is observable using 'seq' -- arityType ArityEnv env (Case CoreExpr scrut Id _ Type _ [Alt Id] alts) | CoreExpr -> Bool exprIsDeadEnd CoreExpr scrut Bool -> Bool -> Bool || [Alt Id] -> Bool forall (t :: * -> *) a. Foldable t => t a -> Bool null [Alt Id] alts = ArityType botArityType -- Do not eta expand -- See Note [Dealing with bottom (1)] | Bool otherwise = case ArityType alts_type of ABot JoinArity n | JoinArity nJoinArity -> JoinArity -> Bool forall a. Ord a => a -> a -> Bool >JoinArity 0 -> [OneShotInfo] -> ArityType ATop [] -- Don't eta expand | Bool otherwise -> ArityType botArityType -- if RHS is bottomming -- See Note [Dealing with bottom (2)] ATop [OneShotInfo] as | Bool -> Bool not (ArityEnv -> Bool ae_ped_bot ArityEnv env) -- See Note [Dealing with bottom (3)] , ArityEnv -> CheapFun ae_cheap_fn ArityEnv env CoreExpr scrut Maybe Type forall a. Maybe a Nothing -> [OneShotInfo] -> ArityType ATop [OneShotInfo] as | CoreExpr -> Bool exprOkForSpeculation CoreExpr scrut -> [OneShotInfo] -> ArityType ATop [OneShotInfo] as | Bool otherwise -> [OneShotInfo] -> ArityType ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo] forall a. (a -> Bool) -> [a] -> [a] takeWhile OneShotInfo -> Bool isOneShotInfo [OneShotInfo] as) where alts_type :: ArityType alts_type = (ArityType -> ArityType -> ArityType) -> [ArityType] -> ArityType forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a foldr1 ArityType -> ArityType -> ArityType andArityType [ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr rhs | (AltCon _,[Id] _,CoreExpr rhs) <- [Alt Id] alts] arityType ArityEnv env (Let (NonRec Id j CoreExpr rhs) CoreExpr body) | Just JoinArity join_arity <- Id -> Maybe JoinArity isJoinId_maybe Id j , ([Id] _, CoreExpr rhs_body) <- JoinArity -> CoreExpr -> ([Id], CoreExpr) forall b. JoinArity -> Expr b -> ([b], Expr b) collectNBinders JoinArity join_arity CoreExpr rhs = -- See Note [Eta-expansion and join points] ArityType -> ArityType -> ArityType andArityType (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr rhs_body) (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env' CoreExpr body) where env' :: ArityEnv env' = ArityEnv -> [Id] -> ArityEnv extendJoinEnv ArityEnv env [Id j] arityType ArityEnv env (Let (Rec [(Id, CoreExpr)] pairs) CoreExpr body) | ((Id j,CoreExpr _):[(Id, CoreExpr)] _) <- [(Id, CoreExpr)] pairs , Id -> Bool isJoinId Id j = -- See Note [Eta-expansion and join points] ((Id, CoreExpr) -> ArityType -> ArityType) -> ArityType -> [(Id, CoreExpr)] -> ArityType forall (t :: * -> *) a b. Foldable t => (a -> b -> b) -> b -> t a -> b foldr (ArityType -> ArityType -> ArityType andArityType (ArityType -> ArityType -> ArityType) -> ((Id, CoreExpr) -> ArityType) -> (Id, CoreExpr) -> ArityType -> ArityType forall b c a. (b -> c) -> (a -> b) -> a -> c . (Id, CoreExpr) -> ArityType do_one) (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env' CoreExpr body) [(Id, CoreExpr)] pairs where env' :: ArityEnv env' = ArityEnv -> [Id] -> ArityEnv extendJoinEnv ArityEnv env (((Id, CoreExpr) -> Id) -> [(Id, CoreExpr)] -> [Id] forall a b. (a -> b) -> [a] -> [b] map (Id, CoreExpr) -> Id forall a b. (a, b) -> a fst [(Id, CoreExpr)] pairs) do_one :: (Id, CoreExpr) -> ArityType do_one (Id j,CoreExpr rhs) | Just JoinArity arity <- Id -> Maybe JoinArity isJoinId_maybe Id j = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env' (CoreExpr -> ArityType) -> CoreExpr -> ArityType forall a b. (a -> b) -> a -> b $ ([Id], CoreExpr) -> CoreExpr forall a b. (a, b) -> b snd (([Id], CoreExpr) -> CoreExpr) -> ([Id], CoreExpr) -> CoreExpr forall a b. (a -> b) -> a -> b $ JoinArity -> CoreExpr -> ([Id], CoreExpr) forall b. JoinArity -> Expr b -> ([b], Expr b) collectNBinders JoinArity arity CoreExpr rhs | Bool otherwise = String -> SDoc -> ArityType forall a. HasCallStack => String -> SDoc -> a pprPanic String "arityType:joinrec" ([(Id, CoreExpr)] -> SDoc forall a. Outputable a => a -> SDoc ppr [(Id, CoreExpr)] pairs) arityType ArityEnv env (Let Bind Id b CoreExpr e) = Bool -> ArityType -> ArityType floatIn (Bind Id -> Bool cheap_bind Bind Id b) (ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e) where cheap_bind :: Bind Id -> Bool cheap_bind (NonRec Id b CoreExpr e) = (Id, CoreExpr) -> Bool is_cheap (Id b,CoreExpr e) cheap_bind (Rec [(Id, CoreExpr)] prs) = ((Id, CoreExpr) -> Bool) -> [(Id, CoreExpr)] -> Bool forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool all (Id, CoreExpr) -> Bool is_cheap [(Id, CoreExpr)] prs is_cheap :: (Id, CoreExpr) -> Bool is_cheap (Id b,CoreExpr e) = ArityEnv -> CheapFun ae_cheap_fn ArityEnv env CoreExpr e (Type -> Maybe Type forall a. a -> Maybe a Just (Id -> Type idType Id b)) arityType ArityEnv env (Tick Tickish Id t CoreExpr e) | Bool -> Bool not (Tickish Id -> Bool forall id. Tickish id -> Bool tickishIsCode Tickish Id t) = ArityEnv -> CoreExpr -> ArityType arityType ArityEnv env CoreExpr e arityType ArityEnv _ CoreExpr _ = ArityType vanillaArityType {- Note [Eta-expansion and join points] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Consider this (#18328) f x = join j y = case y of True -> \a. blah False -> \b. blah in case x of A -> j True B -> \c. blah C -> j False and suppose the join point is too big to inline. Now, what is the arity of f? If we inlined the join point, we'd definitely say "arity 2" because we are prepared to push case-scrutinisation inside a lambda. But currently the join point totally messes all that up, because (thought of as a vanilla let-binding) the arity pinned on 'j' is just 1. Why don't we eta-expand j? Because of Note [Do not eta-expand join points] in GHC.Core.Opt.Simplify.Utils Even if we don't eta-expand j, why is its arity only 1? See invariant 2b in Note [Invariants on join points] in GHC.Core. So we do this: * Treat the RHS of a join-point binding, /after/ stripping off join-arity lambda-binders, as very like the body of the let. More precisely, do andArityType with the arityType from the body of the let. * Dually, when we come to a /call/ of a join point, just no-op by returning botArityType, the bottom element of ArityType, which so that: bot `andArityType` x = x * This works if the join point is bound in the expression we are taking the arityType of. But if it's bound further out, it makes no sense to say that (say) the arityType of (j False) is ABot 0. Bad things happen. So we keep track of the in-scope join-point Ids in ae_join. This will make f, above, have arity 2. Then, we'll eta-expand it thus: f x eta = (join j y = ... in case x of ...) eta and the Simplify will automatically push that application of eta into the join points. An alternative (roughly equivalent) idea would be to carry an environment mapping let-bound Ids to their ArityType. -} idArityType :: Id -> ArityType idArityType :: Id -> ArityType idArityType Id v | StrictSig strict_sig <- Id -> StrictSig idStrictness Id v , Bool -> Bool not (Bool -> Bool) -> Bool -> Bool forall a b. (a -> b) -> a -> b $ StrictSig -> Bool isTopSig StrictSig strict_sig , ([Demand] ds, Divergence res) <- StrictSig -> ([Demand], Divergence) splitStrictSig StrictSig strict_sig , let arity :: JoinArity arity = [Demand] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length [Demand] ds = if Divergence -> Bool isDeadEndDiv Divergence res then JoinArity -> ArityType ABot JoinArity arity else [OneShotInfo] -> ArityType ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo] forall a. JoinArity -> [a] -> [a] take JoinArity arity [OneShotInfo] one_shots) | Bool otherwise = [OneShotInfo] -> ArityType ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo] forall a. JoinArity -> [a] -> [a] take (Id -> JoinArity idArity Id v) [OneShotInfo] one_shots) where one_shots :: [OneShotInfo] -- One-shot-ness derived from the type one_shots :: [OneShotInfo] one_shots = Type -> [OneShotInfo] typeArity (Id -> Type idType Id v) {- %************************************************************************ %* * The main eta-expander %* * %************************************************************************ We go for: f = \x1..xn -> N ==> f = \x1..xn y1..ym -> N y1..ym (n >= 0) where (in both cases) * The xi can include type variables * The yi are all value variables * N is a NORMAL FORM (i.e. no redexes anywhere) wanting a suitable number of extra args. The biggest reason for doing this is for cases like f = \x -> case x of True -> \y -> e1 False -> \y -> e2 Here we want to get the lambdas together. A good example is the nofib program fibheaps, which gets 25% more allocation if you don't do this eta-expansion. We may have to sandwich some coerces between the lambdas to make the types work. exprEtaExpandArity looks through coerces when computing arity; and etaExpand adds the coerces as necessary when actually computing the expansion. Note [No crap in eta-expanded code] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The eta expander is careful not to introduce "crap". In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it returns a CoreExpr satisfying the same invariant. See Note [Eta expansion and the CorePrep invariants] in CorePrep. This means the eta-expander has to do a bit of on-the-fly simplification but it's not too hard. The alternative, of relying on a subsequent clean-up phase of the Simplifier to de-crapify the result, means you can't really use it in CorePrep, which is painful. Note [Eta expansion for join points] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The no-crap rule is very tiresome to guarantee when we have join points. Consider eta-expanding let j :: Int -> Int -> Bool j x = e in b The simple way is \(y::Int). (let j x = e in b) y The no-crap way is \(y::Int). let j' :: Int -> Bool j' x = e y in b[j'/j] y where I have written to stress that j's type has changed. Note that (of course!) we have to push the application inside the RHS of the join as well as into the body. AND if j has an unfolding we have to push it into there too. AND j might be recursive... So for now I'm abandoning the no-crap rule in this case. I think that for the use in CorePrep it really doesn't matter; and if it does, then CoreToStg.myCollectArgs will fall over. (Moreover, I think that casts can make the no-crap rule fail too.) Note [Eta expansion and SCCs] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Note that SCCs are not treated specially by etaExpand. If we have etaExpand 2 (\x -> scc "foo" e) = (\xy -> (scc "foo" e) y) So the costs of evaluating 'e' (not 'e y') are attributed to "foo" Note [Eta expansion and source notes] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ CorePrep puts floatable ticks outside of value applications, but not type applications. As a result we might be trying to eta-expand an expression like (src<...> v) @a which we want to lead to code like \x -> src<...> v @a x This means that we need to look through type applications and be ready to re-add floats on the top. Note [Eta expansion with ArityType] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The etaExpandAT function takes an ArityType (not just an Arity) to guide eta-expansion. Why? Because we want to preserve one-shot info. Consider foo = \x. case x of True -> (\s{os}. blah) |> co False -> wubble We'll get an ArityType for foo of (ATop [NoOneShot,OneShot]). Then we want to eta-expand to foo = \x. (\eta{os}. (case x of ...as before...) eta) |> some_co That 'eta' binder is fresh, and we really want it to have the one-shot flag from the inner \s{osf}. By expanding with the ArityType gotten from analysing the RHS, we achieve this neatly. This makes a big difference to the one-shot monad trick; see Note [The one-shot state monad trick] in GHC.Core.Unify. -} -- | @etaExpand n e@ returns an expression with -- the same meaning as @e@, but with arity @n@. -- -- Given: -- -- > e' = etaExpand n e -- -- We should have that: -- -- > ty = exprType e = exprType e' etaExpand :: Arity -> CoreExpr -> CoreExpr etaExpandAT :: ArityType -> CoreExpr -> CoreExpr etaExpand :: JoinArity -> CoreExpr -> CoreExpr etaExpand JoinArity n CoreExpr orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr eta_expand (JoinArity -> OneShotInfo -> [OneShotInfo] forall a. JoinArity -> a -> [a] replicate JoinArity n OneShotInfo NoOneShotInfo) CoreExpr orig_expr etaExpandAT :: ArityType -> CoreExpr -> CoreExpr etaExpandAT ArityType at CoreExpr orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr eta_expand (ArityType -> [OneShotInfo] arityTypeOneShots ArityType at) CoreExpr orig_expr -- See Note [Eta expansion with ArityType] -- etaExpand arity e = res -- Then 'res' has at least 'arity' lambdas at the top -- See Note [Eta expansion with ArityType] -- -- etaExpand deals with for-alls. For example: -- etaExpand 1 E -- where E :: forall a. a -> a -- would return -- (/\b. \y::a -> E b y) -- -- It deals with coerces too, though they are now rare -- so perhaps the extra code isn't worth it eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr eta_expand [OneShotInfo] one_shots CoreExpr orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr go [OneShotInfo] one_shots CoreExpr orig_expr where -- Strip off existing lambdas and casts before handing off to mkEtaWW -- Note [Eta expansion and SCCs] go :: [OneShotInfo] -> CoreExpr -> CoreExpr go [] CoreExpr expr = CoreExpr expr go oss :: [OneShotInfo] oss@(OneShotInfo _:[OneShotInfo] oss1) (Lam Id v CoreExpr body) | Id -> Bool isTyVar Id v = Id -> CoreExpr -> CoreExpr forall b. b -> Expr b -> Expr b Lam Id v ([OneShotInfo] -> CoreExpr -> CoreExpr go [OneShotInfo] oss CoreExpr body) | Bool otherwise = Id -> CoreExpr -> CoreExpr forall b. b -> Expr b -> Expr b Lam Id v ([OneShotInfo] -> CoreExpr -> CoreExpr go [OneShotInfo] oss1 CoreExpr body) go [OneShotInfo] oss (Cast CoreExpr expr CoercionR co) = CoreExpr -> CoercionR -> CoreExpr forall b. Expr b -> CoercionR -> Expr b Cast ([OneShotInfo] -> CoreExpr -> CoreExpr go [OneShotInfo] oss CoreExpr expr) CoercionR co go [OneShotInfo] oss CoreExpr expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $ CoreExpr -> CoreExpr retick (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr forall a b. (a -> b) -> a -> b $ [EtaInfo] -> CoreExpr -> CoreExpr etaInfoAbs [EtaInfo] etas (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst' CoreExpr sexpr [EtaInfo] etas) where in_scope :: InScopeSet in_scope = IdSet -> InScopeSet mkInScopeSet (CoreExpr -> IdSet exprFreeVars CoreExpr expr) (InScopeSet in_scope', [EtaInfo] etas) = [OneShotInfo] -> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo]) mkEtaWW [OneShotInfo] oss (CoreExpr -> SDoc forall a. Outputable a => a -> SDoc ppr CoreExpr orig_expr) InScopeSet in_scope (CoreExpr -> Type exprType CoreExpr expr) subst' :: Subst subst' = InScopeSet -> Subst mkEmptySubst InScopeSet in_scope' -- Find ticks behind type apps. -- See Note [Eta expansion and source notes] (CoreExpr expr', [CoreExpr] args) = CoreExpr -> (CoreExpr, [CoreExpr]) forall b. Expr b -> (Expr b, [Expr b]) collectArgs CoreExpr expr ([Tickish Id] ticks, CoreExpr expr'') = (Tickish Id -> Bool) -> CoreExpr -> ([Tickish Id], CoreExpr) forall b. (Tickish Id -> Bool) -> Expr b -> ([Tickish Id], Expr b) stripTicksTop Tickish Id -> Bool forall id. Tickish id -> Bool tickishFloatable CoreExpr expr' sexpr :: CoreExpr sexpr = (CoreExpr -> CoreExpr -> CoreExpr) -> CoreExpr -> [CoreExpr] -> CoreExpr forall (t :: * -> *) b a. Foldable t => (b -> a -> b) -> b -> t a -> b foldl' CoreExpr -> CoreExpr -> CoreExpr forall b. Expr b -> Expr b -> Expr b App CoreExpr expr'' [CoreExpr] args retick :: CoreExpr -> CoreExpr retick CoreExpr expr = (Tickish Id -> CoreExpr -> CoreExpr) -> CoreExpr -> [Tickish Id] -> CoreExpr forall (t :: * -> *) a b. Foldable t => (a -> b -> b) -> b -> t a -> b foldr Tickish Id -> CoreExpr -> CoreExpr mkTick CoreExpr expr [Tickish Id] ticks -- Abstraction Application -------------- data EtaInfo = EtaVar Var -- /\a. [] [] a -- \x. [] [] x | EtaCo Coercion -- [] |> sym co [] |> co instance Outputable EtaInfo where ppr :: EtaInfo -> SDoc ppr (EtaVar Id v) = String -> SDoc text String "EtaVar" SDoc -> SDoc -> SDoc <+> Id -> SDoc forall a. Outputable a => a -> SDoc ppr Id v ppr (EtaCo CoercionR co) = String -> SDoc text String "EtaCo" SDoc -> SDoc -> SDoc <+> CoercionR -> SDoc forall a. Outputable a => a -> SDoc ppr CoercionR co pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo] pushCoercion :: CoercionR -> [EtaInfo] -> [EtaInfo] pushCoercion CoercionR co1 (EtaCo CoercionR co2 : [EtaInfo] eis) | CoercionR -> Bool isReflCo CoercionR co = [EtaInfo] eis | Bool otherwise = CoercionR -> EtaInfo EtaCo CoercionR co EtaInfo -> [EtaInfo] -> [EtaInfo] forall a. a -> [a] -> [a] : [EtaInfo] eis where co :: CoercionR co = CoercionR co1 CoercionR -> CoercionR -> CoercionR `mkTransCo` CoercionR co2 pushCoercion CoercionR co [EtaInfo] eis = CoercionR -> EtaInfo EtaCo CoercionR co EtaInfo -> [EtaInfo] -> [EtaInfo] forall a. a -> [a] -> [a] : [EtaInfo] eis -------------- etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr etaInfoAbs [] CoreExpr expr = CoreExpr expr etaInfoAbs (EtaVar Id v : [EtaInfo] eis) CoreExpr expr = Id -> CoreExpr -> CoreExpr forall b. b -> Expr b -> Expr b Lam Id v ([EtaInfo] -> CoreExpr -> CoreExpr etaInfoAbs [EtaInfo] eis CoreExpr expr) etaInfoAbs (EtaCo CoercionR co : [EtaInfo] eis) CoreExpr expr = CoreExpr -> CoercionR -> CoreExpr forall b. Expr b -> CoercionR -> Expr b Cast ([EtaInfo] -> CoreExpr -> CoreExpr etaInfoAbs [EtaInfo] eis CoreExpr expr) (CoercionR -> CoercionR mkSymCo CoercionR co) -------------- etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr -- (etaInfoApp s e eis) returns something equivalent to -- ((substExpr s e) `appliedto` eis) etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst (Lam Id v1 CoreExpr e) (EtaVar Id v2 : [EtaInfo] eis) = Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp (Subst -> Id -> Id -> Subst GHC.Core.Subst.extendSubstWithVar Subst subst Id v1 Id v2) CoreExpr e [EtaInfo] eis etaInfoApp Subst subst (Cast CoreExpr e CoercionR co1) [EtaInfo] eis = Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst CoreExpr e (CoercionR -> [EtaInfo] -> [EtaInfo] pushCoercion CoercionR co' [EtaInfo] eis) where co' :: CoercionR co' = HasCallStack => Subst -> CoercionR -> CoercionR Subst -> CoercionR -> CoercionR GHC.Core.Subst.substCo Subst subst CoercionR co1 etaInfoApp Subst subst (Case CoreExpr e Id b Type ty [Alt Id] alts) [EtaInfo] eis = CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr forall b. Expr b -> b -> Type -> [Alt b] -> Expr b Case (Subst -> CoreExpr -> CoreExpr subst_expr Subst subst CoreExpr e) Id b1 Type ty' [Alt Id] alts' where (Subst subst1, Id b1) = Subst -> Id -> (Subst, Id) substBndr Subst subst Id b alts' :: [Alt Id] alts' = (Alt Id -> Alt Id) -> [Alt Id] -> [Alt Id] forall a b. (a -> b) -> [a] -> [b] map Alt Id -> Alt Id forall {a}. (a, [Id], CoreExpr) -> (a, [Id], CoreExpr) subst_alt [Alt Id] alts ty' :: Type ty' = Type -> [EtaInfo] -> Type etaInfoAppTy (Subst -> Type -> Type GHC.Core.Subst.substTy Subst subst Type ty) [EtaInfo] eis subst_alt :: (a, [Id], CoreExpr) -> (a, [Id], CoreExpr) subst_alt (a con, [Id] bs, CoreExpr rhs) = (a con, [Id] bs', Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst2 CoreExpr rhs [EtaInfo] eis) where (Subst subst2,[Id] bs') = Subst -> [Id] -> (Subst, [Id]) substBndrs Subst subst1 [Id] bs etaInfoApp Subst subst (Let Bind Id b CoreExpr e) [EtaInfo] eis | Bool -> Bool not (Bind Id -> Bool isJoinBind Bind Id b) -- See Note [Eta expansion for join points] = Bind Id -> CoreExpr -> CoreExpr forall b. Bind b -> Expr b -> Expr b Let Bind Id b' (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst' CoreExpr e [EtaInfo] eis) where (Subst subst', Bind Id b') = HasDebugCallStack => Subst -> Bind Id -> (Subst, Bind Id) Subst -> Bind Id -> (Subst, Bind Id) substBindSC Subst subst Bind Id b etaInfoApp Subst subst (Tick Tickish Id t CoreExpr e) [EtaInfo] eis = Tickish Id -> CoreExpr -> CoreExpr forall b. Tickish Id -> Expr b -> Expr b Tick (Subst -> Tickish Id -> Tickish Id substTickish Subst subst Tickish Id t) (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr etaInfoApp Subst subst CoreExpr e [EtaInfo] eis) etaInfoApp Subst subst CoreExpr expr [EtaInfo] _ | (Var Id fun, [CoreExpr] _) <- CoreExpr -> (CoreExpr, [CoreExpr]) forall b. Expr b -> (Expr b, [Expr b]) collectArgs CoreExpr expr , Var Id fun' <- HasDebugCallStack => Subst -> Id -> CoreExpr Subst -> Id -> CoreExpr lookupIdSubst Subst subst Id fun , Id -> Bool isJoinId Id fun' = Subst -> CoreExpr -> CoreExpr subst_expr Subst subst CoreExpr expr etaInfoApp Subst subst CoreExpr e [EtaInfo] eis = CoreExpr -> [EtaInfo] -> CoreExpr forall {b}. Expr b -> [EtaInfo] -> Expr b go (Subst -> CoreExpr -> CoreExpr subst_expr Subst subst CoreExpr e) [EtaInfo] eis where go :: Expr b -> [EtaInfo] -> Expr b go Expr b e [] = Expr b e go Expr b e (EtaVar Id v : [EtaInfo] eis) = Expr b -> [EtaInfo] -> Expr b go (Expr b -> Expr b -> Expr b forall b. Expr b -> Expr b -> Expr b App Expr b e (Id -> Expr b forall b. Id -> Expr b varToCoreExpr Id v)) [EtaInfo] eis go Expr b e (EtaCo CoercionR co : [EtaInfo] eis) = Expr b -> [EtaInfo] -> Expr b go (Expr b -> CoercionR -> Expr b forall b. Expr b -> CoercionR -> Expr b Cast Expr b e CoercionR co) [EtaInfo] eis -------------- etaInfoAppTy :: Type -> [EtaInfo] -> Type -- If e :: ty -- then etaInfoApp e eis :: etaInfoApp ty eis etaInfoAppTy :: Type -> [EtaInfo] -> Type etaInfoAppTy Type ty [] = Type ty etaInfoAppTy Type ty (EtaVar Id v : [EtaInfo] eis) = Type -> [EtaInfo] -> Type etaInfoAppTy (Type -> CoreExpr -> Type applyTypeToArg Type ty (Id -> CoreExpr forall b. Id -> Expr b varToCoreExpr Id v)) [EtaInfo] eis etaInfoAppTy Type _ (EtaCo CoercionR co : [EtaInfo] eis) = Type -> [EtaInfo] -> Type etaInfoAppTy (CoercionR -> Type coercionRKind CoercionR co) [EtaInfo] eis -------------- -- | @mkEtaWW n _ fvs ty@ will compute the 'EtaInfo' necessary for eta-expanding -- an expression @e :: ty@ to take @n@ value arguments, where @fvs@ are the -- free variables of @e@. -- -- Note that this function is entirely unconcerned about cost centres and other -- semantically-irrelevant source annotations, so call sites must take care to -- preserve that info. See Note [Eta expansion and SCCs]. mkEtaWW :: [OneShotInfo] -- ^ How many value arguments to eta-expand -> SDoc -- ^ The pretty-printed original expression, for warnings. -> InScopeSet -- ^ A super-set of the free vars of the expression to eta-expand. -> Type -> (InScopeSet, [EtaInfo]) -- ^ The variables in 'EtaInfo' are fresh wrt. to the incoming 'InScopeSet'. -- The outgoing 'InScopeSet' extends the incoming 'InScopeSet' with the -- fresh variables in 'EtaInfo'. mkEtaWW :: [OneShotInfo] -> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo]) mkEtaWW [OneShotInfo] orig_oss SDoc ppr_orig_expr InScopeSet in_scope Type orig_ty = JoinArity -> [OneShotInfo] -> TCvSubst -> Type -> [EtaInfo] -> (InScopeSet, [EtaInfo]) go JoinArity 0 [OneShotInfo] orig_oss TCvSubst empty_subst Type orig_ty [] where empty_subst :: TCvSubst empty_subst = InScopeSet -> TCvSubst mkEmptyTCvSubst InScopeSet in_scope go :: Int -- For fresh names -> [OneShotInfo] -- Number of value args to expand to -> TCvSubst -> Type -- We are really looking at subst(ty) -> [EtaInfo] -- Accumulating parameter -> (InScopeSet, [EtaInfo]) go :: JoinArity -> [OneShotInfo] -> TCvSubst -> Type -> [EtaInfo] -> (InScopeSet, [EtaInfo]) go JoinArity _ [] TCvSubst subst Type _ [EtaInfo] eis -- See Note [exprArity invariant] ----------- Done! No more expansion needed = (TCvSubst -> InScopeSet getTCvInScope TCvSubst subst, [EtaInfo] -> [EtaInfo] forall a. [a] -> [a] reverse [EtaInfo] eis) go JoinArity n oss :: [OneShotInfo] oss@(OneShotInfo one_shot:[OneShotInfo] oss1) TCvSubst subst Type ty [EtaInfo] eis -- See Note [exprArity invariant] ----------- Forall types (forall a. ty) | Just (Id tcv,Type ty') <- Type -> Maybe (Id, Type) splitForAllTy_maybe Type ty , (TCvSubst subst', Id tcv') <- HasCallStack => TCvSubst -> Id -> (TCvSubst, Id) TCvSubst -> Id -> (TCvSubst, Id) Type.substVarBndr TCvSubst subst Id tcv , let oss' :: [OneShotInfo] oss' | Id -> Bool isTyVar Id tcv = [OneShotInfo] oss | Bool otherwise = [OneShotInfo] oss1 -- A forall can bind a CoVar, in which case -- we consume one of the [OneShotInfo] = JoinArity -> [OneShotInfo] -> TCvSubst -> Type -> [EtaInfo] -> (InScopeSet, [EtaInfo]) go JoinArity n [OneShotInfo] oss' TCvSubst subst' Type ty' (Id -> EtaInfo EtaVar Id tcv' EtaInfo -> [EtaInfo] -> [EtaInfo] forall a. a -> [a] -> [a] : [EtaInfo] eis) ----------- Function types (t1 -> t2) | Just (Type mult, Type arg_ty, Type res_ty) <- Type -> Maybe (Type, Type, Type) splitFunTy_maybe Type ty , Bool -> Bool not (Type -> Bool isTypeLevPoly Type arg_ty) -- See Note [Levity polymorphism invariants] in GHC.Core -- See also test case typecheck/should_run/EtaExpandLevPoly , (TCvSubst subst', Id eta_id) <- JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id) freshEtaId JoinArity n TCvSubst subst (Type -> Type -> Scaled Type forall a. Type -> a -> Scaled a Scaled Type mult Type arg_ty) -- Avoid free vars of the original expression , let eta_id' :: Id eta_id' = Id eta_id Id -> OneShotInfo -> Id `setIdOneShotInfo` OneShotInfo one_shot = JoinArity -> [OneShotInfo] -> TCvSubst -> Type -> [EtaInfo] -> (InScopeSet, [EtaInfo]) go (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a +JoinArity 1) [OneShotInfo] oss1 TCvSubst subst' Type res_ty (Id -> EtaInfo EtaVar Id eta_id' EtaInfo -> [EtaInfo] -> [EtaInfo] forall a. a -> [a] -> [a] : [EtaInfo] eis) ----------- Newtypes -- Given this: -- newtype T = MkT ([T] -> Int) -- Consider eta-expanding this -- eta_expand 1 e T -- We want to get -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x) | Just (CoercionR co, Type ty') <- Type -> Maybe (CoercionR, Type) topNormaliseNewType_maybe Type ty , let co' :: CoercionR co' = HasCallStack => TCvSubst -> CoercionR -> CoercionR TCvSubst -> CoercionR -> CoercionR Coercion.substCo TCvSubst subst CoercionR co -- Remember to apply the substitution to co (#16979) -- (or we could have applied to ty, but then -- we'd have had to zap it for the recursive call) = JoinArity -> [OneShotInfo] -> TCvSubst -> Type -> [EtaInfo] -> (InScopeSet, [EtaInfo]) go JoinArity n [OneShotInfo] oss TCvSubst subst Type ty' (CoercionR -> [EtaInfo] -> [EtaInfo] pushCoercion CoercionR co' [EtaInfo] eis) | Bool otherwise -- We have an expression of arity > 0, -- but its type isn't a function, or a binder -- is levity-polymorphic = WARN( True, (ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr ) (TCvSubst -> InScopeSet getTCvInScope TCvSubst subst, [EtaInfo] -> [EtaInfo] forall a. [a] -> [a] reverse [EtaInfo] eis) -- This *can* legitimately happen: -- e.g. coerce Int (\x. x) Essentially the programmer is -- playing fast and loose with types (Happy does this a lot). -- So we simply decline to eta-expand. Otherwise we'd end up -- with an explicit lambda having a non-function type ------------ subst_expr :: Subst -> CoreExpr -> CoreExpr -- Apply a substitution to an expression. We use substExpr -- not substExprSC (short-cutting substitution) because -- we may be changing the types of join points, so applying -- the in-scope set is necessary. -- -- ToDo: we could instead check if we actually *are* -- changing any join points' types, and if not use substExprSC. subst_expr :: Subst -> CoreExpr -> CoreExpr subst_expr = HasDebugCallStack => Subst -> CoreExpr -> CoreExpr Subst -> CoreExpr -> CoreExpr substExpr -------------- -- | Split an expression into the given number of binders and a body, -- eta-expanding if necessary. Counts value *and* type binders. etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr) etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([Id], CoreExpr) etaExpandToJoinPoint JoinArity join_arity CoreExpr expr = JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr) go JoinArity join_arity [] CoreExpr expr where go :: JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr) go JoinArity 0 [Id] rev_bs CoreExpr e = ([Id] -> [Id] forall a. [a] -> [a] reverse [Id] rev_bs, CoreExpr e) go JoinArity n [Id] rev_bs (Lam Id b CoreExpr e) = JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr) go (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a -JoinArity 1) (Id b Id -> [Id] -> [Id] forall a. a -> [a] -> [a] : [Id] rev_bs) CoreExpr e go JoinArity n [Id] rev_bs CoreExpr e = case JoinArity -> CoreExpr -> ([Id], CoreExpr) etaBodyForJoinPoint JoinArity n CoreExpr e of ([Id] bs, CoreExpr e') -> ([Id] -> [Id] forall a. [a] -> [a] reverse [Id] rev_bs [Id] -> [Id] -> [Id] forall a. [a] -> [a] -> [a] ++ [Id] bs, CoreExpr e') etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule etaExpandToJoinPointRule JoinArity _ rule :: CoreRule rule@(BuiltinRule {}) = WARN(True, (sep [text "Can't eta-expand built-in rule:", ppr rule])) -- How did a local binding get a built-in rule anyway? Probably a plugin. CoreRule rule etaExpandToJoinPointRule JoinArity join_arity rule :: CoreRule rule@(Rule { ru_bndrs :: CoreRule -> [Id] ru_bndrs = [Id] bndrs, ru_rhs :: CoreRule -> CoreExpr ru_rhs = CoreExpr rhs , ru_args :: CoreRule -> [CoreExpr] ru_args = [CoreExpr] args }) | JoinArity need_args JoinArity -> JoinArity -> Bool forall a. Eq a => a -> a -> Bool == JoinArity 0 = CoreRule rule | JoinArity need_args JoinArity -> JoinArity -> Bool forall a. Ord a => a -> a -> Bool < JoinArity 0 = String -> SDoc -> CoreRule forall a. HasCallStack => String -> SDoc -> a pprPanic String "etaExpandToJoinPointRule" (JoinArity -> SDoc forall a. Outputable a => a -> SDoc ppr JoinArity join_arity SDoc -> SDoc -> SDoc $$ CoreRule -> SDoc forall a. Outputable a => a -> SDoc ppr CoreRule rule) | Bool otherwise = CoreRule rule { ru_bndrs :: [Id] ru_bndrs = [Id] bndrs [Id] -> [Id] -> [Id] forall a. [a] -> [a] -> [a] ++ [Id] new_bndrs, ru_args :: [CoreExpr] ru_args = [CoreExpr] args [CoreExpr] -> [CoreExpr] -> [CoreExpr] forall a. [a] -> [a] -> [a] ++ [CoreExpr] forall {b}. [Expr b] new_args , ru_rhs :: CoreExpr ru_rhs = CoreExpr new_rhs } where need_args :: JoinArity need_args = JoinArity join_arity JoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a - [CoreExpr] -> JoinArity forall (t :: * -> *) a. Foldable t => t a -> JoinArity length [CoreExpr] args ([Id] new_bndrs, CoreExpr new_rhs) = JoinArity -> CoreExpr -> ([Id], CoreExpr) etaBodyForJoinPoint JoinArity need_args CoreExpr rhs new_args :: [Expr b] new_args = [Id] -> [Expr b] forall b. [Id] -> [Expr b] varsToCoreExprs [Id] new_bndrs -- Adds as many binders as asked for; assumes expr is not a lambda etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr) etaBodyForJoinPoint :: JoinArity -> CoreExpr -> ([Id], CoreExpr) etaBodyForJoinPoint JoinArity need_args CoreExpr body = JoinArity -> Type -> TCvSubst -> [Id] -> CoreExpr -> ([Id], CoreExpr) forall {b}. JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b) go JoinArity need_args (CoreExpr -> Type exprType CoreExpr body) (CoreExpr -> TCvSubst init_subst CoreExpr body) [] CoreExpr body where go :: JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b) go JoinArity 0 Type _ TCvSubst _ [Id] rev_bs Expr b e = ([Id] -> [Id] forall a. [a] -> [a] reverse [Id] rev_bs, Expr b e) go JoinArity n Type ty TCvSubst subst [Id] rev_bs Expr b e | Just (Id tv, Type res_ty) <- Type -> Maybe (Id, Type) splitForAllTy_maybe Type ty , let (TCvSubst subst', Id tv') = HasCallStack => TCvSubst -> Id -> (TCvSubst, Id) TCvSubst -> Id -> (TCvSubst, Id) Type.substVarBndr TCvSubst subst Id tv = JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b) go (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a -JoinArity 1) Type res_ty TCvSubst subst' (Id tv' Id -> [Id] -> [Id] forall a. a -> [a] -> [a] : [Id] rev_bs) (Expr b e Expr b -> Expr b -> Expr b forall b. Expr b -> Expr b -> Expr b `App` Id -> Expr b forall b. Id -> Expr b varToCoreExpr Id tv') | Just (Type mult, Type arg_ty, Type res_ty) <- Type -> Maybe (Type, Type, Type) splitFunTy_maybe Type ty , let (TCvSubst subst', Id b) = JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id) freshEtaId JoinArity n TCvSubst subst (Type -> Type -> Scaled Type forall a. Type -> a -> Scaled a Scaled Type mult Type arg_ty) = JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b) go (JoinArity nJoinArity -> JoinArity -> JoinArity forall a. Num a => a -> a -> a -JoinArity 1) Type res_ty TCvSubst subst' (Id b Id -> [Id] -> [Id] forall a. a -> [a] -> [a] : [Id] rev_bs) (Expr b e Expr b -> Expr b -> Expr b forall b. Expr b -> Expr b -> Expr b `App` Id -> Expr b forall b. Id -> Expr b Var Id b) | Bool otherwise = String -> SDoc -> ([Id], Expr b) forall a. HasCallStack => String -> SDoc -> a pprPanic String "etaBodyForJoinPoint" (SDoc -> ([Id], Expr b)) -> SDoc -> ([Id], Expr b) forall a b. (a -> b) -> a -> b $ JoinArity -> SDoc int JoinArity need_args SDoc -> SDoc -> SDoc $$ CoreExpr -> SDoc forall a. Outputable a => a -> SDoc ppr CoreExpr body SDoc -> SDoc -> SDoc $$ Type -> SDoc forall a. Outputable a => a -> SDoc ppr (CoreExpr -> Type exprType CoreExpr body) init_subst :: CoreExpr -> TCvSubst init_subst CoreExpr e = InScopeSet -> TCvSubst mkEmptyTCvSubst (IdSet -> InScopeSet mkInScopeSet (CoreExpr -> IdSet exprFreeVars CoreExpr e)) -------------- freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id) -- Make a fresh Id, with specified type (after applying substitution) -- It should be "fresh" in the sense that it's not in the in-scope set -- of the TvSubstEnv; and it should itself then be added to the in-scope -- set of the TvSubstEnv -- -- The Int is just a reasonable starting point for generating a unique; -- it does not necessarily have to be unique itself. freshEtaId :: JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id) freshEtaId JoinArity n TCvSubst subst Scaled Type ty = (TCvSubst subst', Id eta_id') where Scaled Type mult' Type ty' = HasCallStack => TCvSubst -> Scaled Type -> Scaled Type TCvSubst -> Scaled Type -> Scaled Type Type.substScaledTyUnchecked TCvSubst subst Scaled Type ty eta_id' :: Id eta_id' = InScopeSet -> Id -> Id uniqAway (TCvSubst -> InScopeSet getTCvInScope TCvSubst subst) (Id -> Id) -> Id -> Id forall a b. (a -> b) -> a -> b $ FastString -> Unique -> Type -> Type -> Id mkSysLocalOrCoVar (String -> FastString fsLit String "eta") (JoinArity -> Unique mkBuiltinUnique JoinArity n) Type mult' Type ty' -- "OrCoVar" since this can be used to eta-expand -- coercion abstractions subst' :: TCvSubst subst' = TCvSubst -> Id -> TCvSubst extendTCvInScope TCvSubst subst Id eta_id'