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