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