module Exitify ( exitifyProgram ) where
import GhcPrelude
import Var
import Id
import IdInfo
import CoreSyn
import CoreUtils
import State
import Unique
import VarSet
import VarEnv
import CoreFVs
import FastString
import Type
import Util( mapSnd )
import Data.Bifunctor
import Control.Monad
exitifyProgram :: CoreProgram -> CoreProgram
exitifyProgram :: CoreProgram -> CoreProgram
exitifyProgram CoreProgram
binds = (Bind CoreBndr -> Bind CoreBndr) -> CoreProgram -> CoreProgram
forall a b. (a -> b) -> [a] -> [b]
map Bind CoreBndr -> Bind CoreBndr
goTopLvl CoreProgram
binds
where
goTopLvl :: Bind CoreBndr -> Bind CoreBndr
goTopLvl (NonRec CoreBndr
v Expr CoreBndr
e) = CoreBndr -> Expr CoreBndr -> Bind CoreBndr
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
v (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope_toplvl Expr CoreBndr
e)
goTopLvl (Rec [(CoreBndr, Expr CoreBndr)]
pairs) = [(CoreBndr, Expr CoreBndr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec (((CoreBndr, Expr CoreBndr) -> (CoreBndr, Expr CoreBndr))
-> [(CoreBndr, Expr CoreBndr)] -> [(CoreBndr, Expr CoreBndr)]
forall a b. (a -> b) -> [a] -> [b]
map ((Expr CoreBndr -> Expr CoreBndr)
-> (CoreBndr, Expr CoreBndr) -> (CoreBndr, Expr CoreBndr)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope_toplvl)) [(CoreBndr, Expr CoreBndr)]
pairs)
in_scope_toplvl :: InScopeSet
in_scope_toplvl = InScopeSet
emptyInScopeSet InScopeSet -> [CoreBndr] -> InScopeSet
`extendInScopeSetList` CoreProgram -> [CoreBndr]
forall b. [Bind b] -> [b]
bindersOfBinds CoreProgram
binds
go :: InScopeSet -> CoreExpr -> CoreExpr
go :: InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
_ e :: Expr CoreBndr
e@(Var{}) = Expr CoreBndr
e
go InScopeSet
_ e :: Expr CoreBndr
e@(Lit {}) = Expr CoreBndr
e
go InScopeSet
_ e :: Expr CoreBndr
e@(Type {}) = Expr CoreBndr
e
go InScopeSet
_ e :: Expr CoreBndr
e@(Coercion {}) = Expr CoreBndr
e
go InScopeSet
in_scope (Cast Expr CoreBndr
e' Coercion
c) = Expr CoreBndr -> Coercion -> Expr CoreBndr
forall b. Expr b -> Coercion -> Expr b
Cast (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
e') Coercion
c
go InScopeSet
in_scope (Tick Tickish CoreBndr
t Expr CoreBndr
e') = Tickish CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick Tickish CoreBndr
t (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
e')
go InScopeSet
in_scope (App Expr CoreBndr
e1 Expr CoreBndr
e2) = Expr CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Expr b -> Expr b -> Expr b
App (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
e1) (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
e2)
go InScopeSet
in_scope (Lam CoreBndr
v Expr CoreBndr
e')
= CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
v (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope' Expr CoreBndr
e')
where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> CoreBndr -> InScopeSet
`extendInScopeSet` CoreBndr
v
go InScopeSet
in_scope (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 (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
scrut) CoreBndr
bndr Type
ty ((Alt CoreBndr -> Alt CoreBndr) -> [Alt CoreBndr] -> [Alt CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map Alt CoreBndr -> Alt CoreBndr
forall a.
(a, [CoreBndr], Expr CoreBndr) -> (a, [CoreBndr], Expr CoreBndr)
go_alt [Alt CoreBndr]
alts)
where
in_scope1 :: InScopeSet
in_scope1 = InScopeSet
in_scope InScopeSet -> CoreBndr -> InScopeSet
`extendInScopeSet` CoreBndr
bndr
go_alt :: (a, [CoreBndr], Expr CoreBndr) -> (a, [CoreBndr], Expr CoreBndr)
go_alt (a
dc, [CoreBndr]
pats, Expr CoreBndr
rhs) = (a
dc, [CoreBndr]
pats, InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope' Expr CoreBndr
rhs)
where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope1 InScopeSet -> [CoreBndr] -> InScopeSet
`extendInScopeSetList` [CoreBndr]
pats
go InScopeSet
in_scope (Let (NonRec CoreBndr
bndr Expr CoreBndr
rhs) Expr CoreBndr
body)
= Bind CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let (CoreBndr -> Expr CoreBndr -> Bind CoreBndr
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
bndr (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope Expr CoreBndr
rhs)) (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope' Expr CoreBndr
body)
where
in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> CoreBndr -> InScopeSet
`extendInScopeSet` CoreBndr
bndr
go InScopeSet
in_scope (Let (Rec [(CoreBndr, Expr CoreBndr)]
pairs) Expr CoreBndr
body)
| Bool
is_join_rec = CoreProgram -> Expr CoreBndr -> Expr CoreBndr
forall b. [Bind b] -> Expr b -> Expr b
mkLets (InScopeSet -> [(CoreBndr, Expr CoreBndr)] -> CoreProgram
exitifyRec InScopeSet
in_scope' [(CoreBndr, Expr CoreBndr)]
pairs') Expr CoreBndr
body'
| Bool
otherwise = Bind CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let ([(CoreBndr, Expr CoreBndr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs') Expr CoreBndr
body'
where
is_join_rec :: Bool
is_join_rec = ((CoreBndr, Expr CoreBndr) -> Bool)
-> [(CoreBndr, Expr CoreBndr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (CoreBndr -> Bool
isJoinId (CoreBndr -> Bool)
-> ((CoreBndr, Expr CoreBndr) -> CoreBndr)
-> (CoreBndr, Expr CoreBndr)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CoreBndr, Expr CoreBndr) -> CoreBndr
forall a b. (a, b) -> a
fst) [(CoreBndr, Expr CoreBndr)]
pairs
in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> [CoreBndr] -> InScopeSet
`extendInScopeSetList` Bind CoreBndr -> [CoreBndr]
forall b. Bind b -> [b]
bindersOf ([(CoreBndr, Expr CoreBndr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs)
pairs' :: [(CoreBndr, Expr CoreBndr)]
pairs' = (Expr CoreBndr -> Expr CoreBndr)
-> [(CoreBndr, Expr CoreBndr)] -> [(CoreBndr, Expr CoreBndr)]
forall b c a. (b -> c) -> [(a, b)] -> [(a, c)]
mapSnd (InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope') [(CoreBndr, Expr CoreBndr)]
pairs
body' :: Expr CoreBndr
body' = InScopeSet -> Expr CoreBndr -> Expr CoreBndr
go InScopeSet
in_scope' Expr CoreBndr
body
type ExitifyM = State [(JoinId, CoreExpr)]
exitifyRec :: InScopeSet -> [(Var,CoreExpr)] -> [CoreBind]
exitifyRec :: InScopeSet -> [(CoreBndr, Expr CoreBndr)] -> CoreProgram
exitifyRec InScopeSet
in_scope [(CoreBndr, Expr CoreBndr)]
pairs
= [ CoreBndr -> Expr CoreBndr -> Bind CoreBndr
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
xid Expr CoreBndr
rhs | (CoreBndr
xid,Expr CoreBndr
rhs) <- [(CoreBndr, Expr CoreBndr)]
exits ] CoreProgram -> CoreProgram -> CoreProgram
forall a. [a] -> [a] -> [a]
++ [[(CoreBndr, Expr CoreBndr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs']
where
ann_pairs :: [(CoreBndr, CoreExprWithFVs)]
ann_pairs = ((CoreBndr, Expr CoreBndr) -> (CoreBndr, CoreExprWithFVs))
-> [(CoreBndr, Expr CoreBndr)] -> [(CoreBndr, CoreExprWithFVs)]
forall a b. (a -> b) -> [a] -> [b]
map ((Expr CoreBndr -> CoreExprWithFVs)
-> (CoreBndr, Expr CoreBndr) -> (CoreBndr, CoreExprWithFVs)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Expr CoreBndr -> CoreExprWithFVs
freeVars) [(CoreBndr, Expr CoreBndr)]
pairs
recursive_calls :: VarSet
recursive_calls = [CoreBndr] -> VarSet
mkVarSet ([CoreBndr] -> VarSet) -> [CoreBndr] -> VarSet
forall a b. (a -> b) -> a -> b
$ ((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
([(CoreBndr, Expr CoreBndr)]
pairs',[(CoreBndr, Expr CoreBndr)]
exits) = (State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
-> [(CoreBndr, Expr CoreBndr)]
-> ([(CoreBndr, Expr CoreBndr)], [(CoreBndr, Expr CoreBndr)])
forall s a. State s a -> s -> (a, s)
`runState` []) (State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
-> ([(CoreBndr, Expr CoreBndr)], [(CoreBndr, Expr CoreBndr)]))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
-> ([(CoreBndr, Expr CoreBndr)], [(CoreBndr, Expr CoreBndr)])
forall a b. (a -> b) -> a -> b
$ do
[(CoreBndr, CoreExprWithFVs)]
-> ((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(CoreBndr, CoreExprWithFVs)]
ann_pairs (((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)])
-> ((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall a b. (a -> b) -> a -> b
$ \(CoreBndr
x,CoreExprWithFVs
rhs) -> do
let ([CoreBndr]
args, CoreExprWithFVs
body) = Int -> CoreExprWithFVs -> ([CoreBndr], CoreExprWithFVs)
forall bndr annot.
Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectNAnnBndrs (CoreBndr -> Int
idJoinArity CoreBndr
x) CoreExprWithFVs
rhs
Expr CoreBndr
body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go [CoreBndr]
args CoreExprWithFVs
body
let rhs' :: Expr CoreBndr
rhs' = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
args Expr CoreBndr
body'
(CoreBndr, Expr CoreBndr)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr
x, Expr CoreBndr
rhs')
go :: [Var]
-> CoreExprWithFVs
-> ExitifyM CoreExpr
go :: [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go [CoreBndr]
captured CoreExprWithFVs
ann_e
|
let fvs :: VarSet
fvs = DVarSet -> VarSet
dVarSetToVarSet (CoreExprWithFVs -> DVarSet
freeVarsOf CoreExprWithFVs
ann_e)
, VarSet -> VarSet -> Bool
disjointVarSet VarSet
fvs VarSet
recursive_calls
= [CoreBndr] -> Expr CoreBndr -> VarSet -> ExitifyM (Expr CoreBndr)
go_exit [CoreBndr]
captured (CoreExprWithFVs -> Expr CoreBndr
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate CoreExprWithFVs
ann_e) VarSet
fvs
go [CoreBndr]
captured (DVarSet
_, AnnCase CoreExprWithFVs
scrut CoreBndr
bndr Type
ty [AnnAlt CoreBndr DVarSet]
alts) = do
[Alt CoreBndr]
alts' <- [AnnAlt CoreBndr DVarSet]
-> (AnnAlt CoreBndr DVarSet
-> State [(CoreBndr, Expr CoreBndr)] (Alt CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [Alt CoreBndr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [AnnAlt CoreBndr DVarSet]
alts ((AnnAlt CoreBndr DVarSet
-> State [(CoreBndr, Expr CoreBndr)] (Alt CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [Alt CoreBndr])
-> (AnnAlt CoreBndr DVarSet
-> State [(CoreBndr, Expr CoreBndr)] (Alt CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [Alt CoreBndr]
forall a b. (a -> b) -> a -> b
$ \(AltCon
dc, [CoreBndr]
pats, CoreExprWithFVs
rhs) -> do
Expr CoreBndr
rhs' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr
bndr] [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
pats) CoreExprWithFVs
rhs
Alt CoreBndr -> State [(CoreBndr, Expr CoreBndr)] (Alt CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (AltCon
dc, [CoreBndr]
pats, Expr CoreBndr
rhs')
Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> ExitifyM (Expr CoreBndr))
-> Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall a b. (a -> b) -> a -> b
$ Expr CoreBndr
-> CoreBndr -> Type -> [Alt CoreBndr] -> Expr CoreBndr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (CoreExprWithFVs -> Expr CoreBndr
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate CoreExprWithFVs
scrut) CoreBndr
bndr Type
ty [Alt CoreBndr]
alts'
go [CoreBndr]
captured (DVarSet
_, AnnLet AnnBind CoreBndr DVarSet
ann_bind CoreExprWithFVs
body)
| AnnNonRec CoreBndr
j CoreExprWithFVs
rhs <- AnnBind CoreBndr DVarSet
ann_bind
, Just Int
join_arity <- CoreBndr -> Maybe Int
isJoinId_maybe CoreBndr
j
= do let ([CoreBndr]
params, CoreExprWithFVs
join_body) = Int -> CoreExprWithFVs -> ([CoreBndr], CoreExprWithFVs)
forall bndr annot.
Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectNAnnBndrs Int
join_arity CoreExprWithFVs
rhs
Expr CoreBndr
join_body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
params) CoreExprWithFVs
join_body
let rhs' :: Expr CoreBndr
rhs' = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
params Expr CoreBndr
join_body'
Expr CoreBndr
body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr
j]) CoreExprWithFVs
body
Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> ExitifyM (Expr CoreBndr))
-> Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall a b. (a -> b) -> a -> b
$ Bind CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let (CoreBndr -> Expr CoreBndr -> Bind CoreBndr
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
j Expr CoreBndr
rhs') Expr CoreBndr
body'
| AnnRec [(CoreBndr, CoreExprWithFVs)]
pairs <- AnnBind CoreBndr DVarSet
ann_bind
, CoreBndr -> Bool
isJoinId ((CoreBndr, CoreExprWithFVs) -> CoreBndr
forall a b. (a, b) -> a
fst ([(CoreBndr, CoreExprWithFVs)] -> (CoreBndr, CoreExprWithFVs)
forall a. [a] -> a
head [(CoreBndr, CoreExprWithFVs)]
pairs))
= do let js :: [CoreBndr]
js = ((CoreBndr, CoreExprWithFVs) -> CoreBndr)
-> [(CoreBndr, CoreExprWithFVs)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, CoreExprWithFVs) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, CoreExprWithFVs)]
pairs
[(CoreBndr, Expr CoreBndr)]
pairs' <- [(CoreBndr, CoreExprWithFVs)]
-> ((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(CoreBndr, CoreExprWithFVs)]
pairs (((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)])
-> ((CoreBndr, CoreExprWithFVs)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr))
-> State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall a b. (a -> b) -> a -> b
$ \(CoreBndr
j,CoreExprWithFVs
rhs) -> do
let join_arity :: Int
join_arity = CoreBndr -> Int
idJoinArity CoreBndr
j
([CoreBndr]
params, CoreExprWithFVs
join_body) = Int -> CoreExprWithFVs -> ([CoreBndr], CoreExprWithFVs)
forall bndr annot.
Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectNAnnBndrs Int
join_arity CoreExprWithFVs
rhs
Expr CoreBndr
join_body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
js [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
params) CoreExprWithFVs
join_body
let rhs' :: Expr CoreBndr
rhs' = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
params Expr CoreBndr
join_body'
(CoreBndr, Expr CoreBndr)
-> State [(CoreBndr, Expr CoreBndr)] (CoreBndr, Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr
j, Expr CoreBndr
rhs')
Expr CoreBndr
body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ [CoreBndr]
js) CoreExprWithFVs
body
Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> ExitifyM (Expr CoreBndr))
-> Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall a b. (a -> b) -> a -> b
$ Bind CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let ([(CoreBndr, Expr CoreBndr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs') Expr CoreBndr
body'
| Bool
otherwise
= do Expr CoreBndr
body' <- [CoreBndr] -> CoreExprWithFVs -> ExitifyM (Expr CoreBndr)
go ([CoreBndr]
captured [CoreBndr] -> [CoreBndr] -> [CoreBndr]
forall a. [a] -> [a] -> [a]
++ Bind CoreBndr -> [CoreBndr]
forall b. Bind b -> [b]
bindersOf Bind CoreBndr
bind ) CoreExprWithFVs
body
Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> ExitifyM (Expr CoreBndr))
-> Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall a b. (a -> b) -> a -> b
$ Bind CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let Bind CoreBndr
bind Expr CoreBndr
body'
where bind :: Bind CoreBndr
bind = AnnBind CoreBndr DVarSet -> Bind CoreBndr
forall b annot. AnnBind b annot -> Bind b
deAnnBind AnnBind CoreBndr DVarSet
ann_bind
go [CoreBndr]
_ CoreExprWithFVs
ann_e = Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExprWithFVs -> Expr CoreBndr
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate CoreExprWithFVs
ann_e)
go_exit :: [Var]
-> CoreExpr
-> VarSet
-> ExitifyM CoreExpr
go_exit :: [CoreBndr] -> Expr CoreBndr -> VarSet -> ExitifyM (Expr CoreBndr)
go_exit [CoreBndr]
captured Expr CoreBndr
e VarSet
fvs
| (Var CoreBndr
f, [Expr CoreBndr]
args) <- Expr CoreBndr -> (Expr CoreBndr, [Expr CoreBndr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr CoreBndr
e
, CoreBndr -> Bool
isJoinId CoreBndr
f
, (Expr CoreBndr -> Bool) -> [Expr CoreBndr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Expr CoreBndr -> Bool
forall b. Expr b -> Bool
isCapturedVarArg [Expr CoreBndr]
args
= Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return Expr CoreBndr
e
| Bool -> Bool
not Bool
is_interesting
= Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return Expr CoreBndr
e
| Bool
captures_join_points
= Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return Expr CoreBndr
e
| Bool
otherwise
= do {
let rhs :: Expr CoreBndr
rhs = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
abs_vars Expr CoreBndr
e
avoid :: InScopeSet
avoid = InScopeSet
in_scope InScopeSet -> [CoreBndr] -> InScopeSet
`extendInScopeSetList` [CoreBndr]
captured
; CoreBndr
v <- InScopeSet -> Int -> Expr CoreBndr -> ExitifyM CoreBndr
addExit InScopeSet
avoid ([CoreBndr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [CoreBndr]
abs_vars) Expr CoreBndr
rhs
; Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> ExitifyM (Expr CoreBndr))
-> Expr CoreBndr -> ExitifyM (Expr CoreBndr)
forall a b. (a -> b) -> a -> b
$ Expr CoreBndr -> [CoreBndr] -> Expr CoreBndr
forall b. Expr b -> [CoreBndr] -> Expr b
mkVarApps (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
v) [CoreBndr]
abs_vars }
where
isCapturedVarArg :: Expr b -> Bool
isCapturedVarArg (Var CoreBndr
v) = CoreBndr
v CoreBndr -> [CoreBndr] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [CoreBndr]
captured
isCapturedVarArg Expr b
_ = Bool
False
is_interesting :: Bool
is_interesting = (CoreBndr -> Bool) -> VarSet -> Bool
anyVarSet CoreBndr -> Bool
isLocalId (VarSet -> Bool) -> VarSet -> Bool
forall a b. (a -> b) -> a -> b
$
VarSet
fvs VarSet -> VarSet -> VarSet
`minusVarSet` [CoreBndr] -> VarSet
mkVarSet [CoreBndr]
captured
abs_vars :: [CoreBndr]
abs_vars = (VarSet, [CoreBndr]) -> [CoreBndr]
forall a b. (a, b) -> b
snd ((VarSet, [CoreBndr]) -> [CoreBndr])
-> (VarSet, [CoreBndr]) -> [CoreBndr]
forall a b. (a -> b) -> a -> b
$ (CoreBndr -> (VarSet, [CoreBndr]) -> (VarSet, [CoreBndr]))
-> (VarSet, [CoreBndr]) -> [CoreBndr] -> (VarSet, [CoreBndr])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CoreBndr -> (VarSet, [CoreBndr]) -> (VarSet, [CoreBndr])
pick (VarSet
fvs, []) [CoreBndr]
captured
where
pick :: CoreBndr -> (VarSet, [CoreBndr]) -> (VarSet, [CoreBndr])
pick CoreBndr
v (VarSet
fvs', [CoreBndr]
acc) | CoreBndr
v CoreBndr -> VarSet -> Bool
`elemVarSet` VarSet
fvs' = (VarSet
fvs' VarSet -> CoreBndr -> VarSet
`delVarSet` CoreBndr
v, CoreBndr -> CoreBndr
zap CoreBndr
v CoreBndr -> [CoreBndr] -> [CoreBndr]
forall a. a -> [a] -> [a]
: [CoreBndr]
acc)
| Bool
otherwise = (VarSet
fvs', [CoreBndr]
acc)
zap :: CoreBndr -> CoreBndr
zap CoreBndr
v | CoreBndr -> Bool
isId CoreBndr
v = CoreBndr -> IdInfo -> CoreBndr
setIdInfo CoreBndr
v IdInfo
vanillaIdInfo
| Bool
otherwise = CoreBndr
v
captures_join_points :: Bool
captures_join_points = (CoreBndr -> Bool) -> [CoreBndr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any CoreBndr -> Bool
isJoinId [CoreBndr]
abs_vars
mkExitJoinId :: InScopeSet -> Type -> JoinArity -> ExitifyM JoinId
mkExitJoinId :: InScopeSet -> Type -> Int -> ExitifyM CoreBndr
mkExitJoinId InScopeSet
in_scope Type
ty Int
join_arity = do
[(CoreBndr, Expr CoreBndr)]
fs <- State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall s. State s s
get
let avoid :: InScopeSet
avoid = InScopeSet
in_scope InScopeSet -> [CoreBndr] -> InScopeSet
`extendInScopeSetList` (((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)]
fs)
InScopeSet -> CoreBndr -> InScopeSet
`extendInScopeSet` CoreBndr
exit_id_tmpl
CoreBndr -> ExitifyM CoreBndr
forall (m :: * -> *) a. Monad m => a -> m a
return (InScopeSet -> CoreBndr -> CoreBndr
uniqAway InScopeSet
avoid CoreBndr
exit_id_tmpl)
where
exit_id_tmpl :: CoreBndr
exit_id_tmpl = FastString -> Unique -> Type -> CoreBndr
mkSysLocal (String -> FastString
fsLit String
"exit") Unique
initExitJoinUnique Type
ty
CoreBndr -> Int -> CoreBndr
`asJoinId` Int
join_arity
addExit :: InScopeSet -> JoinArity -> CoreExpr -> ExitifyM JoinId
addExit :: InScopeSet -> Int -> Expr CoreBndr -> ExitifyM CoreBndr
addExit InScopeSet
in_scope Int
join_arity Expr CoreBndr
rhs = do
let ty :: Type
ty = Expr CoreBndr -> Type
exprType Expr CoreBndr
rhs
CoreBndr
v <- InScopeSet -> Type -> Int -> ExitifyM CoreBndr
mkExitJoinId InScopeSet
in_scope Type
ty Int
join_arity
[(CoreBndr, Expr CoreBndr)]
fs <- State [(CoreBndr, Expr CoreBndr)] [(CoreBndr, Expr CoreBndr)]
forall s. State s s
get
[(CoreBndr, Expr CoreBndr)] -> State [(CoreBndr, Expr CoreBndr)] ()
forall s. s -> State s ()
put ((CoreBndr
v,Expr CoreBndr
rhs)(CoreBndr, Expr CoreBndr)
-> [(CoreBndr, Expr CoreBndr)] -> [(CoreBndr, Expr CoreBndr)]
forall a. a -> [a] -> [a]
:[(CoreBndr, Expr CoreBndr)]
fs)
CoreBndr -> ExitifyM CoreBndr
forall (m :: * -> *) a. Monad m => a -> m a
return CoreBndr
v