0f392aee2b84d58cc20f9974ae6295f3bec6a5cc

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

117 -- in TyCon

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

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

120 -- (no longer required)

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

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

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

125 --

126 -- AND through a layer of recursive newtypes

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

129 | otherwise

130 = []

132 ---------------

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

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

136 -- float-out

137 exprBotStrictness_maybe e

139 Nothing -> Nothing

141 where

144 -- For this purpose we can be very simple

146 {-

147 Note [exprArity invariant]

148 ~~~~~~~~~~~~~~~~~~~~~~~~~~

149 exprArity has the following invariant:

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

152 then manifestArity (etaExpand e n) = n

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

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

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

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

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

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

164 Why is this important? Because

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

166 each top-level Id, and in

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

168 actually match that arity, which in turn means

169 that the StgRhs has the right number of lambdas

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

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

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

175 Note [Newtype classes and eta expansion]

176 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

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

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

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

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

184 rule mechanism. Consider

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

187 instance C b => C [b] where

188 op x = ...

190 These translate to

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

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

195 $copList d x = ...

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

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

201 Now suppose we have:

209 blah = $copList dCInt

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

225 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

235 have arity decreases.

241 ************************************************************************

242 * *

244 * *

245 ************************************************************************

248 ~~~~~~~~~~~~~~~~~~~~~~~~~~

251 converges rapidly

253 Or, to put it another way

255 there is no work lost in duplicating the partial

270 ~~~~~~~~~~~~~~~~~~~~~~~

271 Consider one-shot lambdas

276 ~~~~~~~~~~~~~~~~~~~~~~~~~~

277 A Big Deal with computing arities is expressions like

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

295 Consider also

308 have several tickets reporting unexpected bahaviour resulting from

319 going to diverge eventually anyway then getting the best arity

330 ~~~~~~~~~~~~~~~~~~~~~~~~

336 Suppose we have

337 e = coerce T f

350 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

351 Suppose we have

352 f = e

354 a usage type like

370 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

383 Extrude the g2

388 Discard args for bottomming function

393 Extrude g1.g3

406 After inlining (>>) we get

422 ~~~~~~~~~~~~~~~~

441 its definition.

445 themselves.

454 Example:

467 -------------------- Main arity code ----------------------------

468 -}

470 -- See Note [ArityType]

472 -- There is always an explicit lambda

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

475 vanillaArityType :: ArityType

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

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

481 -- exprEtaExpandArity is used when eta expanding

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

483 exprEtaExpandArity dflags e

486 ABot n -> n

487 where

492 -- Arity of a divergent function

494 getBotArity _ = Nothing

497 mk_cheap_fn dflags cheap_app

500 | otherwise

504 Just ty -> isDictLikeTy ty

507 ----------------------

509 -- This implements the fixpoint loop for arity analysis

510 -- See Note [Arity analysis]

511 findRhsArity dflags bndr rhs old_arity

513 -- We always call exprEtaExpandArity once, but usually

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

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

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

517 where

518 init_cheap_app :: CheapAppFun

519 init_cheap_app fn n_val_args

524 go cur_arity

528 #ifdef DEBUG

529 pprTrace "Exciting arity"

532 #endif

533 go new_arity

534 where

535 new_arity = rhsEtaExpandArity dflags cheap_app rhs

537 cheap_app :: CheapAppFun

538 cheap_app fn n_val_args

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

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

545 -- exprEtaExpandArity is used when eta expanding

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

547 rhsEtaExpandArity dflags cheap_app e

551 -- Don't expand PAPs/thunks

552 -- Note [Eta expanding thunks]

555 ABot n -> n

556 where

564 {-

565 Note [Arity analysis]

566 ~~~~~~~~~~~~~~~~~~~~~

567 The motivating example for arity analysis is this:

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

570 in \y. ...g...

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

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

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

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

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

578 in Trac #4138.

580 The analysis is easy to achieve because exprEtaExpandArity takes an

581 argument

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

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

584 lambda. And exprIsCheap' in turn takes an argument

585 type CheapAppFun = Id -> Int -> Bool

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

587 write the analysis loop.

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

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

593 Note [Eta expanding through dictionaries]

594 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

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

598 application of a function to its dictionary arguments, which

599 can thereby lose opportunities for fusion. Example:

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

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

602 -- So foo has arity 1

604 f = \x. foo dInt $ bar x

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

607 foo (bar x) = ...

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

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

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

613 isDictLikeTy here rather than isDictTy

615 Note [Eta expanding thunks]

616 ~~~~~~~~~~~~~~~~~~~~~~~~~~~

617 We don't eta-expand

618 * Trivial RHSs x = y

619 * PAPs x = map g

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

622 When we see

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

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

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

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

627 PAPSs

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

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

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

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

632 -}

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

