compiler/coreSyn/CoreArity.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   4 %
   5
   6         Arity and ete expansion
   7
   8 \begin{code}
   9 -- | Arit and eta expansion
  10 module CoreArity (
  11         manifestArity, exprArity,
  12         exprEtaExpandArity, etaExpand
  13     ) where
  14
  15 #include "HsVersions.h"
  16
  17 import CoreSyn
  18 import CoreFVs
  19 import CoreUtils
  20 import qualified CoreSubst
  21 import CoreSubst ( Subst, substBndr, substBndrs, substExpr
  22                  , mkEmptySubst, isEmptySubst )
  23 import Var
  24 import VarEnv
  25 import Id
  26 import Type
  27 import TcType   ( isDictLikeTy )
  28 import Coercion
  29 import BasicTypes
  30 import Unique
  31 import Outputable
  32 import DynFlags
  33 import FastString
  34 import Maybes
  35
  36 import GHC.Exts         -- For `xori`
  37 \end{code}
  38
  39 %************************************************************************
  40 %*                                                                      *
  41               manifestArity and exprArity
  42 %*                                                                      *
  43 %************************************************************************
  44
  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.
  51
  52 Originally I thought that it was enough just to look for top-level lambdas, but
  53 it isn't.  I've seen this
  54
  55         foo = PrelBase.timesInt
  56
  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).
  60
  61 Similarly, see the ok_note check in exprEtaExpandArity.  So
  62         f = __inline_me (\x -> e)
  63 won't be eta-expanded.
  64
  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
  69 Note [exprArity invariant]
  70 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  71 exprArity has the following invariant:
  72         (exprArity e) = n, then manifestArity (etaExpand e n) = n
  73
  74 That is, if exprArity says "the arity is n" then etaExpand really can get
  75 "n" manifest lambdas to the top.
  76
  77 Why is this important?  Because
  78   - In TidyPgm we use exprArity to fix the *final arity* of
  79     each top-level Id, and in
  80   - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
  81     actually match that arity, which in turn means
  82     that the StgRhs has the right number of lambdas
  83
  84 An alternative would be to do the eta-expansion in TidyPgm, at least
  85 for top-level bindings, in which case we would not need the trim_arity
  86 in exprArity.  That is a less local change, so I'm going to leave it for today!
  87
  88
  89 \begin{code}
  90 manifestArity :: CoreExpr -> Arity
  91 -- ^ manifestArity sees how many leading value lambdas there are
  92 manifestArity (Lam v e) | isId v    = 1 + manifestArity e
  93                         | otherwise = manifestArity e
  94 manifestArity (Note _ e)            = manifestArity e
  95 manifestArity (Cast e _)            = manifestArity e
  96 manifestArity _                     = 0
  97
  98 exprArity :: CoreExpr -> Arity
  99 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
 100 exprArity e = go e
 101   where
 102     go (Var v)                   = idArity v
 103     go (Lam x e) | isId x        = go e + 1
 104                  | otherwise     = go e
 105     go (Note _ e)                = go e
 106     go (Cast e co)               = trim_arity (go e) 0 (snd (coercionKind co))
 107     go (App e (Type _))          = go e
 108     go (App f a) | exprIsCheap a = (go f - 1) `max` 0
 109         -- NB: exprIsCheap a!
 110         --      f (fac x) does not have arity 2,
 111         --      even if f has arity 3!
 112         -- NB: `max 0`!  (\x y -> f x) has arity 2, even if f is
 113         --               unknown, hence arity 0
 114     go _                           = 0
 115
 116         -- Note [exprArity invariant]
 117     trim_arity n a ty
 118         | n==a                                        = a
 119         | Just (_, ty') <- splitForAllTy_maybe ty     = trim_arity n a     ty'
 120         | Just (_, ty') <- splitFunTy_maybe ty        = trim_arity n (a+1) ty'
 121         | Just (ty',_)  <- splitNewTypeRepCo_maybe ty = trim_arity n a     ty'
 122         | otherwise                                   = a
 123 \end{code}
 124
 125 %************************************************************************
 126 %*                                                                      *
 127 \subsection{Eta reduction and expansion}
 128 %*                                                                      *
 129 %************************************************************************
 130
 131 exprEtaExpandArity is used when eta expanding
 132         e  ==>  \xy -> e x y
 133
 134 It returns 1 (or more) to:
 135         case x of p -> \s -> ...
 136 because for I/O ish things we really want to get that \s to the top.
 137 We are prepared to evaluate x each time round the loop in order to get that
 138
 139 It's all a bit more subtle than it looks:
 140
 141 1.  One-shot lambdas
 142
 143 Consider one-shot lambdas
 144                 let x = expensive in \y z -> E
 145 We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
 146 Hence the ArityType returned by arityType
 147
 148 2.  The state-transformer hack
 149
 150 The one-shot lambda special cause is particularly important/useful for
 151 IO state transformers, where we often get
 152         let x = E in \ s -> ...
 153
 154 and the \s is a real-world state token abstraction.  Such abstractions
 155 are almost invariably 1-shot, so we want to pull the \s out, past the
 156 let x=E, even if E is expensive.  So we treat state-token lambdas as
 157 one-shot even if they aren't really.  The hack is in Id.isOneShotBndr.
 158
 159 3.  Dealing with bottom
 160
 161 Consider also
 162         f = \x -> error "foo"
 163 Here, arity 1 is fine.  But if it is
 164         f = \x -> case x of
 165                         True  -> error "foo"
 166                         False -> \y -> x+y
 167 then we want to get arity 2.  Tecnically, this isn't quite right, because
 168         (f True) `seq` 1
 169 should diverge, but it'll converge if we eta-expand f.  Nevertheless, we
 170 do so; it improves some programs significantly, and increasing convergence
 171 isn't a bad thing.  Hence the ABot/ATop in ArityType.
 172
 173 Actually, the situation is worse.  Consider
 174         f = \x -> case x of
 175                         True  -> \y -> x+y
 176                         False -> \y -> x-y
 177 Can we eta-expand here?  At first the answer looks like "yes of course", but
 178 consider
 179         (f bot) `seq` 1
 180 This should diverge!  But if we eta-expand, it won't.   Again, we ignore this
 181 "problem", because being scrupulous would lose an important transformation for
 182 many programs.
 183
 184
 185 4. Newtypes
 186
 187 Non-recursive newtypes are transparent, and should not get in the way.
 188 We do (currently) eta-expand recursive newtypes too.  So if we have, say
 189
 190         newtype T = MkT ([T] -> Int)
 191
 192 Suppose we have
 193         e = coerce T f
 194 where f has arity 1.  Then: etaExpandArity e = 1;
 195 that is, etaExpandArity looks through the coerce.
 196
 197 When we eta-expand e to arity 1: eta_expand 1 e T
 198 we want to get:                  coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 199
 200 HOWEVER, note that if you use coerce bogusly you can ge
 201         coerce Int negate
 202 And since negate has arity 2, you might try to eta expand.  But you can't
 203 decopose Int to a function type.   Hence the final case in eta_expand.
 204
 205
 206 \begin{code}
 207 -- ^ The Arity returned is the number of value args the
 208 -- expression can be applied to without doing much work
 209 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
 210 exprEtaExpandArity dflags e = arityDepth (arityType dflags e)
 211
 212 -- A limited sort of function type
 213 data ArityType = AFun Bool ArityType    -- True <=> one-shot
 214                | ATop                   -- Know nothing
 215                | ABot                   -- Diverges
 216
 217 arityDepth :: ArityType -> Arity
 218 arityDepth (AFun _ ty) = 1 + arityDepth ty
 219 arityDepth _           = 0
 220
 221 andArityType :: ArityType -> ArityType -> ArityType
 222 andArityType ABot           at2           = at2
 223 andArityType ATop           _             = ATop
 224 andArityType (AFun t1 at1)  (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
 225 andArityType at1            at2           = andArityType at2 at1
 226
 227 arityType :: DynFlags -> CoreExpr -> ArityType
 228         -- (go1 e) = [b1,..,bn]
 229         -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
 230         -- where bi is True <=> the lambda is one-shot
 231
 232 arityType dflags (Note _ e) = arityType dflags e
 233 --      Not needed any more: etaExpand is cleverer
 234 -- removed: | ok_note n = arityType dflags e
 235 -- removed: | otherwise = ATop
 236
 237 arityType dflags (Cast e _) = arityType dflags e
 238
 239 arityType _ (Var v)
 240   = mk (idArity v) (arg_tys (idType v))
 241   where
 242     mk :: Arity -> [Type] -> ArityType
 243         -- The argument types are only to steer the "state hack"
 244         -- Consider case x of
 245         --              True  -> foo
 246         --              False -> \(s:RealWorld) -> e
 247         -- where foo has arity 1.  Then we want the state hack to
 248         -- apply to foo too, so we can eta expand the case.
 249     mk 0 tys | isBottomingId v                   = ABot
 250              | (ty:_) <- tys, isStateHackType ty = AFun True ATop
 251              | otherwise                         = ATop
 252     mk n (ty:tys) = AFun (isStateHackType ty) (mk (n-1) tys)
 253     mk n []       = AFun False                (mk (n-1) [])
 254
 255     arg_tys :: Type -> [Type]   -- Ignore for-alls
 256     arg_tys ty
 257         | Just (_, ty')  <- splitForAllTy_maybe ty = arg_tys ty'
 258         | Just (arg,res) <- splitFunTy_maybe ty    = arg : arg_tys res
 259         | otherwise                                = []
 260
 261         -- Lambdas; increase arity
 262 arityType dflags (Lam x e)
 263   | isId x    = AFun (isOneShotBndr x) (arityType dflags e)
 264   | otherwise = arityType dflags e
 265
 266         -- Applications; decrease arity
 267 arityType dflags (App f (Type _)) = arityType dflags f
 268 arityType dflags (App f a)
 269    = case arityType dflags f of
 270         ABot -> ABot    -- If function diverges, ignore argument
 271         ATop -> ATop    -- No no info about function
 272         AFun _ xs
 273                 | exprIsCheap a -> xs
 274                 | otherwise     -> ATop
 275
 276         -- Case/Let; keep arity if either the expression is cheap
 277         -- or it's a 1-shot lambda
 278         -- The former is not really right for Haskell
 279         --      f x = case x of { (a,b) -> \y. e }
 280         --  ===>
 281         --      f x y = case x of { (a,b) -> e }
 282         -- The difference is observable using 'seq'
 283 arityType dflags (Case scrut _ _ alts)
 284   = case foldr1 andArityType [arityType dflags rhs | (_,_,rhs) <- alts] of
 285         xs | exprIsCheap scrut     -> xs
 286         AFun one_shot _ | one_shot -> AFun True ATop
 287         _                          -> ATop
 288
 289 arityType dflags (Let b e)
 290   = case arityType dflags e of
 291         xs              | cheap_bind b -> xs
 292         AFun one_shot _ | one_shot     -> AFun True ATop
 293         _                              -> ATop
 294   where
 295     cheap_bind (NonRec b e) = is_cheap (b,e)
 296     cheap_bind (Rec prs)    = all is_cheap prs
 297     is_cheap (b,e) = (dopt Opt_DictsCheap dflags && isDictLikeTy (idType b))
 298                    || exprIsCheap e
 299         -- If the experimental -fdicts-cheap flag is on, we eta-expand through
 300         -- dictionary bindings.  This improves arities. Thereby, it also
 301         -- means that full laziness is less prone to floating out the
 302         -- application of a function to its dictionary arguments, which
 303         -- can thereby lose opportunities for fusion.  Example:
 304         --      foo :: Ord a => a -> ...
 305         --      foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
 306         --              -- So foo has arity 1
 307         --
 308         --      f = \x. foo dInt $ bar x
 309         --
 310         -- The (foo DInt) is floated out, and makes ineffective a RULE
 311         --      foo (bar x) = ...
 312         --
 313         -- One could go further and make exprIsCheap reply True to any
 314         -- dictionary-typed expression, but that's more work.
 315         --
 316         -- See Note [Dictionary-like types] in TcType.lhs for why we use
 317         -- isDictLikeTy here rather than isDictTy
 318
 319 arityType _ _ = ATop
 320 \end{code}
 321
 322
 323 %************************************************************************
 324 %*                                                                      *
 325               The main eta-expander
 326 %*                                                                      *
 327 %************************************************************************
 328
 329 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
 330 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
 331 CorePrep), it returns a CoreExpr satisfying the same invariant. See
 332 Note [Eta expansion and the CorePrep invariants] in CorePrep.
 333
 334 This means the eta-expander has to do a bit of on-the-fly
 335 simplification but it's not too hard.  The alernative, of relying on
 336 a subsequent clean-up phase of the Simplifier to de-crapify the result,
 337 means you can't really use it in CorePrep, which is painful.
 338
 339 \begin{code}
 340 -- | @etaExpand n us e ty@ returns an expression with
 341 -- the same meaning as @e@, but with arity @n@.
 342 --
 343 -- Given:
 344 --
 345 -- > e' = etaExpand n us e ty
 346 --
 347 -- We should have that:
 348 --
 349 -- > ty = exprType e = exprType e'
 350 etaExpand :: Arity              -- ^ Result should have this number of value args
 351           -> CoreExpr           -- ^ Expression to expand
 352           -> CoreExpr
 353 -- Note that SCCs are not treated specially.  If we have
 354 --      etaExpand 2 (\x -> scc "foo" e)
 355 --      = (\xy -> (scc "foo" e) y)
 356 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
 357
 358 -- etaExpand deals with for-alls. For example:
 359 --              etaExpand 1 E
 360 -- where  E :: forall a. a -> a
 361 -- would return
 362 --      (/\b. \y::a -> E b y)
 363 --
 364 -- It deals with coerces too, though they are now rare
 365 -- so perhaps the extra code isn't worth it
 366
 367 etaExpand n orig_expr
 368   | manifestArity orig_expr >= n = orig_expr    -- The no-op case
 369   | otherwise
 370   = go n orig_expr
 371   where
 372       -- Strip off existing lambdas
 373     go 0 expr = expr
 374     go n (Lam v body) | isTyVar v = Lam v (go n     body)
 375                       | otherwise = Lam v (go (n-1) body)
 376     go n (Note InlineMe expr) = Note InlineMe (go n expr)
 377         -- Note [Eta expansion and SCCs]
 378     go n (Cast expr co) = Cast (go n expr) co
 379     go n expr           = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
 380                           etaInfoAbs etas (etaInfoApp subst' expr etas)
 381                         where
 382                             in_scope = mkInScopeSet (exprFreeVars expr)
 383                             (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
 384                             subst' = mkEmptySubst in_scope'
 385
 386                                 -- Wrapper    Unwrapper
 387 --------------
 388 data EtaInfo = EtaVar Var       -- /\a. [],   [] a
 389                                 -- \x.  [],   [] x
 390              | EtaCo Coercion   -- [] |> co,  [] |> (sym co)
 391
 392 instance Outputable EtaInfo where
 393    ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
 394    ppr (EtaCo co) = ptext (sLit "EtaCo")  <+> ppr co
 395
 396 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
 397 pushCoercion co1 (EtaCo co2 : eis)
 398   | isIdentityCoercion co = eis
 399   | otherwise             = EtaCo co : eis
 400   where
 401     co = co1 `mkTransCoercion` co2
 402
 403 pushCoercion co eis = EtaCo co : eis
 404
 405 --------------
 406 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
 407 etaInfoAbs []               expr = expr
 408 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
 409 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
 410
 411 --------------
 412 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
 413 -- (etaInfoApp s e eis) returns something equivalent to
 414 --             ((substExpr s e) `appliedto` eis)
 415
 416 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
 417   = etaInfoApp subst' e eis
 418   where
 419     subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
 420            | otherwise  = CoreSubst.extendIdSubst subst v1 (Var v2)
 421
 422 etaInfoApp subst (Cast e co1) eis
 423   = etaInfoApp subst e (pushCoercion co' eis)
 424   where
 425     co' = CoreSubst.substTy subst co1
 426
 427 etaInfoApp subst (Case e b _ alts) eis
 428   = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
 429   where
 430     (subst1, b1) = substBndr subst b
 431     alts' = map subst_alt alts
 432     subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
 433               where
 434                  (subst2,bs') = substBndrs subst1 bs
 435
 436 etaInfoApp subst (Let b e) eis
 437   = Let b' (etaInfoApp subst' e eis)
 438   where
 439     (subst', b') = subst_bind subst b
 440
 441 etaInfoApp subst (Note note e) eis
 442   = Note note (etaInfoApp subst e eis)
 443
 444 etaInfoApp subst e eis
 445   = go (subst_expr subst e) eis
 446   where
 447     go e []                  = e
 448     go e (EtaVar v    : eis) = go (App e (varToCoreExpr v)) eis
 449     go e (EtaCo co    : eis) = go (Cast e co) eis
 450
 451 --------------
 452 mkEtaWW :: Arity -> InScopeSet -> Type
 453         -> (InScopeSet, [EtaInfo])
 454         -- EtaInfo contains fresh variables,
 455         --   not free in the incoming CoreExpr
 456         -- Outgoing InScopeSet includes the EtaInfo vars
 457         --   and the original free vars
 458
 459 mkEtaWW n in_scope ty
 460   = go n empty_subst ty []
 461   where
 462     empty_subst = mkTvSubst in_scope emptyTvSubstEnv
 463
 464     go n subst ty eis
 465        | n == 0
 466        = (getTvInScope subst, reverse eis)
 467
 468        | Just (tv,ty') <- splitForAllTy_maybe ty
 469        , let (subst', tv') = substTyVarBndr subst tv
 470            -- Avoid free vars of the original expression
 471        = go n subst' ty' (EtaVar tv' : eis)
 472
 473        | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
 474        , let (subst', eta_id') = freshEtaId n subst arg_ty
 475            -- Avoid free vars of the original expression
 476        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
 477
 478        | Just(ty',co) <- splitNewTypeRepCo_maybe ty
 479        =        -- Given this:
 480                 --      newtype T = MkT ([T] -> Int)
 481                 -- Consider eta-expanding this
 482                 --      eta_expand 1 e T
 483                 -- We want to get
 484                 --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 485          go n subst ty' (EtaCo (substTy subst co) : eis)
 486
 487        | otherwise                         -- We have an expression of arity > 0,
 488        = (getTvInScope subst, reverse eis) -- but its type isn't a function.
 489         -- This *can* legitmately happen:
 490         -- e.g.  coerce Int (\x. x) Essentially the programmer is
 491         -- playing fast and loose with types (Happy does this a lot).
 492         -- So we simply decline to eta-expand.  Otherwise we'd end up
 493         -- with an explicit lambda having a non-function type
 494
 495
 496 --------------
 497 -- Avoiding unnecessary substitution
 498
 499 subst_expr :: Subst -> CoreExpr -> CoreExpr
 500 subst_expr s e | isEmptySubst s = e
 501                | otherwise      = substExpr s e
 502
 503 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
 504 subst_bind subst (NonRec b r)
 505   = (subst', NonRec b' (subst_expr subst r))
 506   where
 507     (subst', b') = substBndr subst b
 508 subst_bind subst (Rec prs)
 509   = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
 510   where
 511     (bs, rhss) = unzip prs
 512     (subst', bs1) = substBndrs subst bs
 513
 514
 515 --------------
 516 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
 517 -- Make a fresh Id, with specified type (after applying substitution)
 518 -- It should be "fresh" in the sense that it's not in the in-scope set
 519 -- of the TvSubstEnv; and it should itself then be added to the in-scope
 520 -- set of the TvSubstEnv
 521 --
 522 -- The Int is just a reasonable starting point for generating a unique;
 523 -- it does not necessarily have to be unique itself.
 524 freshEtaId n subst ty
 525       = (subst', eta_id')
 526       where
 527         ty'     = substTy subst ty
 528         eta_id' = uniqAway (getTvInScope subst) $
 529                   mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
 530         subst'  = extendTvInScope subst [eta_id']
 531 \end{code}
 532