{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
module GHC.Core.Opt.Arity
( manifestArity, joinRhsArity, exprArity, typeArity
, exprEtaExpandArity, findRhsArity
, etaExpand, etaExpandAT
, etaExpandToJoinPoint, etaExpandToJoinPointRule
, exprBotStrictness_maybe
, ArityType(..), expandableArityType, arityTypeArity
, maxWithArity, isBotArityType, idArityType
)
where
#include "GhclibHsVersions.h"
import GHC.Prelude
import GHC.Core
import GHC.Core.FVs
import GHC.Core.Utils
import GHC.Core.Subst
import GHC.Types.Demand
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Types.Id
import GHC.Core.Type as Type
import GHC.Core.TyCon ( initRecTc, checkRecTc )
import GHC.Core.Predicate ( isDictTy )
import GHC.Core.Coercion as Coercion
import GHC.Core.Multiplicity
import GHC.Types.Var.Set
import GHC.Types.Basic
import GHC.Types.Unique
import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
import GHC.Utils.Outputable
import GHC.Data.FastString
import GHC.Utils.Misc ( lengthAtLeast )
manifestArity :: CoreExpr -> Arity
manifestArity :: CoreExpr -> Arity
manifestArity (Lam CoreBndr
v CoreExpr
e) | CoreBndr -> Bool
isId CoreBndr
v = Arity
1 Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
+ CoreExpr -> Arity
manifestArity CoreExpr
e
| Bool
otherwise = CoreExpr -> Arity
manifestArity CoreExpr
e
manifestArity (Tick Tickish CoreBndr
t CoreExpr
e) | Bool -> Bool
not (Tickish CoreBndr -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish CoreBndr
t) = CoreExpr -> Arity
manifestArity CoreExpr
e
manifestArity (Cast CoreExpr
e Coercion
_) = CoreExpr -> Arity
manifestArity CoreExpr
e
manifestArity CoreExpr
_ = Arity
0
joinRhsArity :: CoreExpr -> JoinArity
joinRhsArity :: CoreExpr -> Arity
joinRhsArity (Lam CoreBndr
_ CoreExpr
e) = Arity
1 Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
+ CoreExpr -> Arity
joinRhsArity CoreExpr
e
joinRhsArity CoreExpr
_ = Arity
0
exprArity :: CoreExpr -> Arity
exprArity :: CoreExpr -> Arity
exprArity CoreExpr
e = CoreExpr -> Arity
go CoreExpr
e
where
go :: CoreExpr -> Arity
go (Var CoreBndr
v) = CoreBndr -> Arity
idArity CoreBndr
v
go (Lam CoreBndr
x CoreExpr
e) | CoreBndr -> Bool
isId CoreBndr
x = CoreExpr -> Arity
go CoreExpr
e Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
+ Arity
1
| Bool
otherwise = CoreExpr -> Arity
go CoreExpr
e
go (Tick Tickish CoreBndr
t CoreExpr
e) | Bool -> Bool
not (Tickish CoreBndr -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish CoreBndr
t) = CoreExpr -> Arity
go CoreExpr
e
go (Cast CoreExpr
e Coercion
co) = Arity -> Type -> Arity
trim_arity (CoreExpr -> Arity
go CoreExpr
e) (Coercion -> Type
coercionRKind Coercion
co)
go (App CoreExpr
e (Type Type
_)) = CoreExpr -> Arity
go CoreExpr
e
go (App CoreExpr
f CoreExpr
a) | CoreExpr -> Bool
exprIsTrivial CoreExpr
a = (CoreExpr -> Arity
go CoreExpr
f Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
- Arity
1) Arity -> Arity -> Arity
forall a. Ord a => a -> a -> a
`max` Arity
0
go CoreExpr
_ = Arity
0
trim_arity :: Arity -> Type -> Arity
trim_arity :: Arity -> Type -> Arity
trim_arity Arity
arity Type
ty = Arity
arity Arity -> Arity -> Arity
forall a. Ord a => a -> a -> a
`min` [OneShotInfo] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length (Type -> [OneShotInfo]
typeArity Type
ty)
typeArity :: Type -> [OneShotInfo]
typeArity :: Type -> [OneShotInfo]
typeArity Type
ty
= RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
initRecTc Type
ty
where
go :: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
ty
| Just (CoreBndr
_, Type
ty') <- Type -> Maybe (CoreBndr, Type)
splitForAllTy_maybe Type
ty
= RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
ty'
| Just (Type
_,Type
arg,Type
res) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
= Type -> OneShotInfo
typeOneShot Type
arg OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
res
| Just (TyCon
tc,[Type]
tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty
, Just (Type
ty', Coercion
_) <- TyCon -> [Type] -> Maybe (Type, Coercion)
instNewTyCon_maybe TyCon
tc [Type]
tys
, Just RecTcChecker
rec_nts' <- RecTcChecker -> TyCon -> Maybe RecTcChecker
checkRecTc RecTcChecker
rec_nts TyCon
tc
= RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts' Type
ty'
| Bool
otherwise
= []
exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
exprBotStrictness_maybe CoreExpr
e
= case ArityType -> Maybe Arity
getBotArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e) of
Maybe Arity
Nothing -> Maybe (Arity, StrictSig)
forall a. Maybe a
Nothing
Just Arity
ar -> (Arity, StrictSig) -> Maybe (Arity, StrictSig)
forall a. a -> Maybe a
Just (Arity
ar, Arity -> StrictSig
sig Arity
ar)
where
env :: ArityEnv
env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_ped_bot :: Bool
ae_ped_bot = Bool
True
, ae_cheap_fn :: CheapFun
ae_cheap_fn = \ CoreExpr
_ Maybe Type
_ -> Bool
False
, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
sig :: Arity -> StrictSig
sig Arity
ar = [Demand] -> Divergence -> StrictSig
mkClosedStrictSig (Arity -> Demand -> [Demand]
forall a. Arity -> a -> [a]
replicate Arity
ar Demand
topDmd) Divergence
botDiv
data ArityType
= ATop [OneShotInfo]
| ABot Arity
deriving( ArityType -> ArityType -> Bool
(ArityType -> ArityType -> Bool)
-> (ArityType -> ArityType -> Bool) -> Eq ArityType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArityType -> ArityType -> Bool
$c/= :: ArityType -> ArityType -> Bool
== :: ArityType -> ArityType -> Bool
$c== :: ArityType -> ArityType -> Bool
Eq )
instance Outputable ArityType where
ppr :: ArityType -> SDoc
ppr (ATop [OneShotInfo]
os) = String -> SDoc
text String
"ATop" SDoc -> SDoc -> SDoc
<> SDoc -> SDoc
parens (Arity -> SDoc
forall a. Outputable a => a -> SDoc
ppr ([OneShotInfo] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [OneShotInfo]
os))
ppr (ABot Arity
n) = String -> SDoc
text String
"ABot" SDoc -> SDoc -> SDoc
<> SDoc -> SDoc
parens (Arity -> SDoc
forall a. Outputable a => a -> SDoc
ppr Arity
n)
arityTypeArity :: ArityType -> Arity
arityTypeArity :: ArityType -> Arity
arityTypeArity (ATop [OneShotInfo]
oss) = [OneShotInfo] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [OneShotInfo]
oss
arityTypeArity (ABot Arity
ar) = Arity
ar
expandableArityType :: ArityType -> Bool
expandableArityType :: ArityType -> Bool
expandableArityType (ATop [OneShotInfo]
oss) = Bool -> Bool
not ([OneShotInfo] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [OneShotInfo]
oss)
expandableArityType (ABot Arity
ar) = Arity
ar Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
/= Arity
0
isBotArityType :: ArityType -> Bool
isBotArityType :: ArityType -> Bool
isBotArityType (ABot {}) = Bool
True
isBotArityType (ATop {}) = Bool
False
arityTypeOneShots :: ArityType -> [OneShotInfo]
arityTypeOneShots :: ArityType -> [OneShotInfo]
arityTypeOneShots (ATop [OneShotInfo]
oss) = [OneShotInfo]
oss
arityTypeOneShots (ABot Arity
ar) = Arity -> OneShotInfo -> [OneShotInfo]
forall a. Arity -> a -> [a]
replicate Arity
ar OneShotInfo
OneShotLam
botArityType :: ArityType
botArityType :: ArityType
botArityType = Arity -> ArityType
ABot Arity
0
maxWithArity :: ArityType -> Arity -> ArityType
maxWithArity :: ArityType -> Arity -> ArityType
maxWithArity at :: ArityType
at@(ABot {}) Arity
_ = ArityType
at
maxWithArity at :: ArityType
at@(ATop [OneShotInfo]
oss) Arity
ar
| [OneShotInfo]
oss [OneShotInfo] -> Arity -> Bool
forall a. [a] -> Arity -> Bool
`lengthAtLeast` Arity
ar = ArityType
at
| Bool
otherwise = [OneShotInfo] -> ArityType
ATop (Arity -> [OneShotInfo] -> [OneShotInfo]
forall a. Arity -> [a] -> [a]
take Arity
ar ([OneShotInfo]
oss [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
forall a. [a] -> [a] -> [a]
++ OneShotInfo -> [OneShotInfo]
forall a. a -> [a]
repeat OneShotInfo
NoOneShotInfo))
vanillaArityType :: ArityType
vanillaArityType :: ArityType
vanillaArityType = [OneShotInfo] -> ArityType
ATop []
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
exprEtaExpandArity DynFlags
dflags CoreExpr
e
= ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e
where
env :: ArityEnv
env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_cheap_fn :: CheapFun
ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
isCheapApp
, ae_ped_bot :: Bool
ae_ped_bot = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
getBotArity :: ArityType -> Maybe Arity
getBotArity :: ArityType -> Maybe Arity
getBotArity (ABot Arity
n) = Arity -> Maybe Arity
forall a. a -> Maybe a
Just Arity
n
getBotArity ArityType
_ = Maybe Arity
forall a. Maybe a
Nothing
mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
cheap_app
| Bool -> Bool
not (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DictsCheap DynFlags
dflags)
= \CoreExpr
e Maybe Type
_ -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX CheapAppFun
cheap_app CoreExpr
e
| Bool
otherwise
= \CoreExpr
e Maybe Type
mb_ty -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX CheapAppFun
cheap_app CoreExpr
e
Bool -> Bool -> Bool
|| case Maybe Type
mb_ty of
Maybe Type
Nothing -> Bool
False
Just Type
ty -> Type -> Bool
isDictTy Type
ty
findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
findRhsArity :: DynFlags -> CoreBndr -> CoreExpr -> Arity -> ArityType
findRhsArity DynFlags
dflags CoreBndr
bndr CoreExpr
rhs Arity
old_arity
= ArityType -> ArityType
go (CheapAppFun -> ArityType
get_arity CheapAppFun
init_cheap_app)
where
init_cheap_app :: CheapAppFun
init_cheap_app :: CheapAppFun
init_cheap_app CoreBndr
fn Arity
n_val_args
| CoreBndr
fn CoreBndr -> CoreBndr -> Bool
forall a. Eq a => a -> a -> Bool
== CoreBndr
bndr = Bool
True
| Bool
otherwise = CheapAppFun
isCheapApp CoreBndr
fn Arity
n_val_args
go :: ArityType -> ArityType
go :: ArityType -> ArityType
go ArityType
cur_atype
| Arity
cur_arity Arity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
<= Arity
old_arity = ArityType
cur_atype
| ArityType
new_atype ArityType -> ArityType -> Bool
forall a. Eq a => a -> a -> Bool
== ArityType
cur_atype = ArityType
cur_atype
| Bool
otherwise =
#if defined(DEBUG)
pprTrace "Exciting arity"
(vcat [ ppr bndr <+> ppr cur_atype <+> ppr new_atype
, ppr rhs])
#endif
ArityType -> ArityType
go ArityType
new_atype
where
new_atype :: ArityType
new_atype = CheapAppFun -> ArityType
get_arity CheapAppFun
cheap_app
cur_arity :: Arity
cur_arity = ArityType -> Arity
arityTypeArity ArityType
cur_atype
cheap_app :: CheapAppFun
cheap_app :: CheapAppFun
cheap_app CoreBndr
fn Arity
n_val_args
| CoreBndr
fn CoreBndr -> CoreBndr -> Bool
forall a. Eq a => a -> a -> Bool
== CoreBndr
bndr = Arity
n_val_args Arity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
< Arity
cur_arity
| Bool
otherwise = CheapAppFun
isCheapApp CoreBndr
fn Arity
n_val_args
get_arity :: CheapAppFun -> ArityType
get_arity :: CheapAppFun -> ArityType
get_arity CheapAppFun
cheap_app = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs
where
env :: ArityEnv
env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_cheap_fn :: CheapFun
ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
cheap_app
, ae_ped_bot :: Bool
ae_ped_bot = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
arityLam :: Id -> ArityType -> ArityType
arityLam :: CoreBndr -> ArityType -> ArityType
arityLam CoreBndr
id (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop (CoreBndr -> OneShotInfo
idStateHackOneShotInfo CoreBndr
id OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: [OneShotInfo]
as)
arityLam CoreBndr
_ (ABot Arity
n) = Arity -> ArityType
ABot (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
+Arity
1)
floatIn :: Bool -> ArityType -> ArityType
floatIn :: Bool -> ArityType -> ArityType
floatIn Bool
_ (ABot Arity
n) = Arity -> ArityType
ABot Arity
n
floatIn Bool
True (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
floatIn Bool
False (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as)
arityApp :: ArityType -> Bool -> ArityType
arityApp :: ArityType -> Bool -> ArityType
arityApp (ABot Arity
0) Bool
_ = Arity -> ArityType
ABot Arity
0
arityApp (ABot Arity
n) Bool
_ = Arity -> ArityType
ABot (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
-Arity
1)
arityApp (ATop []) Bool
_ = [OneShotInfo] -> ArityType
ATop []
arityApp (ATop (OneShotInfo
_:[OneShotInfo]
as)) Bool
cheap = Bool -> ArityType -> ArityType
floatIn Bool
cheap ([OneShotInfo] -> ArityType
ATop [OneShotInfo]
as)
andArityType :: ArityType -> ArityType -> ArityType
andArityType :: ArityType -> ArityType -> ArityType
andArityType (ABot Arity
n1) (ABot Arity
n2) = Arity -> ArityType
ABot (Arity
n1 Arity -> Arity -> Arity
forall a. Ord a => a -> a -> a
`max` Arity
n2)
andArityType (ATop [OneShotInfo]
as) (ABot Arity
_) = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
andArityType (ABot Arity
_) (ATop [OneShotInfo]
bs) = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
bs
andArityType (ATop [OneShotInfo]
as) (ATop [OneShotInfo]
bs) = [OneShotInfo] -> ArityType
ATop ([OneShotInfo]
as [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
`combine` [OneShotInfo]
bs)
where
combine :: [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
combine (OneShotInfo
a:[OneShotInfo]
as) (OneShotInfo
b:[OneShotInfo]
bs) = (OneShotInfo
a OneShotInfo -> OneShotInfo -> OneShotInfo
`bestOneShot` OneShotInfo
b) OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
combine [OneShotInfo]
as [OneShotInfo]
bs
combine [] [OneShotInfo]
bs = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
bs
combine [OneShotInfo]
as [] = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as
type CheapFun = CoreExpr -> Maybe Type -> Bool
data ArityEnv
= AE { ArityEnv -> CheapFun
ae_cheap_fn :: CheapFun
, ArityEnv -> Bool
ae_ped_bot :: Bool
, ArityEnv -> IdSet
ae_joins :: IdSet
}
extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
extendJoinEnv :: ArityEnv -> [CoreBndr] -> ArityEnv
extendJoinEnv env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins }) [CoreBndr]
join_ids
= ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet
joins IdSet -> [CoreBndr] -> IdSet
`extendVarSetList` [CoreBndr]
join_ids }
arityType :: ArityEnv -> CoreExpr -> ArityType
arityType :: ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env (Cast CoreExpr
e Coercion
co)
= case ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e of
ATop [OneShotInfo]
os -> [OneShotInfo] -> ArityType
ATop (Arity -> [OneShotInfo] -> [OneShotInfo]
forall a. Arity -> [a] -> [a]
take Arity
co_arity [OneShotInfo]
os)
ABot Arity
n | Arity
co_arity Arity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
< Arity
n -> [OneShotInfo] -> ArityType
ATop (Arity -> OneShotInfo -> [OneShotInfo]
forall a. Arity -> a -> [a]
replicate Arity
co_arity OneShotInfo
noOneShotInfo)
| Bool
otherwise -> Arity -> ArityType
ABot Arity
n
where
co_arity :: Arity
co_arity = [OneShotInfo] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length (Type -> [OneShotInfo]
typeArity (Coercion -> Type
coercionRKind Coercion
co))
arityType ArityEnv
env (Var CoreBndr
v)
| CoreBndr
v CoreBndr -> IdSet -> Bool
`elemVarSet` ArityEnv -> IdSet
ae_joins ArityEnv
env
= ArityType
botArityType
| Bool
otherwise
= CoreBndr -> ArityType
idArityType CoreBndr
v
arityType ArityEnv
env (Lam CoreBndr
x CoreExpr
e)
| CoreBndr -> Bool
isId CoreBndr
x = CoreBndr -> ArityType -> ArityType
arityLam CoreBndr
x (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e)
| Bool
otherwise = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e
arityType ArityEnv
env (App CoreExpr
fun (Type Type
_))
= ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun
arityType ArityEnv
env (App CoreExpr
fun CoreExpr
arg )
= ArityType -> Bool -> ArityType
arityApp (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun) (ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
arg Maybe Type
forall a. Maybe a
Nothing)
arityType ArityEnv
env (Case CoreExpr
scrut CoreBndr
_ Type
_ [Alt CoreBndr]
alts)
| CoreExpr -> Bool
exprIsDeadEnd CoreExpr
scrut Bool -> Bool -> Bool
|| [Alt CoreBndr] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Alt CoreBndr]
alts
= ArityType
botArityType
| Bool
otherwise
= case ArityType
alts_type of
ABot Arity
n | Arity
nArity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
>Arity
0 -> [OneShotInfo] -> ArityType
ATop []
| Bool
otherwise -> ArityType
botArityType
ATop [OneShotInfo]
as | Bool -> Bool
not (ArityEnv -> Bool
ae_ped_bot ArityEnv
env)
, ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
scrut Maybe Type
forall a. Maybe a
Nothing -> [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
| CoreExpr -> Bool
exprOkForSpeculation CoreExpr
scrut -> [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
| Bool
otherwise -> [OneShotInfo] -> ArityType
ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as)
where
alts_type :: ArityType
alts_type = (ArityType -> ArityType -> ArityType) -> [ArityType] -> ArityType
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 ArityType -> ArityType -> ArityType
andArityType [ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs | (AltCon
_,[CoreBndr]
_,CoreExpr
rhs) <- [Alt CoreBndr]
alts]
arityType ArityEnv
env (Let (NonRec CoreBndr
j CoreExpr
rhs) CoreExpr
body)
| Just Arity
join_arity <- CoreBndr -> Maybe Arity
isJoinId_maybe CoreBndr
j
, ([CoreBndr]
_, CoreExpr
rhs_body) <- Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
forall b. Arity -> Expr b -> ([b], Expr b)
collectNBinders Arity
join_arity CoreExpr
rhs
=
ArityType -> ArityType -> ArityType
andArityType (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs_body)
(ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
body)
where
env' :: ArityEnv
env' = ArityEnv -> [CoreBndr] -> ArityEnv
extendJoinEnv ArityEnv
env [CoreBndr
j]
arityType ArityEnv
env (Let (Rec [(CoreBndr, CoreExpr)]
pairs) CoreExpr
body)
| ((CoreBndr
j,CoreExpr
_):[(CoreBndr, CoreExpr)]
_) <- [(CoreBndr, CoreExpr)]
pairs
, CoreBndr -> Bool
isJoinId CoreBndr
j
=
((CoreBndr, CoreExpr) -> ArityType -> ArityType)
-> ArityType -> [(CoreBndr, CoreExpr)] -> ArityType
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (ArityType -> ArityType -> ArityType
andArityType (ArityType -> ArityType -> ArityType)
-> ((CoreBndr, CoreExpr) -> ArityType)
-> (CoreBndr, CoreExpr)
-> ArityType
-> ArityType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CoreBndr, CoreExpr) -> ArityType
do_one) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
body) [(CoreBndr, CoreExpr)]
pairs
where
env' :: ArityEnv
env' = ArityEnv -> [CoreBndr] -> ArityEnv
extendJoinEnv ArityEnv
env (((CoreBndr, CoreExpr) -> CoreBndr)
-> [(CoreBndr, CoreExpr)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, CoreExpr) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, CoreExpr)]
pairs)
do_one :: (CoreBndr, CoreExpr) -> ArityType
do_one (CoreBndr
j,CoreExpr
rhs)
| Just Arity
arity <- CoreBndr -> Maybe Arity
isJoinId_maybe CoreBndr
j
= ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' (CoreExpr -> ArityType) -> CoreExpr -> ArityType
forall a b. (a -> b) -> a -> b
$ ([CoreBndr], CoreExpr) -> CoreExpr
forall a b. (a, b) -> b
snd (([CoreBndr], CoreExpr) -> CoreExpr)
-> ([CoreBndr], CoreExpr) -> CoreExpr
forall a b. (a -> b) -> a -> b
$ Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
forall b. Arity -> Expr b -> ([b], Expr b)
collectNBinders Arity
arity CoreExpr
rhs
| Bool
otherwise
= String -> SDoc -> ArityType
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"arityType:joinrec" ([(CoreBndr, CoreExpr)] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [(CoreBndr, CoreExpr)]
pairs)
arityType ArityEnv
env (Let Bind CoreBndr
b CoreExpr
e)
= Bool -> ArityType -> ArityType
floatIn (Bind CoreBndr -> Bool
cheap_bind Bind CoreBndr
b) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e)
where
cheap_bind :: Bind CoreBndr -> Bool
cheap_bind (NonRec CoreBndr
b CoreExpr
e) = (CoreBndr, CoreExpr) -> Bool
is_cheap (CoreBndr
b,CoreExpr
e)
cheap_bind (Rec [(CoreBndr, CoreExpr)]
prs) = ((CoreBndr, CoreExpr) -> Bool) -> [(CoreBndr, CoreExpr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (CoreBndr, CoreExpr) -> Bool
is_cheap [(CoreBndr, CoreExpr)]
prs
is_cheap :: (CoreBndr, CoreExpr) -> Bool
is_cheap (CoreBndr
b,CoreExpr
e) = ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
e (Type -> Maybe Type
forall a. a -> Maybe a
Just (CoreBndr -> Type
idType CoreBndr
b))
arityType ArityEnv
env (Tick Tickish CoreBndr
t CoreExpr
e)
| Bool -> Bool
not (Tickish CoreBndr -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish CoreBndr
t) = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e
arityType ArityEnv
_ CoreExpr
_ = ArityType
vanillaArityType
idArityType :: Id -> ArityType
idArityType :: CoreBndr -> ArityType
idArityType CoreBndr
v
| StrictSig
strict_sig <- CoreBndr -> StrictSig
idStrictness CoreBndr
v
, Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ StrictSig -> Bool
isTopSig StrictSig
strict_sig
, ([Demand]
ds, Divergence
res) <- StrictSig -> ([Demand], Divergence)
splitStrictSig StrictSig
strict_sig
, let arity :: Arity
arity = [Demand] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [Demand]
ds
= if Divergence -> Bool
isDeadEndDiv Divergence
res then Arity -> ArityType
ABot Arity
arity
else [OneShotInfo] -> ArityType
ATop (Arity -> [OneShotInfo] -> [OneShotInfo]
forall a. Arity -> [a] -> [a]
take Arity
arity [OneShotInfo]
one_shots)
| Bool
otherwise
= [OneShotInfo] -> ArityType
ATop (Arity -> [OneShotInfo] -> [OneShotInfo]
forall a. Arity -> [a] -> [a]
take (CoreBndr -> Arity
idArity CoreBndr
v) [OneShotInfo]
one_shots)
where
one_shots :: [OneShotInfo]
one_shots :: [OneShotInfo]
one_shots = Type -> [OneShotInfo]
typeArity (CoreBndr -> Type
idType CoreBndr
v)
etaExpand :: Arity -> CoreExpr -> CoreExpr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpand :: Arity -> CoreExpr -> CoreExpr
etaExpand Arity
n CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (Arity -> OneShotInfo -> [OneShotInfo]
forall a. Arity -> a -> [a]
replicate Arity
n OneShotInfo
NoOneShotInfo) CoreExpr
orig_expr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpandAT ArityType
at CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (ArityType -> [OneShotInfo]
arityTypeOneShots ArityType
at) CoreExpr
orig_expr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots CoreExpr
orig_expr
= [OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
one_shots CoreExpr
orig_expr
where
go :: [OneShotInfo] -> CoreExpr -> CoreExpr
go [] CoreExpr
expr = CoreExpr
expr
go oss :: [OneShotInfo]
oss@(OneShotInfo
_:[OneShotInfo]
oss1) (Lam CoreBndr
v CoreExpr
body) | CoreBndr -> Bool
isTyVar CoreBndr
v = CoreBndr -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
v ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss CoreExpr
body)
| Bool
otherwise = CoreBndr -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
v ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss1 CoreExpr
body)
go [OneShotInfo]
oss (Cast CoreExpr
expr Coercion
co) = CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss CoreExpr
expr) Coercion
co
go [OneShotInfo]
oss CoreExpr
expr
=
CoreExpr -> CoreExpr
retick (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$ [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
etas (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst' CoreExpr
sexpr [EtaInfo]
etas)
where
in_scope :: InScopeSet
in_scope = IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
expr)
(InScopeSet
in_scope', [EtaInfo]
etas) = [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo])
mkEtaWW [OneShotInfo]
oss (CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
orig_expr) InScopeSet
in_scope (CoreExpr -> Type
exprType CoreExpr
expr)
subst' :: Subst
subst' = InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope'
(CoreExpr
expr', [CoreExpr]
args) = CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
([Tickish CoreBndr]
ticks, CoreExpr
expr'') = (Tickish CoreBndr -> Bool)
-> CoreExpr -> ([Tickish CoreBndr], CoreExpr)
forall b.
(Tickish CoreBndr -> Bool)
-> Expr b -> ([Tickish CoreBndr], Expr b)
stripTicksTop Tickish CoreBndr -> Bool
forall id. Tickish id -> Bool
tickishFloatable CoreExpr
expr'
sexpr :: CoreExpr
sexpr = (CoreExpr -> CoreExpr -> CoreExpr)
-> CoreExpr -> [CoreExpr] -> CoreExpr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
expr'' [CoreExpr]
args
retick :: CoreExpr -> CoreExpr
retick CoreExpr
expr = (Tickish CoreBndr -> CoreExpr -> CoreExpr)
-> CoreExpr -> [Tickish CoreBndr] -> CoreExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tickish CoreBndr -> CoreExpr -> CoreExpr
mkTick CoreExpr
expr [Tickish CoreBndr]
ticks
data EtaInfo = EtaVar Var
| EtaCo Coercion
instance Outputable EtaInfo where
ppr :: EtaInfo -> SDoc
ppr (EtaVar CoreBndr
v) = String -> SDoc
text String
"EtaVar" SDoc -> SDoc -> SDoc
<+> CoreBndr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreBndr
v
ppr (EtaCo Coercion
co) = String -> SDoc
text String
"EtaCo" SDoc -> SDoc -> SDoc
<+> Coercion -> SDoc
forall a. Outputable a => a -> SDoc
ppr Coercion
co
pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion Coercion
co1 (EtaCo Coercion
co2 : [EtaInfo]
eis)
| Coercion -> Bool
isReflCo Coercion
co = [EtaInfo]
eis
| Bool
otherwise = Coercion -> EtaInfo
EtaCo Coercion
co EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis
where
co :: Coercion
co = Coercion
co1 Coercion -> Coercion -> Coercion
`mkTransCo` Coercion
co2
pushCoercion Coercion
co [EtaInfo]
eis = Coercion -> EtaInfo
EtaCo Coercion
co EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [] CoreExpr
expr = CoreExpr
expr
etaInfoAbs (EtaVar CoreBndr
v : [EtaInfo]
eis) CoreExpr
expr = CoreBndr -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
v ([EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
eis CoreExpr
expr)
etaInfoAbs (EtaCo Coercion
co : [EtaInfo]
eis) CoreExpr
expr = CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast ([EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
eis CoreExpr
expr) (Coercion -> Coercion
mkSymCo Coercion
co)
etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst (Lam CoreBndr
v1 CoreExpr
e) (EtaVar CoreBndr
v2 : [EtaInfo]
eis)
= Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp (Subst -> CoreBndr -> CoreBndr -> Subst
GHC.Core.Subst.extendSubstWithVar Subst
subst CoreBndr
v1 CoreBndr
v2) CoreExpr
e [EtaInfo]
eis
etaInfoApp Subst
subst (Cast CoreExpr
e Coercion
co1) [EtaInfo]
eis
= Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst CoreExpr
e (Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion Coercion
co' [EtaInfo]
eis)
where
co' :: Coercion
co' = HasCallStack => Subst -> Coercion -> Coercion
Subst -> Coercion -> Coercion
GHC.Core.Subst.substCo Subst
subst Coercion
co1
etaInfoApp Subst
subst (Case CoreExpr
e CoreBndr
b Type
ty [Alt CoreBndr]
alts) [EtaInfo]
eis
= CoreExpr -> CoreBndr -> Type -> [Alt CoreBndr] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
e) CoreBndr
b1 Type
ty' [Alt CoreBndr]
alts'
where
(Subst
subst1, CoreBndr
b1) = Subst -> CoreBndr -> (Subst, CoreBndr)
substBndr Subst
subst CoreBndr
b
alts' :: [Alt CoreBndr]
alts' = (Alt CoreBndr -> Alt CoreBndr) -> [Alt CoreBndr] -> [Alt CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map Alt CoreBndr -> Alt CoreBndr
forall a. (a, [CoreBndr], CoreExpr) -> (a, [CoreBndr], CoreExpr)
subst_alt [Alt CoreBndr]
alts
ty' :: Type
ty' = Type -> [EtaInfo] -> Type
etaInfoAppTy (Subst -> Type -> Type
GHC.Core.Subst.substTy Subst
subst Type
ty) [EtaInfo]
eis
subst_alt :: (a, [CoreBndr], CoreExpr) -> (a, [CoreBndr], CoreExpr)
subst_alt (a
con, [CoreBndr]
bs, CoreExpr
rhs) = (a
con, [CoreBndr]
bs', Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst2 CoreExpr
rhs [EtaInfo]
eis)
where
(Subst
subst2,[CoreBndr]
bs') = Subst -> [CoreBndr] -> (Subst, [CoreBndr])
substBndrs Subst
subst1 [CoreBndr]
bs
etaInfoApp Subst
subst (Let Bind CoreBndr
b CoreExpr
e) [EtaInfo]
eis
| Bool -> Bool
not (Bind CoreBndr -> Bool
isJoinBind Bind CoreBndr
b)
= Bind CoreBndr -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let Bind CoreBndr
b' (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst' CoreExpr
e [EtaInfo]
eis)
where
(Subst
subst', Bind CoreBndr
b') = HasDebugCallStack =>
Subst -> Bind CoreBndr -> (Subst, Bind CoreBndr)
Subst -> Bind CoreBndr -> (Subst, Bind CoreBndr)
substBindSC Subst
subst Bind CoreBndr
b
etaInfoApp Subst
subst (Tick Tickish CoreBndr
t CoreExpr
e) [EtaInfo]
eis
= Tickish CoreBndr -> CoreExpr -> CoreExpr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick (Subst -> Tickish CoreBndr -> Tickish CoreBndr
substTickish Subst
subst Tickish CoreBndr
t) (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst CoreExpr
e [EtaInfo]
eis)
etaInfoApp Subst
subst CoreExpr
expr [EtaInfo]
_
| (Var CoreBndr
fun, [CoreExpr]
_) <- CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
, Var CoreBndr
fun' <- HasDebugCallStack => Subst -> CoreBndr -> CoreExpr
Subst -> CoreBndr -> CoreExpr
lookupIdSubst Subst
subst CoreBndr
fun
, CoreBndr -> Bool
isJoinId CoreBndr
fun'
= Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
expr
etaInfoApp Subst
subst CoreExpr
e [EtaInfo]
eis
= CoreExpr -> [EtaInfo] -> CoreExpr
forall b. Expr b -> [EtaInfo] -> Expr b
go (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
e) [EtaInfo]
eis
where
go :: Expr b -> [EtaInfo] -> Expr b
go Expr b
e [] = Expr b
e
go Expr b
e (EtaVar CoreBndr
v : [EtaInfo]
eis) = Expr b -> [EtaInfo] -> Expr b
go (Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
App Expr b
e (CoreBndr -> Expr b
forall b. CoreBndr -> Expr b
varToCoreExpr CoreBndr
v)) [EtaInfo]
eis
go Expr b
e (EtaCo Coercion
co : [EtaInfo]
eis) = Expr b -> [EtaInfo] -> Expr b
go (Expr b -> Coercion -> Expr b
forall b. Expr b -> Coercion -> Expr b
Cast Expr b
e Coercion
co) [EtaInfo]
eis
etaInfoAppTy :: Type -> [EtaInfo] -> Type
etaInfoAppTy :: Type -> [EtaInfo] -> Type
etaInfoAppTy Type
ty [] = Type
ty
etaInfoAppTy Type
ty (EtaVar CoreBndr
v : [EtaInfo]
eis) = Type -> [EtaInfo] -> Type
etaInfoAppTy (Type -> CoreExpr -> Type
applyTypeToArg Type
ty (CoreBndr -> CoreExpr
forall b. CoreBndr -> Expr b
varToCoreExpr CoreBndr
v)) [EtaInfo]
eis
etaInfoAppTy Type
_ (EtaCo Coercion
co : [EtaInfo]
eis) = Type -> [EtaInfo] -> Type
etaInfoAppTy (Coercion -> Type
coercionRKind Coercion
co) [EtaInfo]
eis
mkEtaWW
:: [OneShotInfo]
-> SDoc
-> InScopeSet
-> Type
-> (InScopeSet, [EtaInfo])
mkEtaWW :: [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo])
mkEtaWW [OneShotInfo]
orig_oss SDoc
ppr_orig_expr InScopeSet
in_scope Type
orig_ty
= Arity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go Arity
0 [OneShotInfo]
orig_oss TCvSubst
empty_subst Type
orig_ty []
where
empty_subst :: TCvSubst
empty_subst = InScopeSet -> TCvSubst
mkEmptyTCvSubst InScopeSet
in_scope
go :: Int
-> [OneShotInfo]
-> TCvSubst -> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go :: Arity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go Arity
_ [] TCvSubst
subst Type
_ [EtaInfo]
eis
= (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [EtaInfo] -> [EtaInfo]
forall a. [a] -> [a]
reverse [EtaInfo]
eis)
go Arity
n oss :: [OneShotInfo]
oss@(OneShotInfo
one_shot:[OneShotInfo]
oss1) TCvSubst
subst Type
ty [EtaInfo]
eis
| Just (CoreBndr
tcv,Type
ty') <- Type -> Maybe (CoreBndr, Type)
splitForAllTy_maybe Type
ty
, (TCvSubst
subst', CoreBndr
tcv') <- HasCallStack => TCvSubst -> CoreBndr -> (TCvSubst, CoreBndr)
TCvSubst -> CoreBndr -> (TCvSubst, CoreBndr)
Type.substVarBndr TCvSubst
subst CoreBndr
tcv
, let oss' :: [OneShotInfo]
oss' | CoreBndr -> Bool
isTyVar CoreBndr
tcv = [OneShotInfo]
oss
| Bool
otherwise = [OneShotInfo]
oss1
= Arity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go Arity
n [OneShotInfo]
oss' TCvSubst
subst' Type
ty' (CoreBndr -> EtaInfo
EtaVar CoreBndr
tcv' EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis)
| Just (Type
mult, Type
arg_ty, Type
res_ty) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
, Bool -> Bool
not (Type -> Bool
isTypeLevPoly Type
arg_ty)
, (TCvSubst
subst', CoreBndr
eta_id) <- Arity -> TCvSubst -> Scaled Type -> (TCvSubst, CoreBndr)
freshEtaId Arity
n TCvSubst
subst (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
, let eta_id' :: CoreBndr
eta_id' = CoreBndr
eta_id CoreBndr -> OneShotInfo -> CoreBndr
`setIdOneShotInfo` OneShotInfo
one_shot
= Arity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
+Arity
1) [OneShotInfo]
oss1 TCvSubst
subst' Type
res_ty (CoreBndr -> EtaInfo
EtaVar CoreBndr
eta_id' EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis)
| Just (Coercion
co, Type
ty') <- Type -> Maybe (Coercion, Type)
topNormaliseNewType_maybe Type
ty
, let co' :: Coercion
co' = HasCallStack => TCvSubst -> Coercion -> Coercion
TCvSubst -> Coercion -> Coercion
Coercion.substCo TCvSubst
subst Coercion
co
= Arity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go Arity
n [OneShotInfo]
oss TCvSubst
subst Type
ty' (Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion Coercion
co' [EtaInfo]
eis)
| Bool
otherwise
= WARN( True, (ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr )
(TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [EtaInfo] -> [EtaInfo]
forall a. [a] -> [a]
reverse [EtaInfo]
eis)
subst_expr :: Subst -> CoreExpr -> CoreExpr
subst_expr :: Subst -> CoreExpr -> CoreExpr
subst_expr = HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Subst -> CoreExpr -> CoreExpr
substExpr
etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaExpandToJoinPoint :: Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaExpandToJoinPoint Arity
join_arity CoreExpr
expr
= Arity -> [CoreBndr] -> CoreExpr -> ([CoreBndr], CoreExpr)
go Arity
join_arity [] CoreExpr
expr
where
go :: Arity -> [CoreBndr] -> CoreExpr -> ([CoreBndr], CoreExpr)
go Arity
0 [CoreBndr]
rev_bs CoreExpr
e = ([CoreBndr] -> [CoreBndr]
forall a. [a] -> [a]
reverse [CoreBndr]
rev_bs, CoreExpr
e)
go Arity
n [CoreBndr]
rev_bs (Lam CoreBndr
b CoreExpr
e) = Arity -> [CoreBndr] -> CoreExpr -> ([CoreBndr], CoreExpr)
go (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
-Arity
1) (CoreBndr
b CoreBndr -> [CoreBndr] -> [CoreBndr]
forall a. a -> [a] -> [a]
: [CoreBndr]
rev_bs) CoreExpr
e
go Arity
n [CoreBndr]
rev_bs CoreExpr
e = case Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint Arity
n CoreExpr
e of
([CoreBndr]
bs, CoreExpr
e') -> ([CoreBndr] -> [CoreBndr]
forall a. [a] -> [a]
reverse [CoreBndr]
rev_bs [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
bs, CoreExpr
e')
etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule :: Arity -> CoreRule -> CoreRule
etaExpandToJoinPointRule Arity
_ rule :: CoreRule
rule@(BuiltinRule {})
= WARN(True, (sep [text "Can't eta-expand built-in rule:", ppr rule]))
CoreRule
rule
etaExpandToJoinPointRule Arity
join_arity rule :: CoreRule
rule@(Rule { ru_bndrs :: CoreRule -> [CoreBndr]
ru_bndrs = [CoreBndr]
bndrs, ru_rhs :: CoreRule -> CoreExpr
ru_rhs = CoreExpr
rhs
, ru_args :: CoreRule -> [CoreExpr]
ru_args = [CoreExpr]
args })
| Arity
need_args Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
0
= CoreRule
rule
| Arity
need_args Arity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
< Arity
0
= String -> SDoc -> CoreRule
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaExpandToJoinPointRule" (Arity -> SDoc
forall a. Outputable a => a -> SDoc
ppr Arity
join_arity SDoc -> SDoc -> SDoc
$$ CoreRule -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreRule
rule)
| Bool
otherwise
= CoreRule
rule { ru_bndrs :: [CoreBndr]
ru_bndrs = [CoreBndr]
bndrs [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
new_bndrs, ru_args :: [CoreExpr]
ru_args = [CoreExpr]
args [CoreExpr] -> [CoreExpr] -> [CoreExpr]
forall a. [a] -> [a] -> [a]
++ [CoreExpr]
forall b. [Expr b]
new_args
, ru_rhs :: CoreExpr
ru_rhs = CoreExpr
new_rhs }
where
need_args :: Arity
need_args = Arity
join_arity Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
- [CoreExpr] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [CoreExpr]
args
([CoreBndr]
new_bndrs, CoreExpr
new_rhs) = Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint Arity
need_args CoreExpr
rhs
new_args :: [Expr b]
new_args = [CoreBndr] -> [Expr b]
forall b. [CoreBndr] -> [Expr b]
varsToCoreExprs [CoreBndr]
new_bndrs
etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint :: Arity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint Arity
need_args CoreExpr
body
= Arity
-> Type
-> TCvSubst
-> [CoreBndr]
-> CoreExpr
-> ([CoreBndr], CoreExpr)
forall b.
Arity
-> Type -> TCvSubst -> [CoreBndr] -> Expr b -> ([CoreBndr], Expr b)
go Arity
need_args (CoreExpr -> Type
exprType CoreExpr
body) (CoreExpr -> TCvSubst
init_subst CoreExpr
body) [] CoreExpr
body
where
go :: Arity
-> Type -> TCvSubst -> [CoreBndr] -> Expr b -> ([CoreBndr], Expr b)
go Arity
0 Type
_ TCvSubst
_ [CoreBndr]
rev_bs Expr b
e
= ([CoreBndr] -> [CoreBndr]
forall a. [a] -> [a]
reverse [CoreBndr]
rev_bs, Expr b
e)
go Arity
n Type
ty TCvSubst
subst [CoreBndr]
rev_bs Expr b
e
| Just (CoreBndr
tv, Type
res_ty) <- Type -> Maybe (CoreBndr, Type)
splitForAllTy_maybe Type
ty
, let (TCvSubst
subst', CoreBndr
tv') = HasCallStack => TCvSubst -> CoreBndr -> (TCvSubst, CoreBndr)
TCvSubst -> CoreBndr -> (TCvSubst, CoreBndr)
Type.substVarBndr TCvSubst
subst CoreBndr
tv
= Arity
-> Type -> TCvSubst -> [CoreBndr] -> Expr b -> ([CoreBndr], Expr b)
go (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
-Arity
1) Type
res_ty TCvSubst
subst' (CoreBndr
tv' CoreBndr -> [CoreBndr] -> [CoreBndr]
forall a. a -> [a] -> [a]
: [CoreBndr]
rev_bs) (Expr b
e Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
`App` CoreBndr -> Expr b
forall b. CoreBndr -> Expr b
varToCoreExpr CoreBndr
tv')
| Just (Type
mult, Type
arg_ty, Type
res_ty) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
, let (TCvSubst
subst', CoreBndr
b) = Arity -> TCvSubst -> Scaled Type -> (TCvSubst, CoreBndr)
freshEtaId Arity
n TCvSubst
subst (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
= Arity
-> Type -> TCvSubst -> [CoreBndr] -> Expr b -> ([CoreBndr], Expr b)
go (Arity
nArity -> Arity -> Arity
forall a. Num a => a -> a -> a
-Arity
1) Type
res_ty TCvSubst
subst' (CoreBndr
b CoreBndr -> [CoreBndr] -> [CoreBndr]
forall a. a -> [a] -> [a]
: [CoreBndr]
rev_bs) (Expr b
e Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
`App` CoreBndr -> Expr b
forall b. CoreBndr -> Expr b
Var CoreBndr
b)
| Bool
otherwise
= String -> SDoc -> ([CoreBndr], Expr b)
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaBodyForJoinPoint" (SDoc -> ([CoreBndr], Expr b)) -> SDoc -> ([CoreBndr], Expr b)
forall a b. (a -> b) -> a -> b
$ Arity -> SDoc
int Arity
need_args SDoc -> SDoc -> SDoc
$$
CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
body SDoc -> SDoc -> SDoc
$$ Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr (CoreExpr -> Type
exprType CoreExpr
body)
init_subst :: CoreExpr -> TCvSubst
init_subst CoreExpr
e = InScopeSet -> TCvSubst
mkEmptyTCvSubst (IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
e))
freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
freshEtaId :: Arity -> TCvSubst -> Scaled Type -> (TCvSubst, CoreBndr)
freshEtaId Arity
n TCvSubst
subst Scaled Type
ty
= (TCvSubst
subst', CoreBndr
eta_id')
where
Scaled Type
mult' Type
ty' = HasCallStack => TCvSubst -> Scaled Type -> Scaled Type
TCvSubst -> Scaled Type -> Scaled Type
Type.substScaledTyUnchecked TCvSubst
subst Scaled Type
ty
eta_id' :: CoreBndr
eta_id' = InScopeSet -> CoreBndr -> CoreBndr
uniqAway (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst) (CoreBndr -> CoreBndr) -> CoreBndr -> CoreBndr
forall a b. (a -> b) -> a -> b
$
FastString -> Unique -> Type -> Type -> CoreBndr
mkSysLocalOrCoVar (String -> FastString
fsLit String
"eta") (Arity -> Unique
mkBuiltinUnique Arity
n) Type
mult' Type
ty'
subst' :: TCvSubst
subst' = TCvSubst -> CoreBndr -> TCvSubst
extendTCvInScope TCvSubst
subst CoreBndr
eta_id'