-- | Arit and eta expansion
module CoreArity (
manifestArity, exprArity, exprBotStrictness_maybe,
- exprEtaExpandArity, etaExpand
+ exprEtaExpandArity, CheapFun, etaExpand
) where
#include "HsVersions.h"
import Id
import Type
import TyCon ( isRecursiveTyCon, isClassTyCon )
-import TcType ( isDictLikeTy )
import Coercion
import BasicTypes
import Unique
import Outputable
-import DynFlags
import FastString
+import Pair
\end{code}
%************************************************************************
go (Lam x e) | isId x = go e + 1
| otherwise = go e
go (Note n e) | notSccNote n = go e
- go (Cast e co) = go e `min` length (typeArity (snd (coercionKind co)))
- -- Note [exprArity invariant]
+ go (Cast e co) = go e `min` length (typeArity (pSnd (coercionKind co)))
+ -- Note [exprArity invariant]
go (App e (Type _)) = go e
go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
-- See Note [exprArity for applications]
+ -- NB: coercions count as a value argument
+
go _ = 0
-- and gives them a suitable strictness signatures. It's used during
-- float-out
exprBotStrictness_maybe e
- = case getBotArity (arityType False e) of
+ = case getBotArity (arityType is_cheap e) of
Nothing -> Nothing
Just ar -> Just (ar, mkStrictSig (mkTopDmdType (replicate ar topDmd) BotRes))
+ where
+ is_cheap _ _ = False -- Irrelevant for this purpose
\end{code}
Note [exprArity invariant]
-- ^ The Arity returned is the number of value args the [_$_]
-- expression can be applied to without doing much work
-exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
+exprEtaExpandArity :: CheapFun -> CoreExpr -> Arity
-- exprEtaExpandArity is used when eta expanding
-- e ==> \xy -> e x y
-exprEtaExpandArity dflags e
- = case (arityType dicts_cheap e) of
+exprEtaExpandArity cheap_fun e
+ = case (arityType cheap_fun e) of
ATop (os:oss)
| os || has_lam e -> 1 + length oss -- Note [Eta expanding thunks]
| otherwise -> 0
ATop [] -> 0
ABot n -> n
where
- dicts_cheap = dopt Opt_DictsCheap dflags
has_lam (Note _ e) = has_lam e
has_lam (Lam b e) = isId b || has_lam e
has_lam _ = False
floatIn False (ATop as) = ATop (takeWhile id as)
-- If E is not cheap, keep arity only for one-shots
-arityApp :: ArityType -> CoreExpr -> ArityType
+arityApp :: ArityType -> Bool -> ArityType
-- Processing (fun arg) where at is the ArityType of fun,
-- Knock off an argument and behave like 'let'
-arityApp (ABot 0) _ = ABot 0
-arityApp (ABot n) _ = ABot (n-1)
-arityApp (ATop []) _ = ATop []
-arityApp (ATop (_:as)) arg = floatIn (exprIsCheap arg) (ATop as)
+arityApp (ABot 0) _ = ABot 0
+arityApp (ABot n) _ = ABot (n-1)
+arityApp (ATop []) _ = ATop []
+arityApp (ATop (_:as)) cheap = floatIn cheap (ATop as)
andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
andArityType (ABot n1) (ABot n2)
\begin{code}
---------------------------
-arityType :: Bool -> CoreExpr -> ArityType
+type CheapFun = CoreExpr -> Maybe Type -> Bool
+ -- How to decide if an expression is cheap
+ -- If the Maybe is Just, the type is the type
+ -- of the expression; Nothing means "don't know"
+
+arityType :: CheapFun -> CoreExpr -> ArityType
arityType _ (Var v)
| Just strict_sig <- idStrictness_maybe v
, (ds, res) <- splitStrictSig strict_sig
one_shots = typeArity (idType v)
-- Lambdas; increase arity
-arityType dicts_cheap (Lam x e)
- | isId x = arityLam x (arityType dicts_cheap e)
- | otherwise = arityType dicts_cheap e
+arityType cheap_fn (Lam x e)
+ | isId x = arityLam x (arityType cheap_fn e)
+ | otherwise = arityType cheap_fn e
- -- Applications; decrease arity
-arityType dicts_cheap (App fun (Type _))
- = arityType dicts_cheap fun
-arityType dicts_cheap (App fun arg )
- = arityApp (arityType dicts_cheap fun) arg
+ -- Applications; decrease arity, except for types
+arityType cheap_fn (App fun (Type _))
+ = arityType cheap_fn fun
+arityType cheap_fn (App fun arg )
+ = arityApp (arityType cheap_fn fun) (cheap_fn arg Nothing)
-- Case/Let; keep arity if either the expression is cheap
-- or it's a 1-shot lambda
-- ===>
-- f x y = case x of { (a,b) -> e }
-- The difference is observable using 'seq'
-arityType dicts_cheap (Case scrut _ _ alts)
- = floatIn (exprIsCheap scrut)
- (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
+arityType cheap_fn (Case scrut bndr _ alts)
+ = floatIn (cheap_fn scrut (Just (idType bndr)))
+ (foldr1 andArityType [arityType cheap_fn rhs | (_,_,rhs) <- alts])
-arityType dicts_cheap (Let b e)
- = floatIn (cheap_bind b) (arityType dicts_cheap e)
+arityType cheap_fn (Let b e)
+ = floatIn (cheap_bind b) (arityType cheap_fn e)
where
cheap_bind (NonRec b e) = is_cheap (b,e)
cheap_bind (Rec prs) = all is_cheap prs
- is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
- || exprIsCheap e
- -- If the experimental -fdicts-cheap flag is on, we eta-expand through
- -- dictionary bindings. This improves arities. Thereby, it also
- -- means that full laziness is less prone to floating out the
- -- application of a function to its dictionary arguments, which
- -- can thereby lose opportunities for fusion. Example:
- -- foo :: Ord a => a -> ...
- -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
- -- -- So foo has arity 1
- --
- -- f = \x. foo dInt $ bar x
- --
- -- The (foo DInt) is floated out, and makes ineffective a RULE
- -- foo (bar x) = ...
- --
- -- One could go further and make exprIsCheap reply True to any
- -- dictionary-typed expression, but that's more work.
- --
- -- See Note [Dictionary-like types] in TcType.lhs for why we use
- -- isDictLikeTy here rather than isDictTy
-
-arityType dicts_cheap (Note n e)
- | notSccNote n = arityType dicts_cheap e
-arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
-arityType _ _ = vanillaArityType
+ is_cheap (b,e) = cheap_fn e (Just (idType b))
+
+arityType cheap_fn (Note n e)
+ | notSccNote n = arityType cheap_fn e
+arityType cheap_fn (Cast e _) = arityType cheap_fn e
+arityType _ _ = vanillaArityType
\end{code}
-- Strip off existing lambdas and casts
-- Note [Eta expansion and SCCs]
go 0 expr = expr
- go n (Lam v body) | isTyCoVar v = Lam v (go n body)
- | otherwise = Lam v (go (n-1) body)
+ go n (Lam v body) | isTyVar v = Lam v (go n body)
+ | otherwise = Lam v (go (n-1) body)
go n (Cast expr co) = Cast (go n expr) co
go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
etaInfoAbs etas (etaInfoApp subst' expr etas)
where
in_scope = mkInScopeSet (exprFreeVars expr)
- (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
+ (in_scope', etas) = mkEtaWW n orig_expr in_scope (exprType expr)
subst' = mkEmptySubst in_scope'
-- Wrapper Unwrapper
pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion co1 (EtaCo co2 : eis)
- | isIdentityCoercion co = eis
- | otherwise = EtaCo co : eis
+ | isReflCo co = eis
+ | otherwise = EtaCo co : eis
where
- co = co1 `mkTransCoercion` co2
+ co = co1 `mkTransCo` co2
pushCoercion co eis = EtaCo co : eis
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [] expr = expr
etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
-etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
+etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCo co)
--------------
etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
-- ((substExpr s e) `appliedto` eis)
etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
- = etaInfoApp subst' e eis
- where
- subst' | isTyCoVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
- | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
+ = etaInfoApp (CoreSubst.extendSubstWithVar subst v1 v2) e eis
etaInfoApp subst (Cast e co1) eis
= etaInfoApp subst e (pushCoercion co' eis)
where
- co' = CoreSubst.substTy subst co1
+ co' = CoreSubst.substCo subst co1
etaInfoApp subst (Case e b _ alts) eis
= Case (subst_expr subst e) b1 (coreAltsType alts') alts'
go e (EtaCo co : eis) = go (Cast e co) eis
--------------
-mkEtaWW :: Arity -> InScopeSet -> Type
+mkEtaWW :: Arity -> CoreExpr -> InScopeSet -> Type
-> (InScopeSet, [EtaInfo])
-- EtaInfo contains fresh variables,
-- not free in the incoming CoreExpr
-- Outgoing InScopeSet includes the EtaInfo vars
-- and the original free vars
-mkEtaWW orig_n in_scope orig_ty
+mkEtaWW orig_n orig_expr in_scope orig_ty
= go orig_n empty_subst orig_ty []
where
- empty_subst = mkTvSubst in_scope emptyTvSubstEnv
+ empty_subst = TvSubst in_scope emptyTvSubstEnv
go n subst ty eis -- See Note [exprArity invariant]
| n == 0
= (getTvInScope subst, reverse eis)
| Just (tv,ty') <- splitForAllTy_maybe ty
- , let (subst', tv') = substTyVarBndr subst tv
+ , let (subst', tv') = Type.substTyVarBndr subst tv
-- Avoid free vars of the original expression
= go n subst' ty' (EtaVar tv' : eis)
-- eta_expand 1 e T
-- We want to get
-- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
- go n subst ty' (EtaCo (Type.substTy subst co) : eis)
+ go n subst ty' (EtaCo co : eis)
| otherwise -- We have an expression of arity > 0,
-- but its type isn't a function.
- = WARN( True, ppr orig_n <+> ppr orig_ty )
+ = WARN( True, (ppr orig_n <+> ppr orig_ty) $$ ppr orig_expr )
(getTvInScope subst, reverse eis)
-- This *can* legitmately happen:
-- e.g. coerce Int (\x. x) Essentially the programmer is
projects / ghc-hetmet.git / blobdiff