640 -- where b has the given arity type.

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

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

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

665 {-

666 Note [Combining case branches]

667 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

668 Consider

669 go = \x. let z = go e0

670 go2 = \x. case x of

671 True -> z

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

673 in go2 x

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

675 When combining the barnches of the case we have

676 ATop [] `andAT` ATop [OneShotLam]

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

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

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

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

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

683 -}

685 ---------------------------

687 -- How to decide if an expression is cheap

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

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

691 data ArityEnv

694 }

702 where

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

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

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

707 -- Trac #5441 is a nice demo

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

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

712 | strict_sig <- idStrictness v

718 | otherwise

720 where

724 -- Lambdas; increase arity

729 -- Applications; decrease arity, except for types

731 = arityType env fun

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

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

737 -- The former is not really right for Haskell

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

739 -- ===>

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

741 -- The difference is observable using 'seq'

742 --

744 | exprIsBottom scrut || null alts

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

747 | otherwise

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

757 where

762 where

770 arityType _ _ = vanillaArityType

772 {-

773 %************************************************************************

774 %* *

775 The main eta-expander

776 %* *

777 %************************************************************************

779 We go for:

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

781 (n >= 0)

783 where (in both cases)

785 * The xi can include type variables

787 * The yi are all value variables

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

790 wanting a suitable number of extra args.

792 The biggest reason for doing this is for cases like

794 f = \x -> case x of

795 True -> \y -> e1

796 False -> \y -> e2

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

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

800 eta-expansion.

802 We may have to sandwich some coerces between the lambdas

803 to make the types work. exprEtaExpandArity looks through coerces

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

805 actually computing the expansion.

807 Note [No crap in eta-expanded code]

808 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

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

812 expansion and the CorePrep invariants] in CorePrep.

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

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

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

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

819 Note [Eta expansion and SCCs]

820 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

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

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

826 Note [Eta expansion and source notes]

827 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

830 expression like

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

834 which we want to lead to code like

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

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

839 to re-add floats on the top.

841 -}

843 -- | @etaExpand n us e ty@ returns an expression with

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

845 --

846 -- Given:

847 --

848 -- > e' = etaExpand n us e ty

849 --

850 -- We should have that:

851 --

852 -- > ty = exprType e = exprType e'

855 -> CoreExpr

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

857 -- etaExpand 1 E

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

859 -- would return

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

861 --

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

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

865 etaExpand n orig_expr

866 = go n orig_expr

867 where

868 -- Strip off existing lambdas and casts

869 -- Note [Eta expansion and SCCs]

874 go n expr

877 where

882 -- Find ticks behind type apps.

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

889 -- Wrapper Unwrapper

890 --------------

892 -- \x. [], [] x

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

901 | isReflCo co = eis

903 where

904 co = co1 `mkTransCo` co2

908 --------------

914 --------------

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

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

924 where

929 where

933 where

936 mk_alts_ty ty [] = ty

942 where

948 etaInfoApp subst e eis

950 where

951 go e [] = e

955 --------------

958 -- EtaInfo contains fresh variables,

959 -- not free in the incoming CoreExpr

960 -- Outgoing InScopeSet includes the EtaInfo vars

961 -- and the original free vars

963 mkEtaWW orig_n orig_expr in_scope orig_ty

965 where

966 empty_subst = mkEmptyTCvSubst in_scope

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

976 in

977 -- Avoid free vars of the original expression

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

983 -- Consider eta-expanding this

984 -- eta_expand 1 e T

985 -- We want to get

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

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

993 -- This *can* legitmately happen:

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

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

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

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

1000 --------------

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

1007 subst_bind = substBindSC

1010 --------------

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

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

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

1015 -- set of the TvSubstEnv

1016 --

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

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

1019 freshEtaVar n subst ty

1021 where