1 {-

2 (c) The University of Glasgow 2006

3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998

6 Arity and eta expansion

7 -}

9 {-# LANGUAGE CPP #-}

11 -- | Arity and eta expansion

19 import CoreSyn

20 import CoreFVs

21 import CoreUtils

22 import CoreSubst

23 import Demand

24 import Var

25 import VarEnv

26 import Id

27 import Type

29 import Coercion

30 import BasicTypes

31 import Unique

33 import Outputable

34 import FastString

35 import Pair

38 {-

39 ************************************************************************

40 * *

41 manifestArity and exprArity

42 * *

43 ************************************************************************

45 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.

46 It tells how many things the expression can be applied to before doing

47 any work. It doesn't look inside cases, lets, etc. The idea is that

48 exprEtaExpandArity will do the hard work, leaving something that's easy

49 for exprArity to grapple with. In particular, Simplify uses exprArity to

50 compute the ArityInfo for the Id.

52 Originally I thought that it was enough just to look for top-level lambdas, but

53 it isn't. I've seen this

55 foo = PrelBase.timesInt

57 We want foo to get arity 2 even though the eta-expander will leave it

58 unchanged, in the expectation that it'll be inlined. But occasionally it

59 isn't, because foo is blacklisted (used in a rule).

61 Similarly, see the ok_note check in exprEtaExpandArity. So

62 f = __inline_me (\x -> e)

63 won't be eta-expanded.

65 And in any case it seems more robust to have exprArity be a bit more intelligent.

66 But note that (\x y z -> f x y z)

67 should have arity 3, regardless of f's arity.

68 -}

71 -- ^ manifestArity sees how many leading value lambdas there are,

72 -- after looking through casts

79 ---------------

81 -- ^ An approximate, fast, version of 'exprEtaExpandArity'

82 exprArity e = go e

83 where

89 -- Note [exprArity invariant]

92 -- See Note [exprArity for applications]

93 -- NB: coercions count as a value argument

100 ---------------

102 -- How many value arrows are visible in the type?

103 -- We look through foralls, and newtypes

104 -- See Note [exprArity invariant]

105 typeArity ty

106 = go initRecTc ty

107 where

108 go rec_nts ty

118 -- in TyCon

119 -- , not (isClassTyCon tc) -- Do not eta-expand through newtype classes

120 -- -- See Note [Newtype classes and eta expansion]

121 -- (no longer required)

123 -- Important to look through non-recursive newtypes, so that, eg

124 -- (f x) where f has arity 2, f :: Int -> IO ()

125 -- Here we want to get arity 1 for the result!

126 --

127 -- AND through a layer of recursive newtypes

128 -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))

130 | otherwise

131 = []

133 ---------------

135 -- A cheap and cheerful function that identifies bottoming functions

136 -- and gives them a suitable strictness signatures. It's used during

137 -- float-out

138 exprBotStrictness_maybe e

140 Nothing -> Nothing

142 where

145 -- For this purpose we can be very simple

146 -- exnRes is a bit less aggressive than botRes

