{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
module GHC.Core.Opt.Arity
( manifestArity, joinRhsArity, exprArity, typeArity
, exprEtaExpandArity, findRhsArity
, etaExpand, etaExpandAT
, exprBotStrictness_maybe
, ArityType(..), mkBotArityType, mkTopArityType, expandableArityType
, arityTypeArity, maxWithArity, idArityType
, etaExpandToJoinPoint, etaExpandToJoinPointRule
, pushCoArg, pushCoArgs, pushCoValArg, pushCoTyArg
, pushCoercionIntoLambda, pushCoDataCon, collectBindersPushingCo
)
where
#include "HsVersions.h"
import GHC.Prelude
import GHC.Driver.Ppr
import GHC.Core
import GHC.Core.FVs
import GHC.Core.Utils
import GHC.Types.Demand
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Types.Id
import GHC.Core.Subst as Core
import GHC.Core.Type as Type
import GHC.Core.Coercion as Type
import GHC.Core.DataCon
import GHC.Core.TyCon ( tyConArity )
import GHC.Core.TyCon.RecWalk ( initRecTc, checkRecTc )
import GHC.Core.Predicate ( isDictTy )
import GHC.Core.Multiplicity
import GHC.Types.Var.Set
import GHC.Types.Basic
import GHC.Types.Tickish
import GHC.Builtin.Uniques
import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
import GHC.Utils.Outputable
import GHC.Utils.Panic
import GHC.Data.FastString
import GHC.Data.Pair
import GHC.Utils.Misc
manifestArity :: CoreExpr -> Arity
manifestArity :: CoreExpr -> Int
manifestArity (Lam TyVar
v CoreExpr
e) | TyVar -> Bool
isId TyVar
v = Int
1 forall a. Num a => a -> a -> a
+ CoreExpr -> Int
manifestArity CoreExpr
e
| Bool
otherwise = CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity (Tick CoreTickish
t CoreExpr
e) | Bool -> Bool
not (forall (pass :: TickishPass). GenTickish pass -> Bool
tickishIsCode CoreTickish
t) = CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity (Cast CoreExpr
e Coercion
_) = CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity CoreExpr
_ = Int
0
joinRhsArity :: CoreExpr -> JoinArity
joinRhsArity :: CoreExpr -> Int
joinRhsArity (Lam TyVar
_ CoreExpr
e) = Int
1 forall a. Num a => a -> a -> a
+ CoreExpr -> Int
joinRhsArity CoreExpr
e
joinRhsArity CoreExpr
_ = Int
0
exprArity :: CoreExpr -> Arity
exprArity :: CoreExpr -> Int
exprArity CoreExpr
e = CoreExpr -> Int
go CoreExpr
e
where
go :: CoreExpr -> Int
go (Var TyVar
v) = TyVar -> Int
idArity TyVar
v
go (Lam TyVar
x CoreExpr
e) | TyVar -> Bool
isId TyVar
x = CoreExpr -> Int
go CoreExpr
e forall a. Num a => a -> a -> a
+ Int
1
| Bool
otherwise = CoreExpr -> Int
go CoreExpr
e
go (Tick CoreTickish
t CoreExpr
e) | Bool -> Bool
not (forall (pass :: TickishPass). GenTickish pass -> Bool
tickishIsCode CoreTickish
t) = CoreExpr -> Int
go CoreExpr
e
go (Cast CoreExpr
e Coercion
co) = Int -> Type -> Int
trim_arity (CoreExpr -> Int
go CoreExpr
e) (Coercion -> Type
coercionRKind Coercion
co)
go (App CoreExpr
e (Type Type
_)) = CoreExpr -> Int
go CoreExpr
e
go (App CoreExpr
f CoreExpr
a) | CoreExpr -> Bool
exprIsTrivial CoreExpr
a = (CoreExpr -> Int
go CoreExpr
f forall a. Num a => a -> a -> a
- Int
1) forall a. Ord a => a -> a -> a
`max` Int
0
go CoreExpr
_ = Int
0
trim_arity :: Arity -> Type -> Arity
trim_arity :: Int -> Type -> Int
trim_arity Int
arity Type
ty = Int
arity forall a. Ord a => a -> a -> a
`min` forall (t :: * -> *) a. Foldable t => t a -> Int
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 (TyVar
_, Type
ty') <- Type -> Maybe (TyVar, Type)
splitForAllTyCoVar_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 forall a. a -> [a] -> [a]
: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
res
| Just (TyCon
tc,[Type]
tys) <- HasDebugCallStack => 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 (Int, StrictSig)
exprBotStrictness_maybe CoreExpr
e
= case ArityType -> Maybe Int
getBotArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
botStrictnessArityEnv CoreExpr
e) of
Maybe Int
Nothing -> forall a. Maybe a
Nothing
Just Int
ar -> forall a. a -> Maybe a
Just (Int
ar, Int -> StrictSig
sig Int
ar)
where
sig :: Int -> StrictSig
sig Int
ar = [Demand] -> Divergence -> StrictSig
mkClosedStrictSig (forall a. Int -> a -> [a]
replicate Int
ar Demand
topDmd) Divergence
botDiv
data ArityType
= AT ![OneShotInfo] !Divergence
deriving ArityType -> ArityType -> Bool
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 (AT [OneShotInfo]
oss Divergence
div)
| forall (t :: * -> *) a. Foldable t => t a -> Bool
null [OneShotInfo]
oss = Divergence -> SDoc
pp_div Divergence
div
| Bool
otherwise = Char -> SDoc
char Char
'\\' SDoc -> SDoc -> SDoc
<> [SDoc] -> SDoc
hcat (forall a b. (a -> b) -> [a] -> [b]
map OneShotInfo -> SDoc
pp_os [OneShotInfo]
oss) SDoc -> SDoc -> SDoc
<> SDoc
dot SDoc -> SDoc -> SDoc
<> Divergence -> SDoc
pp_div Divergence
div
where
pp_div :: Divergence -> SDoc
pp_div Divergence
Diverges = Char -> SDoc
char Char
'⊥'
pp_div Divergence
ExnOrDiv = Char -> SDoc
char Char
'x'
pp_div Divergence
Dunno = Char -> SDoc
char Char
'T'
pp_os :: OneShotInfo -> SDoc
pp_os OneShotInfo
OneShotLam = Char -> SDoc
char Char
'1'
pp_os OneShotInfo
NoOneShotInfo = Char -> SDoc
char Char
'?'
mkBotArityType :: [OneShotInfo] -> ArityType
mkBotArityType :: [OneShotInfo] -> ArityType
mkBotArityType [OneShotInfo]
oss = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
botDiv
botArityType :: ArityType
botArityType :: ArityType
botArityType = [OneShotInfo] -> ArityType
mkBotArityType []
mkTopArityType :: [OneShotInfo] -> ArityType
mkTopArityType :: [OneShotInfo] -> ArityType
mkTopArityType [OneShotInfo]
oss = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
topDiv
topArityType :: ArityType
topArityType :: ArityType
topArityType = [OneShotInfo] -> ArityType
mkTopArityType []
arityTypeArity :: ArityType -> Arity
arityTypeArity :: ArityType -> Int
arityTypeArity (AT [OneShotInfo]
oss Divergence
_) = forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss
expandableArityType :: ArityType -> Bool
expandableArityType :: ArityType -> Bool
expandableArityType ArityType
at = ArityType -> Int
arityTypeArity ArityType
at forall a. Eq a => a -> a -> Bool
/= Int
0
isDeadEndArityType :: ArityType -> Bool
isDeadEndArityType :: ArityType -> Bool
isDeadEndArityType (AT [OneShotInfo]
_ Divergence
div) = Divergence -> Bool
isDeadEndDiv Divergence
div
maxWithArity :: ArityType -> Arity -> ArityType
maxWithArity :: ArityType -> Int -> ArityType
maxWithArity at :: ArityType
at@(AT [OneShotInfo]
oss Divergence
div) !Int
ar
| ArityType -> Bool
isDeadEndArityType ArityType
at = ArityType
at
| [OneShotInfo]
oss forall a. [a] -> Int -> Bool
`lengthAtLeast` Int
ar = ArityType
at
| Bool
otherwise = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
ar forall a b. (a -> b) -> a -> b
$ [OneShotInfo]
oss forall a. [a] -> [a] -> [a]
++ forall a. a -> [a]
repeat OneShotInfo
NoOneShotInfo) Divergence
div
minWithArity :: ArityType -> Arity -> ArityType
minWithArity :: ArityType -> Int -> ArityType
minWithArity at :: ArityType
at@(AT [OneShotInfo]
oss Divergence
_) Int
ar
| [OneShotInfo]
oss forall a. [a] -> Int -> Bool
`lengthAtMost` Int
ar = ArityType
at
| Bool
otherwise = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
ar [OneShotInfo]
oss) Divergence
topDiv
takeWhileOneShot :: ArityType -> ArityType
takeWhileOneShot :: ArityType -> ArityType
takeWhileOneShot (AT [OneShotInfo]
oss Divergence
div)
| Divergence -> Bool
isDeadEndDiv Divergence
div = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
oss) Divergence
topDiv
| Bool
otherwise = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
oss) Divergence
div
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
exprEtaExpandArity DynFlags
dflags CoreExpr
e = ArityEnv -> CoreExpr -> ArityType
arityType (DynFlags -> ArityEnv
etaExpandArityEnv DynFlags
dflags) CoreExpr
e
getBotArity :: ArityType -> Maybe Arity
getBotArity :: ArityType -> Maybe Int
getBotArity (AT [OneShotInfo]
oss Divergence
div)
| Divergence -> Bool
isDeadEndDiv Divergence
div = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss
| Bool
otherwise = forall a. Maybe a
Nothing
findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
findRhsArity :: DynFlags -> TyVar -> CoreExpr -> Int -> ArityType
findRhsArity DynFlags
dflags TyVar
bndr CoreExpr
rhs Int
old_arity
= Int -> ArityType -> ArityType
go Int
0 ArityType
botArityType
where
go :: Int -> ArityType -> ArityType
go :: Int -> ArityType -> ArityType
go !Int
n cur_at :: ArityType
cur_at@(AT [OneShotInfo]
oss Divergence
div)
| Bool -> Bool
not (Divergence -> Bool
isDeadEndDiv Divergence
div)
, forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss forall a. Ord a => a -> a -> Bool
<= Int
old_arity = ArityType
cur_at
| ArityType
next_at forall a. Eq a => a -> a -> Bool
== ArityType
cur_at = ArityType
cur_at
| Bool
otherwise =
WARN( debugIsOn && n > 2, text "Exciting arity"
$$ nest 2 (
ppr bndr <+> ppr cur_at <+> ppr next_at
$$ ppr rhs) )
Int -> ArityType -> ArityType
go (Int
nforall a. Num a => a -> a -> a
+Int
1) ArityType
next_at
where
next_at :: ArityType
next_at = ArityType -> ArityType
step ArityType
cur_at
step :: ArityType -> ArityType
step :: ArityType -> ArityType
step ArityType
at =
ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs
where
env :: ArityEnv
env = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv (DynFlags -> ArityEnv
findRhsArityEnv DynFlags
dflags) TyVar
bndr ArityType
at
arityLam :: Id -> ArityType -> ArityType
arityLam :: TyVar -> ArityType -> ArityType
arityLam TyVar
id (AT [OneShotInfo]
oss Divergence
div) = [OneShotInfo] -> Divergence -> ArityType
AT (TyVar -> OneShotInfo
idStateHackOneShotInfo TyVar
id forall a. a -> [a] -> [a]
: [OneShotInfo]
oss) Divergence
div
floatIn :: Bool -> ArityType -> ArityType
floatIn :: Bool -> ArityType -> ArityType
floatIn Bool
cheap ArityType
at
| ArityType -> Bool
isDeadEndArityType ArityType
at Bool -> Bool -> Bool
|| Bool
cheap = ArityType
at
| Bool
otherwise = ArityType -> ArityType
takeWhileOneShot ArityType
at
arityApp :: ArityType -> Bool -> ArityType
arityApp :: ArityType -> Bool -> ArityType
arityApp (AT (OneShotInfo
_:[OneShotInfo]
oss) Divergence
div) Bool
cheap = Bool -> ArityType -> ArityType
floatIn Bool
cheap ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
div)
arityApp ArityType
at Bool
_ = ArityType
at
andArityType :: ArityType -> ArityType -> ArityType
andArityType :: ArityType -> ArityType -> ArityType
andArityType (AT (OneShotInfo
os1:[OneShotInfo]
oss1) Divergence
div1) (AT (OneShotInfo
os2:[OneShotInfo]
oss2) Divergence
div2)
| AT [OneShotInfo]
oss' Divergence
div' <- ArityType -> ArityType -> ArityType
andArityType ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss1 Divergence
div1) ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss2 Divergence
div2)
= [OneShotInfo] -> Divergence -> ArityType
AT ((OneShotInfo
os1 OneShotInfo -> OneShotInfo -> OneShotInfo
`bestOneShot` OneShotInfo
os2) forall a. a -> [a] -> [a]
: [OneShotInfo]
oss') Divergence
div'
andArityType at1 :: ArityType
at1@(AT [] Divergence
div1) ArityType
at2
| Divergence -> Bool
isDeadEndDiv Divergence
div1 = ArityType
at2
| Bool
otherwise = ArityType
at1
andArityType ArityType
at1 at2 :: ArityType
at2@(AT [] Divergence
div2)
| Divergence -> Bool
isDeadEndDiv Divergence
div2 = ArityType
at1
| Bool
otherwise = ArityType
at2
data AnalysisMode
= BotStrictness
| EtaExpandArity { AnalysisMode -> Bool
am_ped_bot :: !Bool
, AnalysisMode -> Bool
am_dicts_cheap :: !Bool }
| FindRhsArity { am_ped_bot :: !Bool
, am_dicts_cheap :: !Bool
, AnalysisMode -> IdEnv ArityType
am_sigs :: !(IdEnv ArityType) }
data ArityEnv
= AE
{ ArityEnv -> AnalysisMode
ae_mode :: !AnalysisMode
, ArityEnv -> IdSet
ae_joins :: !IdSet
}
botStrictnessArityEnv :: ArityEnv
botStrictnessArityEnv :: ArityEnv
botStrictnessArityEnv = AE { ae_mode :: AnalysisMode
ae_mode = AnalysisMode
BotStrictness, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
etaExpandArityEnv :: DynFlags -> ArityEnv
etaExpandArityEnv :: DynFlags -> ArityEnv
etaExpandArityEnv DynFlags
dflags
= AE { ae_mode :: AnalysisMode
ae_mode = EtaExpandArity { am_ped_bot :: Bool
am_ped_bot = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
, am_dicts_cheap :: Bool
am_dicts_cheap = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DictsCheap DynFlags
dflags }
, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
findRhsArityEnv :: DynFlags -> ArityEnv
findRhsArityEnv :: DynFlags -> ArityEnv
findRhsArityEnv DynFlags
dflags
= AE { ae_mode :: AnalysisMode
ae_mode = FindRhsArity { am_ped_bot :: Bool
am_ped_bot = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
, am_dicts_cheap :: Bool
am_dicts_cheap = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DictsCheap DynFlags
dflags
, am_sigs :: IdEnv ArityType
am_sigs = forall a. VarEnv a
emptyVarEnv }
, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv IdEnv ArityType -> IdEnv ArityType
f env :: ArityEnv
env@AE { ae_mode :: ArityEnv -> AnalysisMode
ae_mode = am :: AnalysisMode
am@FindRhsArity{am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs} } =
ArityEnv
env { ae_mode :: AnalysisMode
ae_mode = AnalysisMode
am { am_sigs :: IdEnv ArityType
am_sigs = IdEnv ArityType -> IdEnv ArityType
f IdEnv ArityType
sigs } }
modifySigEnv IdEnv ArityType -> IdEnv ArityType
_ ArityEnv
env = ArityEnv
env
{-# INLINE modifySigEnv #-}
del_sig_env :: Id -> ArityEnv -> ArityEnv
del_sig_env :: TyVar -> ArityEnv -> ArityEnv
del_sig_env TyVar
id = (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> TyVar -> VarEnv a
delVarEnv IdEnv ArityType
sigs TyVar
id)
{-# INLINE del_sig_env #-}
del_sig_env_list :: [Id] -> ArityEnv -> ArityEnv
del_sig_env_list :: [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
ids = (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> [TyVar] -> VarEnv a
delVarEnvList IdEnv ArityType
sigs [TyVar]
ids)
{-# INLINE del_sig_env_list #-}
del_join_env :: JoinId -> ArityEnv -> ArityEnv
del_join_env :: TyVar -> ArityEnv -> ArityEnv
del_join_env TyVar
id env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins })
= ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet -> TyVar -> IdSet
delVarSet IdSet
joins TyVar
id }
{-# INLINE del_join_env #-}
del_join_env_list :: [JoinId] -> ArityEnv -> ArityEnv
del_join_env_list :: [TyVar] -> ArityEnv -> ArityEnv
del_join_env_list [TyVar]
ids env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins })
= ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet -> [TyVar] -> IdSet
delVarSetList IdSet
joins [TyVar]
ids }
{-# INLINE del_join_env_list #-}
extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
extendJoinEnv :: ArityEnv -> [TyVar] -> ArityEnv
extendJoinEnv env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins }) [TyVar]
join_ids
= [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
join_ids
forall a b. (a -> b) -> a -> b
$ ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet
joins IdSet -> [TyVar] -> IdSet
`extendVarSetList` [TyVar]
join_ids }
extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
extendSigEnv :: ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
id ArityType
ar_ty
= TyVar -> ArityEnv -> ArityEnv
del_join_env TyVar
id ((IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> TyVar -> a -> VarEnv a
extendVarEnv IdEnv ArityType
sigs TyVar
id ArityType
ar_ty) ArityEnv
env)
delInScope :: ArityEnv -> Id -> ArityEnv
delInScope :: ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
id = TyVar -> ArityEnv -> ArityEnv
del_join_env TyVar
id forall a b. (a -> b) -> a -> b
$ TyVar -> ArityEnv -> ArityEnv
del_sig_env TyVar
id ArityEnv
env
delInScopeList :: ArityEnv -> [Id] -> ArityEnv
delInScopeList :: ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env [TyVar]
ids = [TyVar] -> ArityEnv -> ArityEnv
del_join_env_list [TyVar]
ids forall a b. (a -> b) -> a -> b
$ [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
ids ArityEnv
env
lookupSigEnv :: ArityEnv -> Id -> Maybe ArityType
lookupSigEnv :: ArityEnv -> TyVar -> Maybe ArityType
lookupSigEnv AE{ ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode } TyVar
id = case AnalysisMode
mode of
AnalysisMode
BotStrictness -> forall a. Maybe a
Nothing
EtaExpandArity{} -> forall a. Maybe a
Nothing
FindRhsArity{ am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs } -> forall a. VarEnv a -> TyVar -> Maybe a
lookupVarEnv IdEnv ArityType
sigs TyVar
id
pedanticBottoms :: ArityEnv -> Bool
pedanticBottoms :: ArityEnv -> Bool
pedanticBottoms AE{ ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode } = case AnalysisMode
mode of
AnalysisMode
BotStrictness -> Bool
True
EtaExpandArity{ am_ped_bot :: AnalysisMode -> Bool
am_ped_bot = Bool
ped_bot } -> Bool
ped_bot
FindRhsArity{ am_ped_bot :: AnalysisMode -> Bool
am_ped_bot = Bool
ped_bot } -> Bool
ped_bot
myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap AE{ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode} CoreExpr
e Maybe Type
mb_ty = case AnalysisMode
mode of
AnalysisMode
BotStrictness -> Bool
False
AnalysisMode
_ -> Bool
cheap_dict Bool -> Bool -> Bool
|| CoreExpr -> Bool
cheap_fun CoreExpr
e
where
cheap_dict :: Bool
cheap_dict = AnalysisMode -> Bool
am_dicts_cheap AnalysisMode
mode Bool -> Bool -> Bool
&& forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Bool
isDictTy Maybe Type
mb_ty forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just Bool
True
cheap_fun :: CoreExpr -> Bool
cheap_fun CoreExpr
e = case AnalysisMode
mode of
#if __GLASGOW_HASKELL__ <= 900
BotStrictness -> panic "impossible"
#endif
EtaExpandArity{} -> CoreExpr -> Bool
exprIsCheap CoreExpr
e
FindRhsArity{am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs} -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX (IdEnv ArityType -> CheapAppFun
myIsCheapApp IdEnv ArityType
sigs) CoreExpr
e
myIsCheapApp :: IdEnv ArityType -> CheapAppFun
myIsCheapApp :: IdEnv ArityType -> CheapAppFun
myIsCheapApp IdEnv ArityType
sigs TyVar
fn Int
n_val_args = case forall a. VarEnv a -> TyVar -> Maybe a
lookupVarEnv IdEnv ArityType
sigs TyVar
fn of
Maybe ArityType
Nothing -> CheapAppFun
isCheapApp TyVar
fn Int
n_val_args
Just (AT [OneShotInfo]
oss Divergence
div)
| Divergence -> Bool
isDeadEndDiv Divergence
div -> Bool
True
| Int
n_val_args forall a. Ord a => a -> a -> Bool
< forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss -> Bool
True
| Bool
otherwise -> Bool
False
arityType :: ArityEnv -> CoreExpr -> ArityType
arityType :: ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env (Cast CoreExpr
e Coercion
co)
= ArityType -> Int -> ArityType
minWithArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e) Int
co_arity
where
co_arity :: Int
co_arity = forall (t :: * -> *) a. Foldable t => t a -> Int
length (Type -> [OneShotInfo]
typeArity (Coercion -> Type
coercionRKind Coercion
co))
arityType ArityEnv
env (Var TyVar
v)
| TyVar
v TyVar -> IdSet -> Bool
`elemVarSet` ArityEnv -> IdSet
ae_joins ArityEnv
env
= ArityType
botArityType
| Just ArityType
at <- ArityEnv -> TyVar -> Maybe ArityType
lookupSigEnv ArityEnv
env TyVar
v
= ArityType
at
| Bool
otherwise
= TyVar -> ArityType
idArityType TyVar
v
arityType ArityEnv
env (Lam TyVar
x CoreExpr
e)
| TyVar -> Bool
isId TyVar
x = TyVar -> ArityType -> ArityType
arityLam TyVar
x (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
| Bool
otherwise = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e
where
env' :: ArityEnv
env' = ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
x
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 -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
arg forall a. Maybe a
Nothing)
arityType ArityEnv
env (Case CoreExpr
scrut TyVar
bndr Type
_ [Alt TyVar]
alts)
| CoreExpr -> Bool
exprIsDeadEnd CoreExpr
scrut Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Alt TyVar]
alts
= ArityType
botArityType
| Bool -> Bool
not (ArityEnv -> Bool
pedanticBottoms ArityEnv
env)
, ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
scrut (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
bndr))
= ArityType
alts_type
| CoreExpr -> Bool
exprOkForSpeculation CoreExpr
scrut
= ArityType
alts_type
| Bool
otherwise
= ArityType -> ArityType
takeWhileOneShot ArityType
alts_type
where
env' :: ArityEnv
env' = ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
bndr
arity_type_alt :: Alt TyVar -> ArityType
arity_type_alt (Alt AltCon
_con [TyVar]
bndrs CoreExpr
rhs) = ArityEnv -> CoreExpr -> ArityType
arityType (ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env' [TyVar]
bndrs) CoreExpr
rhs
alts_type :: ArityType
alts_type = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 ArityType -> ArityType -> ArityType
andArityType (forall a b. (a -> b) -> [a] -> [b]
map Alt TyVar -> ArityType
arity_type_alt [Alt TyVar]
alts)
arityType ArityEnv
env (Let (NonRec TyVar
j CoreExpr
rhs) CoreExpr
body)
| Just Int
join_arity <- TyVar -> Maybe Int
isJoinId_maybe TyVar
j
, ([TyVar]
_, CoreExpr
rhs_body) <- forall b. Int -> Expr b -> ([b], Expr b)
collectNBinders Int
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 -> [TyVar] -> ArityEnv
extendJoinEnv ArityEnv
env [TyVar
j]
arityType ArityEnv
env (Let (Rec [(TyVar, CoreExpr)]
pairs) CoreExpr
body)
| ((TyVar
j,CoreExpr
_):[(TyVar, CoreExpr)]
_) <- [(TyVar, CoreExpr)]
pairs
, TyVar -> Bool
isJoinId TyVar
j
=
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (ArityType -> ArityType -> ArityType
andArityType forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TyVar, CoreExpr) -> ArityType
do_one) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
body) [(TyVar, CoreExpr)]
pairs
where
env' :: ArityEnv
env' = ArityEnv -> [TyVar] -> ArityEnv
extendJoinEnv ArityEnv
env (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(TyVar, CoreExpr)]
pairs)
do_one :: (TyVar, CoreExpr) -> ArityType
do_one (TyVar
j,CoreExpr
rhs)
| Just Int
arity <- TyVar -> Maybe Int
isJoinId_maybe TyVar
j
= ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall b. Int -> Expr b -> ([b], Expr b)
collectNBinders Int
arity CoreExpr
rhs
| Bool
otherwise
= forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"arityType:joinrec" (forall a. Outputable a => a -> SDoc
ppr [(TyVar, CoreExpr)]
pairs)
arityType ArityEnv
env (Let (NonRec TyVar
b CoreExpr
r) CoreExpr
e)
= Bool -> ArityType -> ArityType
floatIn Bool
cheap_rhs (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
where
cheap_rhs :: Bool
cheap_rhs = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
r (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))
env' :: ArityEnv
env' = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
b (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
r)
arityType ArityEnv
env (Let (Rec [(TyVar, CoreExpr)]
prs) CoreExpr
e)
= Bool -> ArityType -> ArityType
floatIn (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (TyVar, CoreExpr) -> Bool
is_cheap [(TyVar, CoreExpr)]
prs) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
where
env' :: ArityEnv
env' = ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(TyVar, CoreExpr)]
prs)
is_cheap :: (TyVar, CoreExpr) -> Bool
is_cheap (TyVar
b,CoreExpr
e) = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env' CoreExpr
e (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))
arityType ArityEnv
env (Tick CoreTickish
t CoreExpr
e)
| Bool -> Bool
not (forall (pass :: TickishPass). GenTickish pass -> Bool
tickishIsCode CoreTickish
t) = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e
arityType ArityEnv
_ CoreExpr
_ = ArityType
topArityType
idArityType :: Id -> ArityType
idArityType :: TyVar -> ArityType
idArityType TyVar
v
| StrictSig
strict_sig <- TyVar -> StrictSig
idStrictness TyVar
v
, Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ StrictSig -> Bool
isTopSig StrictSig
strict_sig
, ([Demand]
ds, Divergence
div) <- StrictSig -> ([Demand], Divergence)
splitStrictSig StrictSig
strict_sig
, let arity :: Int
arity = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Demand]
ds
= [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
arity [OneShotInfo]
one_shots) Divergence
div
| Bool
otherwise
= [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take (TyVar -> Int
idArity TyVar
v) [OneShotInfo]
one_shots) Divergence
topDiv
where
one_shots :: [OneShotInfo]
one_shots :: [OneShotInfo]
one_shots = Type -> [OneShotInfo]
typeArity (TyVar -> Type
idType TyVar
v)
etaExpand :: Arity -> CoreExpr -> CoreExpr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpand :: Int -> CoreExpr -> CoreExpr
etaExpand Int
n CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (forall a. Int -> a -> [a]
replicate Int
n OneShotInfo
NoOneShotInfo) CoreExpr
orig_expr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpandAT (AT [OneShotInfo]
oss Divergence
_) CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
oss CoreExpr
orig_expr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots (Cast CoreExpr
expr Coercion
co)
= CoreExpr -> Coercion -> CoreExpr
mkCast ([OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots CoreExpr
expr) Coercion
co
eta_expand [OneShotInfo]
one_shots CoreExpr
orig_expr
= [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
one_shots [] CoreExpr
orig_expr
where
go :: [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [] [TyVar]
_ CoreExpr
_ = CoreExpr
orig_expr
go oss :: [OneShotInfo]
oss@(OneShotInfo
_:[OneShotInfo]
oss1) [TyVar]
vs (Lam TyVar
v CoreExpr
body)
| TyVar -> Bool
isTyVar TyVar
v = [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss (TyVar
vforall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body
| Bool
otherwise = [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss1 (TyVar
vforall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body
go [OneShotInfo]
oss [TyVar]
rev_vs CoreExpr
expr
=
CoreExpr -> CoreExpr
retick forall a b. (a -> b) -> a -> b
$ EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs EtaInfo
top_eis forall a b. (a -> b) -> a -> b
$
InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp InScopeSet
in_scope' CoreExpr
sexpr EtaInfo
eis
where
in_scope :: InScopeSet
in_scope = IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
expr)
(InScopeSet
in_scope', eis :: EtaInfo
eis@(EI [TyVar]
eta_bndrs MCoercionR
mco))
= [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, EtaInfo)
mkEtaWW [OneShotInfo]
oss (forall a. Outputable a => a -> SDoc
ppr CoreExpr
orig_expr) InScopeSet
in_scope (CoreExpr -> Type
exprType CoreExpr
expr)
top_bndrs :: [TyVar]
top_bndrs = forall a. [a] -> [a]
reverse [TyVar]
rev_vs
top_eis :: EtaInfo
top_eis = [TyVar] -> MCoercionR -> EtaInfo
EI ([TyVar]
top_bndrs forall a. [a] -> [a] -> [a]
++ [TyVar]
eta_bndrs) ([TyVar] -> MCoercionR -> MCoercionR
mkPiMCos [TyVar]
top_bndrs MCoercionR
mco)
(CoreExpr
expr', [CoreExpr]
args) = forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
([CoreTickish]
ticks, CoreExpr
expr'') = forall b.
(CoreTickish -> Bool) -> Expr b -> ([CoreTickish], Expr b)
stripTicksTop forall (pass :: TickishPass). GenTickish pass -> Bool
tickishFloatable CoreExpr
expr'
sexpr :: CoreExpr
sexpr = forall b. Expr b -> [Expr b] -> Expr b
mkApps CoreExpr
expr'' [CoreExpr]
args
retick :: CoreExpr -> CoreExpr
retick CoreExpr
expr = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CoreTickish -> CoreExpr -> CoreExpr
mkTick CoreExpr
expr [CoreTickish]
ticks
data EtaInfo = EI [Var] MCoercionR
etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp InScopeSet
in_scope CoreExpr
expr EtaInfo
eis
= Subst -> CoreExpr -> EtaInfo -> CoreExpr
go (InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope) CoreExpr
expr EtaInfo
eis
where
go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst (Tick CoreTickish
t CoreExpr
e) EtaInfo
eis
= forall b. CoreTickish -> Expr b -> Expr b
Tick (Subst -> CoreTickish -> CoreTickish
substTickish Subst
subst CoreTickish
t) (Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst CoreExpr
e EtaInfo
eis)
go Subst
subst (Cast CoreExpr
e Coercion
co) (EI [TyVar]
bs MCoercionR
mco)
= Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst CoreExpr
e ([TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs (HasCallStack => Subst -> Coercion -> Coercion
Core.substCo Subst
subst Coercion
co Coercion -> MCoercionR -> MCoercionR
`mkTransMCoR` MCoercionR
mco))
go Subst
subst (Case CoreExpr
e TyVar
b Type
ty [Alt TyVar]
alts) EtaInfo
eis
= forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Core.substExprSC Subst
subst CoreExpr
e) TyVar
b1 Type
ty' [Alt TyVar]
alts'
where
(Subst
subst1, TyVar
b1) = Subst -> TyVar -> (Subst, TyVar)
Core.substBndr Subst
subst TyVar
b
alts' :: [Alt TyVar]
alts' = forall a b. (a -> b) -> [a] -> [b]
map Alt TyVar -> Alt TyVar
subst_alt [Alt TyVar]
alts
ty' :: Type
ty' = Type -> EtaInfo -> Type
etaInfoAppTy (Subst -> Type -> Type
Core.substTy Subst
subst Type
ty) EtaInfo
eis
subst_alt :: Alt TyVar -> Alt TyVar
subst_alt (Alt AltCon
con [TyVar]
bs CoreExpr
rhs) = forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con [TyVar]
bs' (Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst2 CoreExpr
rhs EtaInfo
eis)
where
(Subst
subst2,[TyVar]
bs') = Subst -> [TyVar] -> (Subst, [TyVar])
Core.substBndrs Subst
subst1 [TyVar]
bs
go Subst
subst (Let Bind TyVar
b CoreExpr
e) EtaInfo
eis
| Bool -> Bool
not (Bind TyVar -> Bool
isJoinBind Bind TyVar
b)
= forall b. Bind b -> Expr b -> Expr b
Let Bind TyVar
b' (Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst' CoreExpr
e EtaInfo
eis)
where
(Subst
subst', Bind TyVar
b') = HasDebugCallStack => Subst -> Bind TyVar -> (Subst, Bind TyVar)
Core.substBindSC Subst
subst Bind TyVar
b
go Subst
subst (Lam TyVar
v CoreExpr
e) (EI (TyVar
b:[TyVar]
bs) MCoercionR
mco)
| Just (CoreExpr
arg,MCoercionR
mco') <- MCoercionR -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushMCoArg MCoercionR
mco (forall b. TyVar -> Expr b
varToCoreExpr TyVar
b)
= Subst -> CoreExpr -> EtaInfo -> CoreExpr
go (Subst -> TyVar -> CoreExpr -> Subst
Core.extendSubst Subst
subst TyVar
v CoreExpr
arg) CoreExpr
e ([TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs MCoercionR
mco')
go Subst
subst CoreExpr
e (EI [TyVar]
bs MCoercionR
mco) = HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Core.substExprSC Subst
subst CoreExpr
e
CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR
mco
forall b. Expr b -> [TyVar] -> Expr b
`mkVarApps` [TyVar]
bs
etaInfoAppTy :: Type -> EtaInfo -> Type
etaInfoAppTy :: Type -> EtaInfo -> Type
etaInfoAppTy Type
ty (EI [TyVar]
bs MCoercionR
mco)
= SDoc -> Type -> [CoreExpr] -> Type
applyTypeToArgs (String -> SDoc
text String
"etaInfoAppTy") Type
ty1 (forall a b. (a -> b) -> [a] -> [b]
map forall b. TyVar -> Expr b
varToCoreExpr [TyVar]
bs)
where
ty1 :: Type
ty1 = case MCoercionR
mco of
MCoercionR
MRefl -> Type
ty
MCo Coercion
co -> Coercion -> Type
coercionRKind Coercion
co
etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs (EI [TyVar]
bs MCoercionR
mco) CoreExpr
expr = (forall b. [b] -> Expr b -> Expr b
mkLams [TyVar]
bs CoreExpr
expr) CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR -> MCoercionR
mkSymMCo MCoercionR
mco
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
= Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
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
-> (InScopeSet, EtaInfo)
go :: Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
_ [] TCvSubst
subst Type
_
= (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [TyVar] -> MCoercionR -> EtaInfo
EI [] MCoercionR
MRefl)
go Int
n oss :: [OneShotInfo]
oss@(OneShotInfo
one_shot:[OneShotInfo]
oss1) TCvSubst
subst Type
ty
| Just (TyVar
tcv,Type
ty') <- Type -> Maybe (TyVar, Type)
splitForAllTyCoVar_maybe Type
ty
, (TCvSubst
subst', TyVar
tcv') <- HasCallStack => TCvSubst -> TyVar -> (TCvSubst, TyVar)
Type.substVarBndr TCvSubst
subst TyVar
tcv
, let oss' :: [OneShotInfo]
oss' | TyVar -> Bool
isTyVar TyVar
tcv = [OneShotInfo]
oss
| Bool
otherwise = [OneShotInfo]
oss1
, (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
n [OneShotInfo]
oss' TCvSubst
subst' Type
ty'
= (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI (TyVar
tcv' forall a. a -> [a] -> [a]
: [TyVar]
bs) (TyVar -> MCoercionR -> MCoercionR
mkHomoForAllMCo TyVar
tcv' MCoercionR
mco))
| 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', TyVar
eta_id) <- Int -> TCvSubst -> Scaled Type -> (TCvSubst, TyVar)
freshEtaId Int
n TCvSubst
subst (forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
, let eta_id' :: TyVar
eta_id' = TyVar
eta_id TyVar -> OneShotInfo -> TyVar
`setIdOneShotInfo` OneShotInfo
one_shot
, (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go (Int
nforall a. Num a => a -> a -> a
+Int
1) [OneShotInfo]
oss1 TCvSubst
subst' Type
res_ty
= (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI (TyVar
eta_id' forall a. a -> [a] -> [a]
: [TyVar]
bs) (Scaled Type -> MCoercionR -> MCoercionR
mkFunResMCo (TyVar -> Scaled Type
idScaledType TyVar
eta_id') MCoercionR
mco))
| Just (Coercion
co, Type
ty') <- Type -> Maybe (Coercion, Type)
topNormaliseNewType_maybe Type
ty
,
let co' :: Coercion
co' = HasCallStack => TCvSubst -> Coercion -> Coercion
Type.substCo TCvSubst
subst Coercion
co
, (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
n [OneShotInfo]
oss TCvSubst
subst Type
ty'
= (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs (Coercion -> MCoercionR -> MCoercionR
mkTransMCoR Coercion
co' MCoercionR
mco))
| Bool
otherwise
= WARN( True, (ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr )
(TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [TyVar] -> MCoercionR -> EtaInfo
EI [] MCoercionR
MRefl)
pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], MCoercion)
pushCoArgs :: Coercion -> [CoreExpr] -> Maybe ([CoreExpr], MCoercionR)
pushCoArgs Coercion
co [] = forall (m :: * -> *) a. Monad m => a -> m a
return ([], Coercion -> MCoercionR
MCo Coercion
co)
pushCoArgs Coercion
co (CoreExpr
arg:[CoreExpr]
args) = do { (CoreExpr
arg', MCoercionR
m_co1) <- Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg Coercion
co CoreExpr
arg
; case MCoercionR
m_co1 of
MCo Coercion
co1 -> do { ([CoreExpr]
args', MCoercionR
m_co2) <- Coercion -> [CoreExpr] -> Maybe ([CoreExpr], MCoercionR)
pushCoArgs Coercion
co1 [CoreExpr]
args
; forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'forall a. a -> [a] -> [a]
:[CoreExpr]
args', MCoercionR
m_co2) }
MCoercionR
MRefl -> forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'forall a. a -> [a] -> [a]
:[CoreExpr]
args, MCoercionR
MRefl) }
pushMCoArg :: MCoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
pushMCoArg :: MCoercionR -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushMCoArg MCoercionR
MRefl CoreExpr
arg = forall a. a -> Maybe a
Just (CoreExpr
arg, MCoercionR
MRefl)
pushMCoArg (MCo Coercion
co) CoreExpr
arg = Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg Coercion
co CoreExpr
arg
pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
pushCoArg :: Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg Coercion
co (Type Type
ty) = do { (Type
ty', MCoercionR
m_co') <- Coercion -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg Coercion
co Type
ty
; forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Type -> Expr b
Type Type
ty', MCoercionR
m_co') }
pushCoArg Coercion
co CoreExpr
val_arg = do { (MCoercionR
arg_co, MCoercionR
m_co') <- Coercion -> Maybe (MCoercionR, MCoercionR)
pushCoValArg Coercion
co
; forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
val_arg CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR
arg_co, MCoercionR
m_co') }
pushCoTyArg :: CoercionR -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg :: Coercion -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg Coercion
co Type
ty
| Coercion -> Bool
isReflCo Coercion
co
= forall a. a -> Maybe a
Just (Type
ty, MCoercionR
MRefl)
| Type -> Bool
isForAllTy_ty Type
tyL
= ASSERT2( isForAllTy_ty tyR, ppr co $$ ppr ty )
forall a. a -> Maybe a
Just (Type
ty Type -> Coercion -> Type
`mkCastTy` Coercion
co1, Coercion -> MCoercionR
MCo Coercion
co2)
| Bool
otherwise
= forall a. Maybe a
Nothing
where
Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
co1 :: Coercion
co1 = Coercion -> Coercion
mkSymCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)
co2 :: Coercion
co2 = Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Role -> Type -> Coercion -> Coercion
mkGReflLeftCo Role
Nominal Type
ty Coercion
co1)
pushCoValArg :: CoercionR -> Maybe (MCoercionR, MCoercionR)
pushCoValArg :: Coercion -> Maybe (MCoercionR, MCoercionR)
pushCoValArg Coercion
co
| Coercion -> Bool
isReflCo Coercion
co
= forall a. a -> Maybe a
Just (MCoercionR
MRefl, MCoercionR
MRefl)
| Type -> Bool
isFunTy Type
tyL
, (Coercion
co_mult, Coercion
co1, Coercion
co2) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
, Coercion -> Bool
isReflexiveCo Coercion
co_mult
= ASSERT2( isFunTy tyR, ppr co $$ ppr arg )
forall a. a -> Maybe a
Just (Coercion -> MCoercionR
coToMCo (Coercion -> Coercion
mkSymCo Coercion
co1), Coercion -> MCoercionR
coToMCo Coercion
co2)
| Bool
otherwise
= forall a. Maybe a
Nothing
where
arg :: Type
arg = Type -> Type
funArgTy Type
tyR
Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
pushCoercionIntoLambda
:: HasDebugCallStack => InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
pushCoercionIntoLambda :: HasDebugCallStack =>
InScopeSet
-> TyVar -> CoreExpr -> Coercion -> Maybe (TyVar, CoreExpr)
pushCoercionIntoLambda InScopeSet
in_scope TyVar
x CoreExpr
e Coercion
co
| ASSERT(not (isTyVar x) && not (isCoVar x)) True
, Pair Type
s1s2 Type
t1t2 <- Coercion -> Pair Type
coercionKind Coercion
co
, Just (Type
_, Type
_s1,Type
_s2) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
s1s2
, Just (Type
w1, Type
t1,Type
_t2) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
t1t2
, (Coercion
co_mult, Coercion
co1, Coercion
co2) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
, Coercion -> Bool
isReflexiveCo Coercion
co_mult
= let
x' :: TyVar
x' = TyVar
x TyVar -> Type -> TyVar
`setIdType` Type
t1 TyVar -> Type -> TyVar
`setIdMult` Type
w1
in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> TyVar -> InScopeSet
`extendInScopeSet` TyVar
x'
subst :: Subst
subst = Subst -> TyVar -> CoreExpr -> Subst
extendIdSubst (InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope')
TyVar
x
(CoreExpr -> Coercion -> CoreExpr
mkCast (forall b. TyVar -> Expr b
Var TyVar
x') Coercion
co1)
in forall a. a -> Maybe a
Just (TyVar
x', HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
substExpr Subst
subst CoreExpr
e CoreExpr -> Coercion -> CoreExpr
`mkCast` Coercion
co2)
| Bool
otherwise
=
forall a. Maybe a
Nothing
pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
-> Maybe (DataCon
, [Type]
, [CoreExpr])
pushCoDataCon :: DataCon
-> [CoreExpr] -> Coercion -> Maybe (DataCon, [Type], [CoreExpr])
pushCoDataCon DataCon
dc [CoreExpr]
dc_args Coercion
co
| Coercion -> Bool
isReflCo Coercion
co Bool -> Bool -> Bool
|| Type
from_ty Type -> Type -> Bool
`eqType` Type
to_ty
, let ([CoreExpr]
univ_ty_args, [CoreExpr]
rest_args) = forall b a. [b] -> [a] -> ([a], [a])
splitAtList (DataCon -> [TyVar]
dataConUnivTyVars DataCon
dc) [CoreExpr]
dc_args
= forall a. a -> Maybe a
Just (DataCon
dc, forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType [CoreExpr]
univ_ty_args, [CoreExpr]
rest_args)
| Just (TyCon
to_tc, [Type]
to_tc_arg_tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
to_ty
, TyCon
to_tc forall a. Eq a => a -> a -> Bool
== DataCon -> TyCon
dataConTyCon DataCon
dc
= let
tc_arity :: Int
tc_arity = TyCon -> Int
tyConArity TyCon
to_tc
dc_univ_tyvars :: [TyVar]
dc_univ_tyvars = DataCon -> [TyVar]
dataConUnivTyVars DataCon
dc
dc_ex_tcvars :: [TyVar]
dc_ex_tcvars = DataCon -> [TyVar]
dataConExTyCoVars DataCon
dc
arg_tys :: [Scaled Type]
arg_tys = DataCon -> [Scaled Type]
dataConRepArgTys DataCon
dc
non_univ_args :: [CoreExpr]
non_univ_args = forall b a. [b] -> [a] -> [a]
dropList [TyVar]
dc_univ_tyvars [CoreExpr]
dc_args
([CoreExpr]
ex_args, [CoreExpr]
val_args) = forall b a. [b] -> [a] -> ([a], [a])
splitAtList [TyVar]
dc_ex_tcvars [CoreExpr]
non_univ_args
omegas :: [Coercion]
omegas = Int -> Coercion -> [Role] -> [Coercion]
decomposeCo Int
tc_arity Coercion
co (TyCon -> [Role]
tyConRolesRepresentational TyCon
to_tc)
(Type -> Coercion
psi_subst, [Type]
to_ex_arg_tys)
= Role
-> [TyVar]
-> [Coercion]
-> [TyVar]
-> [Type]
-> (Type -> Coercion, [Type])
liftCoSubstWithEx Role
Representational
[TyVar]
dc_univ_tyvars
[Coercion]
omegas
[TyVar]
dc_ex_tcvars
(forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType [CoreExpr]
ex_args)
new_val_args :: [CoreExpr]
new_val_args = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> CoreExpr -> CoreExpr
cast_arg (forall a b. (a -> b) -> [a] -> [b]
map forall a. Scaled a -> a
scaledThing [Scaled Type]
arg_tys) [CoreExpr]
val_args
cast_arg :: Type -> CoreExpr -> CoreExpr
cast_arg Type
arg_ty CoreExpr
arg = CoreExpr -> Coercion -> CoreExpr
mkCast CoreExpr
arg (Type -> Coercion
psi_subst Type
arg_ty)
to_ex_args :: [CoreExpr]
to_ex_args = forall a b. (a -> b) -> [a] -> [b]
map forall b. Type -> Expr b
Type [Type]
to_ex_arg_tys
dump_doc :: SDoc
dump_doc = [SDoc] -> SDoc
vcat [forall a. Outputable a => a -> SDoc
ppr DataCon
dc, forall a. Outputable a => a -> SDoc
ppr [TyVar]
dc_univ_tyvars, forall a. Outputable a => a -> SDoc
ppr [TyVar]
dc_ex_tcvars,
forall a. Outputable a => a -> SDoc
ppr [Scaled Type]
arg_tys, forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
dc_args,
forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
ex_args, forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
val_args, forall a. Outputable a => a -> SDoc
ppr Coercion
co, forall a. Outputable a => a -> SDoc
ppr Type
from_ty, forall a. Outputable a => a -> SDoc
ppr Type
to_ty, forall a. Outputable a => a -> SDoc
ppr TyCon
to_tc
, forall a. Outputable a => a -> SDoc
ppr forall a b. (a -> b) -> a -> b
$ TyCon -> [Type] -> Type
mkTyConApp TyCon
to_tc (forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType forall a b. (a -> b) -> a -> b
$ forall b a. [b] -> [a] -> [a]
takeList [TyVar]
dc_univ_tyvars [CoreExpr]
dc_args) ]
in
ASSERT2( eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args)), dump_doc )
ASSERT2( equalLength val_args arg_tys, dump_doc )
forall a. a -> Maybe a
Just (DataCon
dc, [Type]
to_tc_arg_tys, [CoreExpr]
to_ex_args forall a. [a] -> [a] -> [a]
++ [CoreExpr]
new_val_args)
| Bool
otherwise
= forall a. Maybe a
Nothing
where
Pair Type
from_ty Type
to_ty = Coercion -> Pair Type
coercionKind Coercion
co
collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
collectBindersPushingCo :: CoreExpr -> ([TyVar], CoreExpr)
collectBindersPushingCo CoreExpr
e
= [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go [] CoreExpr
e
where
go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
go :: [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go [TyVar]
bs (Lam TyVar
b CoreExpr
e) = [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e
go [TyVar]
bs (Cast CoreExpr
e Coercion
co) = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs CoreExpr
e Coercion
co
go [TyVar]
bs CoreExpr
e = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr
e)
go_c :: [Var] -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
go_c :: [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs (Cast CoreExpr
e Coercion
co1) Coercion
co2 = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs CoreExpr
e (Coercion
co1 Coercion -> Coercion -> Coercion
`mkTransCo` Coercion
co2)
go_c [TyVar]
bs (Lam TyVar
b CoreExpr
e) Coercion
co = [TyVar] -> TyVar -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_lam [TyVar]
bs TyVar
b CoreExpr
e Coercion
co
go_c [TyVar]
bs CoreExpr
e Coercion
co = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr -> Coercion -> CoreExpr
mkCast CoreExpr
e Coercion
co)
go_lam :: [Var] -> Var -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
go_lam :: [TyVar] -> TyVar -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_lam [TyVar]
bs TyVar
b CoreExpr
e Coercion
co
| TyVar -> Bool
isTyVar TyVar
b
, let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
, ASSERT( isForAllTy_ty tyL )
Type -> Bool
isForAllTy_ty Type
tyR
, Coercion -> Bool
isReflCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)
= [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e (Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Type -> Coercion
mkNomReflCo (TyVar -> Type
mkTyVarTy TyVar
b)))
| TyVar -> Bool
isCoVar TyVar
b
, let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
, ASSERT( isForAllTy_co tyL )
Type -> Bool
isForAllTy_co Type
tyR
, Coercion -> Bool
isReflCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)
, let cov :: Coercion
cov = TyVar -> Coercion
mkCoVarCo TyVar
b
= [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e (Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Type -> Coercion
mkNomReflCo (Coercion -> Type
mkCoercionTy Coercion
cov)))
| TyVar -> Bool
isId TyVar
b
, let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
, ASSERT( isFunTy tyL) isFunTy tyR
, (Coercion
co_mult, Coercion
co_arg, Coercion
co_res) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
, Coercion -> Bool
isReflCo Coercion
co_mult
, Coercion -> Bool
isReflCo Coercion
co_arg
= [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e Coercion
co_res
| Bool
otherwise = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr -> Coercion -> CoreExpr
mkCast (forall b. b -> Expr b -> Expr b
Lam TyVar
b CoreExpr
e) Coercion
co)
etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaExpandToJoinPoint :: Int -> CoreExpr -> ([TyVar], CoreExpr)
etaExpandToJoinPoint Int
join_arity CoreExpr
expr
= Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
join_arity [] CoreExpr
expr
where
go :: Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
0 [TyVar]
rev_bs CoreExpr
e = (forall a. [a] -> [a]
reverse [TyVar]
rev_bs, CoreExpr
e)
go Int
n [TyVar]
rev_bs (Lam TyVar
b CoreExpr
e) = Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nforall a. Num a => a -> a -> a
-Int
1) (TyVar
b forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) CoreExpr
e
go Int
n [TyVar]
rev_bs CoreExpr
e = case Int -> CoreExpr -> ([TyVar], CoreExpr)
etaBodyForJoinPoint Int
n CoreExpr
e of
([TyVar]
bs, CoreExpr
e') -> (forall a. [a] -> [a]
reverse [TyVar]
rev_bs forall a. [a] -> [a] -> [a]
++ [TyVar]
bs, CoreExpr
e')
etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule :: Int -> CoreRule -> CoreRule
etaExpandToJoinPointRule Int
_ rule :: CoreRule
rule@(BuiltinRule {})
= WARN(True, (sep [text "Can't eta-expand built-in rule:", ppr rule]))
CoreRule
rule
etaExpandToJoinPointRule Int
join_arity rule :: CoreRule
rule@(Rule { ru_bndrs :: CoreRule -> [TyVar]
ru_bndrs = [TyVar]
bndrs, ru_rhs :: CoreRule -> CoreExpr
ru_rhs = CoreExpr
rhs
, ru_args :: CoreRule -> [CoreExpr]
ru_args = [CoreExpr]
args })
| Int
need_args forall a. Eq a => a -> a -> Bool
== Int
0
= CoreRule
rule
| Int
need_args forall a. Ord a => a -> a -> Bool
< Int
0
= forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaExpandToJoinPointRule" (forall a. Outputable a => a -> SDoc
ppr Int
join_arity SDoc -> SDoc -> SDoc
$$ forall a. Outputable a => a -> SDoc
ppr CoreRule
rule)
| Bool
otherwise
= CoreRule
rule { ru_bndrs :: [TyVar]
ru_bndrs = [TyVar]
bndrs forall a. [a] -> [a] -> [a]
++ [TyVar]
new_bndrs, ru_args :: [CoreExpr]
ru_args = [CoreExpr]
args forall a. [a] -> [a] -> [a]
++ [CoreExpr]
new_args
, ru_rhs :: CoreExpr
ru_rhs = CoreExpr
new_rhs }
where
need_args :: Int
need_args = Int
join_arity forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Int
length [CoreExpr]
args
([TyVar]
new_bndrs, CoreExpr
new_rhs) = Int -> CoreExpr -> ([TyVar], CoreExpr)
etaBodyForJoinPoint Int
need_args CoreExpr
rhs
new_args :: [CoreExpr]
new_args = forall b. [TyVar] -> [Expr b]
varsToCoreExprs [TyVar]
new_bndrs
etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint :: Int -> CoreExpr -> ([TyVar], CoreExpr)
etaBodyForJoinPoint Int
need_args CoreExpr
body
= Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
need_args (CoreExpr -> Type
exprType CoreExpr
body) (CoreExpr -> TCvSubst
init_subst CoreExpr
body) [] CoreExpr
body
where
go :: Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
0 Type
_ TCvSubst
_ [TyVar]
rev_bs CoreExpr
e
= (forall a. [a] -> [a]
reverse [TyVar]
rev_bs, CoreExpr
e)
go Int
n Type
ty TCvSubst
subst [TyVar]
rev_bs CoreExpr
e
| Just (TyVar
tv, Type
res_ty) <- Type -> Maybe (TyVar, Type)
splitForAllTyCoVar_maybe Type
ty
, let (TCvSubst
subst', TyVar
tv') = HasCallStack => TCvSubst -> TyVar -> (TCvSubst, TyVar)
substVarBndr TCvSubst
subst TyVar
tv
= Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nforall a. Num a => a -> a -> a
-Int
1) Type
res_ty TCvSubst
subst' (TyVar
tv' forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) (CoreExpr
e forall b. Expr b -> Expr b -> Expr b
`App` forall b. TyVar -> Expr b
varToCoreExpr TyVar
tv')
| Just (Type
mult, Type
arg_ty, Type
res_ty) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
, let (TCvSubst
subst', TyVar
b) = Int -> TCvSubst -> Scaled Type -> (TCvSubst, TyVar)
freshEtaId Int
n TCvSubst
subst (forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
= Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nforall a. Num a => a -> a -> a
-Int
1) Type
res_ty TCvSubst
subst' (TyVar
b forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) (CoreExpr
e forall b. Expr b -> Expr b -> Expr b
`App` forall b. TyVar -> Expr b
Var TyVar
b)
| Bool
otherwise
= forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaBodyForJoinPoint" forall a b. (a -> b) -> a -> b
$ Int -> SDoc
int Int
need_args SDoc -> SDoc -> SDoc
$$
forall a. Outputable a => a -> SDoc
ppr CoreExpr
body SDoc -> SDoc -> 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 :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, TyVar)
freshEtaId Int
n TCvSubst
subst Scaled Type
ty
= (TCvSubst
subst', TyVar
eta_id')
where
Scaled Type
mult' Type
ty' = HasCallStack => TCvSubst -> Scaled Type -> Scaled Type
Type.substScaledTyUnchecked TCvSubst
subst Scaled Type
ty
eta_id' :: TyVar
eta_id' = InScopeSet -> TyVar -> TyVar
uniqAway (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst) forall a b. (a -> b) -> a -> b
$
FastString -> Unique -> Type -> Type -> TyVar
mkSysLocalOrCoVar (String -> FastString
fsLit String
"eta") (Int -> Unique
mkBuiltinUnique Int
n) Type
mult' Type
ty'
subst' :: TCvSubst
subst' = TCvSubst -> TyVar -> TCvSubst
extendTCvInScope TCvSubst
subst TyVar
eta_id'