{-# LANGUAGE CPP #-}
module GHC.Core.Opt.LiberateCase ( liberateCase ) where
#include "HsVersions.h"
import GHC.Prelude
import GHC.Driver.Session
import GHC.Core
import GHC.Core.Unfold ( couldBeSmallEnoughToInline )
import GHC.Builtin.Types ( unitDataConId )
import GHC.Types.Id
import GHC.Types.Var.Env
import GHC.Utils.Misc ( notNull )
liberateCase :: DynFlags -> CoreProgram -> CoreProgram
liberateCase :: DynFlags -> CoreProgram -> CoreProgram
liberateCase DynFlags
dflags CoreProgram
binds = LibCaseEnv -> CoreProgram -> CoreProgram
do_prog (DynFlags -> LibCaseEnv
initEnv DynFlags
dflags) CoreProgram
binds
where
do_prog :: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog LibCaseEnv
_ [] = []
do_prog LibCaseEnv
env (CoreBind
bind:CoreProgram
binds) = CoreBind
bind' CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog LibCaseEnv
env' CoreProgram
binds
where
(LibCaseEnv
env', CoreBind
bind') = LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env CoreBind
bind
libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env (NonRec CoreBndr
binder Expr CoreBndr
rhs)
= (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr
binder], CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
binder (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
rhs))
libCaseBind LibCaseEnv
env (Rec [(CoreBndr, Expr CoreBndr)]
pairs)
= (LibCaseEnv
env_body, [(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs')
where
binders :: [CoreBndr]
binders = ((CoreBndr, Expr CoreBndr) -> CoreBndr)
-> [(CoreBndr, Expr CoreBndr)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, Expr CoreBndr) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, Expr CoreBndr)]
pairs
env_body :: LibCaseEnv
env_body = LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr]
binders
pairs' :: [(CoreBndr, Expr CoreBndr)]
pairs' = [(CoreBndr
binder, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_rhs Expr CoreBndr
rhs) | (CoreBndr
binder,Expr CoreBndr
rhs) <- [(CoreBndr, Expr CoreBndr)]
pairs]
env_rhs :: LibCaseEnv
env_rhs | Bool
is_dupable_bind = LibCaseEnv -> [(CoreBndr, Expr CoreBndr)] -> LibCaseEnv
addRecBinds LibCaseEnv
env [(CoreBndr, Expr CoreBndr)]
dup_pairs
| Bool
otherwise = LibCaseEnv
env
dup_pairs :: [(CoreBndr, Expr CoreBndr)]
dup_pairs = [ (CoreBndr -> CoreBndr
localiseId CoreBndr
binder, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_body Expr CoreBndr
rhs)
| (CoreBndr
binder, Expr CoreBndr
rhs) <- [(CoreBndr, Expr CoreBndr)]
pairs ]
is_dupable_bind :: Bool
is_dupable_bind = Bool
small_enough Bool -> Bool -> Bool
&& ((CoreBndr, Expr CoreBndr) -> Bool)
-> [(CoreBndr, Expr CoreBndr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (CoreBndr, Expr CoreBndr) -> Bool
forall {b}. (CoreBndr, b) -> Bool
ok_pair [(CoreBndr, Expr CoreBndr)]
pairs
small_enough :: Bool
small_enough = case LibCaseEnv -> Maybe LibCaseLevel
bombOutSize LibCaseEnv
env of
Maybe LibCaseLevel
Nothing -> Bool
True
Just LibCaseLevel
size -> DynFlags -> LibCaseLevel -> Expr CoreBndr -> Bool
couldBeSmallEnoughToInline (LibCaseEnv -> DynFlags
lc_dflags LibCaseEnv
env) LibCaseLevel
size (Expr CoreBndr -> Bool) -> Expr CoreBndr -> Bool
forall a b. (a -> b) -> a -> b
$
CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
dup_pairs) (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
unitDataConId)
ok_pair :: (CoreBndr, b) -> Bool
ok_pair (CoreBndr
id,b
_)
= CoreBndr -> LibCaseLevel
idArity CoreBndr
id LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
> LibCaseLevel
0
Bool -> Bool -> Bool
&& Bool -> Bool
not (CoreBndr -> Bool
isDeadEndId CoreBndr
id)
libCase :: LibCaseEnv
-> CoreExpr
-> CoreExpr
libCase :: LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env (Var CoreBndr
v) = LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v []
libCase LibCaseEnv
_ (Lit Literal
lit) = Literal -> Expr CoreBndr
forall b. Literal -> Expr b
Lit Literal
lit
libCase LibCaseEnv
_ (Type Type
ty) = Type -> Expr CoreBndr
forall b. Type -> Expr b
Type Type
ty
libCase LibCaseEnv
_ (Coercion Coercion
co) = Coercion -> Expr CoreBndr
forall b. Coercion -> Expr b
Coercion Coercion
co
libCase LibCaseEnv
env e :: Expr CoreBndr
e@(App {}) | let (Expr CoreBndr
fun, [Expr CoreBndr]
args) = Expr CoreBndr -> (Expr CoreBndr, [Expr CoreBndr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr CoreBndr
e
, Var CoreBndr
v <- Expr CoreBndr
fun
= LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v [Expr CoreBndr]
args
libCase LibCaseEnv
env (App Expr CoreBndr
fun Expr CoreBndr
arg) = Expr CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Expr b -> Expr b -> Expr b
App (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
fun) (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
arg)
libCase LibCaseEnv
env (Tick Tickish CoreBndr
tickish Expr CoreBndr
body) = Tickish CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick Tickish CoreBndr
tickish (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
body)
libCase LibCaseEnv
env (Cast Expr CoreBndr
e Coercion
co) = Expr CoreBndr -> Coercion -> Expr CoreBndr
forall b. Expr b -> Coercion -> Expr b
Cast (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
e) Coercion
co
libCase LibCaseEnv
env (Lam CoreBndr
binder Expr CoreBndr
body)
= CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
binder (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr
binder]) Expr CoreBndr
body)
libCase LibCaseEnv
env (Let CoreBind
bind Expr CoreBndr
body)
= CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind' (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_body Expr CoreBndr
body)
where
(LibCaseEnv
env_body, CoreBind
bind') = LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env CoreBind
bind
libCase LibCaseEnv
env (Case Expr CoreBndr
scrut CoreBndr
bndr Type
ty [Alt CoreBndr]
alts)
= Expr CoreBndr
-> CoreBndr -> Type -> [Alt CoreBndr] -> Expr CoreBndr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
scrut) CoreBndr
bndr Type
ty ((Alt CoreBndr -> Alt CoreBndr) -> [Alt CoreBndr] -> [Alt CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
libCaseAlt LibCaseEnv
env_alts) [Alt CoreBndr]
alts)
where
env_alts :: LibCaseEnv
env_alts = LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders (Expr CoreBndr -> LibCaseEnv
forall {b}. Expr b -> LibCaseEnv
mk_alt_env Expr CoreBndr
scrut) [CoreBndr
bndr]
mk_alt_env :: Expr b -> LibCaseEnv
mk_alt_env (Var CoreBndr
scrut_var) = LibCaseEnv -> CoreBndr -> LibCaseEnv
addScrutedVar LibCaseEnv
env CoreBndr
scrut_var
mk_alt_env (Cast Expr b
scrut Coercion
_) = Expr b -> LibCaseEnv
mk_alt_env Expr b
scrut
mk_alt_env Expr b
_ = LibCaseEnv
env
libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
-> (AltCon, [CoreBndr], CoreExpr)
libCaseAlt :: LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
libCaseAlt LibCaseEnv
env (AltCon
con,[CoreBndr]
args,Expr CoreBndr
rhs) = (AltCon
con, [CoreBndr]
args, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr]
args) Expr CoreBndr
rhs)
libCaseApp :: LibCaseEnv -> Id -> [CoreExpr] -> CoreExpr
libCaseApp :: LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v [Expr CoreBndr]
args
| Just CoreBind
the_bind <- LibCaseEnv -> CoreBndr -> Maybe CoreBind
lookupRecId LibCaseEnv
env CoreBndr
v
, [CoreBndr] -> Bool
forall a. [a] -> Bool
notNull [CoreBndr]
free_scruts
= CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
the_bind Expr CoreBndr
expr'
| Bool
otherwise
= Expr CoreBndr
expr'
where
rec_id_level :: LibCaseLevel
rec_id_level = LibCaseEnv -> CoreBndr -> LibCaseLevel
lookupLevel LibCaseEnv
env CoreBndr
v
free_scruts :: [CoreBndr]
free_scruts = LibCaseEnv -> LibCaseLevel -> [CoreBndr]
freeScruts LibCaseEnv
env LibCaseLevel
rec_id_level
expr' :: Expr CoreBndr
expr' = Expr CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
v) ((Expr CoreBndr -> Expr CoreBndr)
-> [Expr CoreBndr] -> [Expr CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env) [Expr CoreBndr]
args)
freeScruts :: LibCaseEnv
-> LibCaseLevel
-> [Id]
freeScruts :: LibCaseEnv -> LibCaseLevel -> [CoreBndr]
freeScruts LibCaseEnv
env LibCaseLevel
rec_bind_lvl
= [CoreBndr
v | (CoreBndr
v, LibCaseLevel
scrut_bind_lvl, LibCaseLevel
scrut_at_lvl) <- LibCaseEnv -> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts LibCaseEnv
env
, LibCaseLevel
scrut_bind_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
<= LibCaseLevel
rec_bind_lvl
, LibCaseLevel
scrut_at_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
> LibCaseLevel
rec_bind_lvl]
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders env :: LibCaseEnv
env@(LibCaseEnv { lc_lvl :: LibCaseEnv -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env }) [CoreBndr]
binders
= LibCaseEnv
env { lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env' }
where
lvl_env' :: IdEnv LibCaseLevel
lvl_env' = IdEnv LibCaseLevel
-> [(CoreBndr, LibCaseLevel)] -> IdEnv LibCaseLevel
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv LibCaseLevel
lvl_env ([CoreBndr]
binders [CoreBndr] -> [LibCaseLevel] -> [(CoreBndr, LibCaseLevel)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` LibCaseLevel -> [LibCaseLevel]
forall a. a -> [a]
repeat LibCaseLevel
lvl)
addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
addRecBinds :: LibCaseEnv -> [(CoreBndr, Expr CoreBndr)] -> LibCaseEnv
addRecBinds env :: LibCaseEnv
env@(LibCaseEnv {lc_lvl :: LibCaseEnv -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env,
lc_rec_env :: LibCaseEnv -> IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env}) [(CoreBndr, Expr CoreBndr)]
pairs
= LibCaseEnv
env { lc_lvl :: LibCaseLevel
lc_lvl = LibCaseLevel
lvl', lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env', lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env' }
where
lvl' :: LibCaseLevel
lvl' = LibCaseLevel
lvl LibCaseLevel -> LibCaseLevel -> LibCaseLevel
forall a. Num a => a -> a -> a
+ LibCaseLevel
1
lvl_env' :: IdEnv LibCaseLevel
lvl_env' = IdEnv LibCaseLevel
-> [(CoreBndr, LibCaseLevel)] -> IdEnv LibCaseLevel
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv LibCaseLevel
lvl_env [(CoreBndr
binder,LibCaseLevel
lvl) | (CoreBndr
binder,Expr CoreBndr
_) <- [(CoreBndr, Expr CoreBndr)]
pairs]
rec_env' :: IdEnv CoreBind
rec_env' = IdEnv CoreBind -> [(CoreBndr, CoreBind)] -> IdEnv CoreBind
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv CoreBind
rec_env [(CoreBndr
binder, [(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs) | (CoreBndr
binder,Expr CoreBndr
_) <- [(CoreBndr, Expr CoreBndr)]
pairs]
addScrutedVar :: LibCaseEnv
-> Id
-> LibCaseEnv
addScrutedVar :: LibCaseEnv -> CoreBndr -> LibCaseEnv
addScrutedVar env :: LibCaseEnv
env@(LibCaseEnv { lc_lvl :: LibCaseEnv -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env,
lc_scruts :: LibCaseEnv -> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts = [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts }) CoreBndr
scrut_var
| LibCaseLevel
bind_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
< LibCaseLevel
lvl
= LibCaseEnv
env { lc_scruts :: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts = [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts' }
| Bool
otherwise = LibCaseEnv
env
where
scruts' :: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts' = (CoreBndr
scrut_var, LibCaseLevel
bind_lvl, LibCaseLevel
lvl) (CoreBndr, LibCaseLevel, LibCaseLevel)
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
forall a. a -> [a] -> [a]
: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts
bind_lvl :: LibCaseLevel
bind_lvl = case IdEnv LibCaseLevel -> CoreBndr -> Maybe LibCaseLevel
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv IdEnv LibCaseLevel
lvl_env CoreBndr
scrut_var of
Just LibCaseLevel
lvl -> LibCaseLevel
lvl
Maybe LibCaseLevel
Nothing -> LibCaseLevel
topLevel
lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
lookupRecId :: LibCaseEnv -> CoreBndr -> Maybe CoreBind
lookupRecId LibCaseEnv
env CoreBndr
id = IdEnv CoreBind -> CoreBndr -> Maybe CoreBind
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv (LibCaseEnv -> IdEnv CoreBind
lc_rec_env LibCaseEnv
env) CoreBndr
id
lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
lookupLevel :: LibCaseEnv -> CoreBndr -> LibCaseLevel
lookupLevel LibCaseEnv
env CoreBndr
id
= case IdEnv LibCaseLevel -> CoreBndr -> Maybe LibCaseLevel
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv (LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env LibCaseEnv
env) CoreBndr
id of
Just LibCaseLevel
lvl -> LibCaseLevel
lvl
Maybe LibCaseLevel
Nothing -> LibCaseLevel
topLevel
type LibCaseLevel = Int
topLevel :: LibCaseLevel
topLevel :: LibCaseLevel
topLevel = LibCaseLevel
0
data LibCaseEnv
= LibCaseEnv {
LibCaseEnv -> DynFlags
lc_dflags :: DynFlags,
LibCaseEnv -> LibCaseLevel
lc_lvl :: LibCaseLevel,
LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env :: IdEnv LibCaseLevel,
LibCaseEnv -> IdEnv CoreBind
lc_rec_env :: IdEnv CoreBind,
LibCaseEnv -> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
}
initEnv :: DynFlags -> LibCaseEnv
initEnv :: DynFlags -> LibCaseEnv
initEnv DynFlags
dflags
= LibCaseEnv :: DynFlags
-> LibCaseLevel
-> IdEnv LibCaseLevel
-> IdEnv CoreBind
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
-> LibCaseEnv
LibCaseEnv { lc_dflags :: DynFlags
lc_dflags = DynFlags
dflags,
lc_lvl :: LibCaseLevel
lc_lvl = LibCaseLevel
0,
lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
forall a. VarEnv a
emptyVarEnv,
lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
forall a. VarEnv a
emptyVarEnv,
lc_scruts :: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts = [] }
bombOutSize :: LibCaseEnv -> Maybe Int
bombOutSize :: LibCaseEnv -> Maybe LibCaseLevel
bombOutSize = DynFlags -> Maybe LibCaseLevel
liberateCaseThreshold (DynFlags -> Maybe LibCaseLevel)
-> (LibCaseEnv -> DynFlags) -> LibCaseEnv -> Maybe LibCaseLevel
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LibCaseEnv -> DynFlags
lc_dflags