148 {-

149 Note [exprArity invariant]

150 ~~~~~~~~~~~~~~~~~~~~~~~~~~

151 exprArity has the following invariant:

153 (1) If typeArity (exprType e) = n,

154 then manifestArity (etaExpand e n) = n

156 That is, etaExpand can always expand as much as typeArity says

157 So the case analysis in etaExpand and in typeArity must match

159 (2) exprArity e <= typeArity (exprType e)

161 (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n

163 That is, if exprArity says "the arity is n" then etaExpand really

164 can get "n" manifest lambdas to the top.

166 Why is this important? Because

167 - In TidyPgm we use exprArity to fix the *final arity* of

168 each top-level Id, and in

169 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas

170 actually match that arity, which in turn means

171 that the StgRhs has the right number of lambdas

173 An alternative would be to do the eta-expansion in TidyPgm, at least

174 for top-level bindings, in which case we would not need the trim_arity

175 in exprArity. That is a less local change, so I'm going to leave it for today!

177 Note [Newtype classes and eta expansion]

178 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

179 NB: this nasty special case is no longer required, because

180 for newtype classes we don't use the class-op rule mechanism

181 at all. See Note [Single-method classes] in TcInstDcls. SLPJ May 2013

183 -------- Old out of date comments, just for interest -----------

184 We have to be careful when eta-expanding through newtypes. In general

185 it's a good idea, but annoyingly it interacts badly with the class-op

186 rule mechanism. Consider

188 class C a where { op :: a -> a }

189 instance C b => C [b] where

190 op x = ...

192 These translate to

194 co :: forall a. (a->a) ~ C a

196 $copList :: C b -> [b] -> [b]

197 $copList d x = ...

199 $dfList :: C b -> C [b]

200 {-# DFunUnfolding = [$copList] #-}

203 Now suppose we have:

211 blah = $copList dCInt

223 -------- End of old out of date comments, just for interest -----------

227 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

228 When we come to an application we check that the arg is trivial.

237 have arity decreases.

243 ************************************************************************

244 * *

246 * *

247 ************************************************************************

250 ~~~~~~~~~~~~~~~~~~~~~~~~~~

253 converges rapidly

255 Or, to put it another way

257 there is no work lost in duplicating the partial

272 ~~~~~~~~~~~~~~~~~~~~~~~

273 Consider one-shot lambdas

278 ~~~~~~~~~~~~~~~~~~~~~~~~~~

279 A Big Deal with computing arities is expressions like

294 would lose an important transformation for many programs. (See

297 Consider also

310 have several tickets reporting unexpected bahaviour resulting from

321 going to diverge eventually anyway then getting the best arity

332 ~~~~~~~~~~~~~~~~~~~~~~~~

338 Suppose we have

339 e = coerce T f

352 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

353 Suppose we have

354 f = e

356 a usage type like

372 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

385 Extrude the g2

390 Discard args for bottomming function

395 Extrude g1.g3

408 After inlining (>>) we get

424 ~~~~~~~~~~~~~~~~

443 its definition.

447 themselves.

456 Example:

469 -------------------- Main arity code ----------------------------

470 -}

472 -- See Note [ArityType]

474 -- There is always an explicit lambda

475 -- to justify the [OneShot], or the Arity

477 vanillaArityType :: ArityType

480 -- ^ The Arity returned is the number of value args the

481 -- expression can be applied to without doing much work

483 -- exprEtaExpandArity is used when eta expanding

484 -- e ==> \xy -> e x y

485 exprEtaExpandArity dflags e

488 ABot n -> n

489 where

494 -- Arity of a divergent function

496 getBotArity _ = Nothing

499 mk_cheap_fn dflags cheap_app

502 | otherwise

506 Just ty -> isDictLikeTy ty

509 ----------------------

511 -- This implements the fixpoint loop for arity analysis

512 -- See Note [Arity analysis]

513 findRhsArity dflags bndr rhs old_arity

515 -- We always call exprEtaExpandArity once, but usually

516 -- that produces a result equal to old_arity, and then

517 -- we stop right away (since arities should not decrease)

518 -- Result: the common case is that there is just one iteration

519 where

520 init_cheap_app :: CheapAppFun

521 init_cheap_app fn n_val_args

526 go cur_arity

530 #ifdef DEBUG

531 pprTrace "Exciting arity"

534 #endif

535 go new_arity

536 where

537 new_arity = rhsEtaExpandArity dflags cheap_app rhs

539 cheap_app :: CheapAppFun

540 cheap_app fn n_val_args

544 -- ^ The Arity returned is the number of value args the

545 -- expression can be applied to without doing much work

547 -- exprEtaExpandArity is used when eta expanding

548 -- e ==> \xy -> e x y

549 rhsEtaExpandArity dflags cheap_app e

553 -- Don't expand PAPs/thunks

554 -- Note [Eta expanding thunks]

557 ABot n -> n

558 where

566 {-

567 Note [Arity analysis]

568 ~~~~~~~~~~~~~~~~~~~~~

569 The motivating example for arity analysis is this:

571 f = \x. let g = f (x+1)

572 in \y. ...g...

574 What arity does f have? Really it should have arity 2, but a naive

575 look at the RHS won't see that. You need a fixpoint analysis which

576 says it has arity "infinity" the first time round.

578 This example happens a lot; it first showed up in Andy Gill's thesis,

579 fifteen years ago! It also shows up in the code for 'rnf' on lists

580 in Trac #4138.

582 The analysis is easy to achieve because exprEtaExpandArity takes an

583 argument

584 type CheapFun = CoreExpr -> Maybe Type -> Bool

585 used to decide if an expression is cheap enough to push inside a

586 lambda. And exprIsCheap' in turn takes an argument

587 type CheapAppFun = Id -> Int -> Bool

588 which tells when an application is cheap. This makes it easy to

589 write the analysis loop.

591 The analysis is cheap-and-cheerful because it doesn't deal with

592 mutual recursion. But the self-recursive case is the important one.

595 Note [Eta expanding through dictionaries]

596 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

597 If the experimental -fdicts-cheap flag is on, we eta-expand through

598 dictionary bindings. This improves arities. Thereby, it also

599 means that full laziness is less prone to floating out the

600 application of a function to its dictionary arguments, which

601 can thereby lose opportunities for fusion. Example:

602 foo :: Ord a => a -> ...

603 foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....

604 -- So foo has arity 1

606 f = \x. foo dInt $ bar x

608 The (foo DInt) is floated out, and makes ineffective a RULE

609 foo (bar x) = ...

611 One could go further and make exprIsCheap reply True to any

612 dictionary-typed expression, but that's more work.

614 See Note [Dictionary-like types] in TcType.hs for why we use

615 isDictLikeTy here rather than isDictTy

617 Note [Eta expanding thunks]

618 ~~~~~~~~~~~~~~~~~~~~~~~~~~~

619 We don't eta-expand

620 * Trivial RHSs x = y

621 * PAPs x = map g

622 * Thunks f = case y of p -> \x -> blah

624 When we see

625 f = case y of p -> \x -> blah

626 should we eta-expand it? Well, if 'x' is a one-shot state token

627 then 'yes' because 'f' will only be applied once. But otherwise

628 we (conservatively) say no. My main reason is to avoid expanding

629 PAPSs

630 f = g d ==> f = \x. g d x

631 because that might in turn make g inline (if it has an inline pragma),

632 which we might not want. After all, INLINE pragmas say "inline only

633 when saturated" so we don't want to be too gung-ho about saturating!

634 -}

641 -- We have something like (let x = E in b),

642 -- where b has the given arity type.

646 -- If E is not cheap, keep arity only for one-shots

649 -- Processing (fun arg) where at is the ArityType of fun,

650 -- Knock off an argument and behave like 'let'

667 {-

668 Note [Combining case branches]

669 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

670 Consider

671 go = \x. let z = go e0

672 go2 = \x. case x of

673 True -> z

674 False -> \s(one-shot). e1

675 in go2 x

676 We *really* want to eta-expand go and go2.

677 When combining the barnches of the case we have

678 ATop [] `andAT` ATop [OneShotLam]

679 and we want to get ATop [OneShotLam]. But if the inner

680 lambda wasn't one-shot we don't want to do this.

681 (We need a proper arity analysis to justify that.)

683 So we combine the best of the two branches, on the (slightly dodgy)

684 basis that if we know one branch is one-shot, then they all must be.

685 -}

687 ---------------------------

689 -- How to decide if an expression is cheap

690 -- If the Maybe is Just, the type is the type

691 -- of the expression; Nothing means "don't know"

693 data ArityEnv

696 }

704 where

706 -- See Note [exprArity invariant] (2); must be true of

707 -- arityType too, since that is how we compute the arity

708 -- of variables, and they in turn affect result of exprArity

709 -- Trac #5441 is a nice demo

710 -- However, do make sure that ATop -> ATop and ABot -> ABot!

711 -- Casts don't affect that part. Getting this wrong provoked #5475

714 | strict_sig <- idStrictness v

720 | otherwise

722 where

726 -- Lambdas; increase arity

731 -- Applications; decrease arity, except for types

733 = arityType env fun

737 -- Case/Let; keep arity if either the expression is cheap

738 -- or it's a 1-shot lambda

739 -- The former is not really right for Haskell

740 -- f x = case x of { (a,b) -> \y. e }

741 -- ===>

742 -- f x y = case x of { (a,b) -> e }

743 -- The difference is observable using 'seq'

744 --

746 | exprIsBottom scrut || null alts

748 -- See Note [Dealing with bottom (1)]

749 | otherwise

753 -- See Note [Dealing with bottom (2)]

759 where

764 where

772 arityType _ _ = vanillaArityType

774 {-

775 %************************************************************************

776 %* *

777 The main eta-expander

778 %* *

779 %************************************************************************

781 We go for:

782 f = \x1..xn -> N ==> f = \x1..xn y1..ym -> N y1..ym

783 (n >= 0)

785 where (in both cases)

787 * The xi can include type variables

789 * The yi are all value variables

791 * N is a NORMAL FORM (i.e. no redexes anywhere)

792 wanting a suitable number of extra args.

794 The biggest reason for doing this is for cases like

796 f = \x -> case x of

797 True -> \y -> e1

798 False -> \y -> e2

800 Here we want to get the lambdas together. A good example is the nofib

801 program fibheaps, which gets 25% more allocation if you don't do this

802 eta-expansion.

804 We may have to sandwich some coerces between the lambdas

805 to make the types work. exprEtaExpandArity looks through coerces

806 when computing arity; and etaExpand adds the coerces as necessary when

807 actually computing the expansion.

809 Note [No crap in eta-expanded code]

810 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

811 The eta expander is careful not to introduce "crap". In particular,

812 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it

813 returns a CoreExpr satisfying the same invariant. See Note [Eta

814 expansion and the CorePrep invariants] in CorePrep.

816 This means the eta-expander has to do a bit of on-the-fly

817 simplification but it's not too hard. The alernative, of relying on

818 a subsequent clean-up phase of the Simplifier to de-crapify the result,

819 means you can't really use it in CorePrep, which is painful.

821 Note [Eta expansion and SCCs]

822 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

823 Note that SCCs are not treated specially by etaExpand. If we have

824 etaExpand 2 (\x -> scc "foo" e)

825 = (\xy -> (scc "foo" e) y)

826 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"

828 Note [Eta expansion and source notes]

829 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

830 CorePrep puts floatable ticks outside of value applications, but not

831 type applications. As a result we might be trying to eta-expand an

832 expression like

834 (src<...> v) @a

836 which we want to lead to code like

838 \x -> src<...> v @a x

840 This means that we need to look through type applications and be ready

841 to re-add floats on the top.

843 -}

845 -- | @etaExpand n e@ returns an expression with

846 -- the same meaning as @e@, but with arity @n@.

847 --

848 -- Given:

849 --

850 -- > e' = etaExpand n e

851 --

852 -- We should have that:

853 --

854 -- > ty = exprType e = exprType e'

857 -> CoreExpr

858 -- etaExpand deals with for-alls. For example:

859 -- etaExpand 1 E

860 -- where E :: forall a. a -> a

861 -- would return

862 -- (/\b. \y::a -> E b y)

863 --

864 -- It deals with coerces too, though they are now rare

865 -- so perhaps the extra code isn't worth it

867 etaExpand n orig_expr

868 = go n orig_expr

869 where

870 -- Strip off existing lambdas and casts

871 -- Note [Eta expansion and SCCs]

876 go n expr

879 where

884 -- Find ticks behind type apps.

885 -- See Note [Eta expansion and source notes]

891 -- Wrapper Unwrapper

892 --------------

894 -- \x. [], [] x

895 | EtaCo Coercion -- [] |> co, [] |> (sym co)

903 | isReflCo co = eis

905 where

906 co = co1 `mkTransCo` co2

910 --------------

916 --------------

918 -- (etaInfoApp s e eis) returns something equivalent to

919 -- ((substExpr s e) `appliedto` eis)

926 where

931 where

935 where

938 mk_alts_ty ty [] = ty

944 where

950 etaInfoApp subst e eis

952 where

953 go e [] = e

957 --------------

960 -- EtaInfo contains fresh variables,

961 -- not free in the incoming CoreExpr

962 -- Outgoing InScopeSet includes the EtaInfo vars

963 -- and the original free vars

965 mkEtaWW orig_n orig_expr in_scope orig_ty

967 where

968 empty_subst = mkEmptyTCvSubst in_scope

970 go n subst ty eis -- See Note [exprArity invariant]

976 -- Avoid free vars of the original expression

981 -- Avoid free vars of the original expression

986 -- newtype T = MkT ([T] -> Int)

987 -- Consider eta-expanding this

988 -- eta_expand 1 e T

989 -- We want to get

990 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)

994 -- but its type isn't a function.

997 -- This *can* legitmately happen:

998 -- e.g. coerce Int (\x. x) Essentially the programmer is

999 -- playing fast and loose with types (Happy does this a lot).

1000 -- So we simply decline to eta-expand. Otherwise we'd end up

1001 -- with an explicit lambda having a non-function type

1004 --------------

1005 -- Avoiding unnecessary substitution; use short-cutting versions

1011 subst_bind = substBindSC

1014 --------------

1016 -- Make a fresh Id, with specified type (after applying substitution)

1017 -- It should be "fresh" in the sense that it's not in the in-scope set

1018 -- of the TvSubstEnv; and it should itself then be added to the in-scope

1019 -- set of the TvSubstEnv

1020 --

1021 -- The Int is just a reasonable starting point for generating a unique;

1022 -- it does not necessarily have to be unique itself.

1023 freshEtaId n subst ty

1025 where