{-# LANGUAGE DeriveFunctor, PatternGuards #-}
module IRTS.LangOpts(inlineAll) where
import Idris.Core.CaseTree
import Idris.Core.TT
import IRTS.Lang
import Control.Monad.State hiding (lift)
inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
inlineAll [(Name, LDecl)]
lds = let defs :: Ctxt LDecl
defs = forall a. [(Name, a)] -> Ctxt a -> Ctxt a
addAlist [(Name, LDecl)]
lds forall {k} {a}. Map k a
emptyContext in
forall a b. (a -> b) -> [a] -> [b]
map (\ (Name
n, LDecl
def) -> (Name
n, Ctxt LDecl -> LDecl -> LDecl
doInline Ctxt LDecl
defs LDecl
def)) [(Name, LDecl)]
lds
nextN :: State Int Name
nextN :: State Int Name
nextN = do Int
i <- forall s (m :: * -> *). MonadState s m => m s
get
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i forall a. Num a => a -> a -> a
+ Int
1)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Int -> String -> Name
sMN Int
i String
"in"
doInline :: LDefs -> LDecl -> LDecl
doInline :: Ctxt LDecl -> LDecl -> LDecl
doInline = Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' Int
1
doInline' :: Int -> LDefs -> LDecl -> LDecl
doInline' :: Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' Int
0 Ctxt LDecl
defs LDecl
d = LDecl
d
doInline' Int
i Ctxt LDecl
defs d :: LDecl
d@(LConstructor Name
_ Int
_ Int
_) = LDecl
d
doInline' Int
i Ctxt LDecl
defs (LFun [LOpt]
opts Name
topn [Name]
args LExp
exp)
= let inl :: LExp
inl = forall s a. State s a -> s -> a
evalState ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
initEnv [Name
topn] Ctxt LDecl
defs LExp
exp)
(forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
args)
res :: LExp
res = LExp -> LExp
eta forall a b. (a -> b) -> a -> b
$ forall {t}. (Eq t, Num t) => t -> LExp -> LExp
caseFloats Integer
10 LExp
inl in
case LExp
res of
LLam [Name]
args' LExp
body ->
Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' (Int
i forall a. Num a => a -> a -> a
- Int
1) Ctxt LDecl
defs forall a b. (a -> b) -> a -> b
$
[LOpt] -> Name -> [Name] -> LExp -> LDecl
LFun [LOpt]
opts Name
topn (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, Name)]
initNames forall a. [a] -> [a] -> [a]
++ [Name]
args') LExp
body
LExp
_ -> Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' (Int
i forall a. Num a => a -> a -> a
- Int
1) Ctxt LDecl
defs forall a b. (a -> b) -> a -> b
$
[LOpt] -> Name -> [Name] -> LExp -> LDecl
LFun [LOpt]
opts Name
topn (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, Name)]
initNames) LExp
res
where
caseFloats :: t -> LExp -> LExp
caseFloats t
0 LExp
tm = LExp
tm
caseFloats t
n LExp
tm
= let res :: LExp
res = LExp -> LExp
caseFloat LExp
tm in
if LExp
res forall a. Eq a => a -> a -> Bool
== LExp
tm
then LExp
res
else t -> LExp -> LExp
caseFloats (t
nforall a. Num a => a -> a -> a
-t
1) LExp
res
initNames :: [(Name, Name)]
initNames = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Name
n Int
i -> (Name
n, Name -> Int -> Name
newn Name
n Int
i)) [Name]
args [Int
0..]
initEnv :: [(Name, LExp)]
initEnv = forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
initNames
newn :: Name -> Int -> Name
newn (UN Text
n) Int
i = Int -> Text -> Name
MN Int
i Text
n
newn Name
_ Int
i = Int -> String -> Name
sMN Int
i String
"arg"
unload :: [LExp] -> LExp -> LExp
unload :: [LExp] -> LExp -> LExp
unload [] LExp
e = LExp
e
unload [LExp]
stk (LApp Bool
tc LExp
e [LExp]
args) = Bool -> LExp -> [LExp] -> LExp
LApp Bool
tc LExp
e ([LExp]
args forall a. [a] -> [a] -> [a]
++ [LExp]
stk)
unload [LExp]
stk LExp
e = Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
e [LExp]
stk
takeStk :: [(Name, LExp)] -> [Name] -> [LExp] ->
([(Name, LExp)], [Name], [LExp])
takeStk :: [(Name, LExp)]
-> [Name] -> [LExp] -> ([(Name, LExp)], [Name], [LExp])
takeStk [(Name, LExp)]
env (Name
a : [Name]
args) (LExp
v : [LExp]
stk) = [(Name, LExp)]
-> [Name] -> [LExp] -> ([(Name, LExp)], [Name], [LExp])
takeStk ((Name
a, LExp
v) forall a. a -> [a] -> [a]
: [(Name, LExp)]
env) [Name]
args [LExp]
stk
takeStk [(Name, LExp)]
env [Name]
args [LExp]
stk = ([(Name, LExp)]
env, [Name]
args, [LExp]
stk)
eval :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp -> State Int LExp
eval :: [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LLazyApp Name
n [LExp]
es)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> [LExp] -> LExp
LLazyApp Name
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es)
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LForce LExp
e)
= do LExp
e' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
case LExp
e' of
LLazyExp LExp
forced -> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
forced
LLazyApp Name
n [LExp]
es -> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False (Name -> LExp
LV Name
n) [LExp]
es)
LExp
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk (LExp -> LExp
LForce LExp
e'))
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LLazyExp LExp
e)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LExp -> LExp
LLazyExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LApp Bool
t (LV Name
n) [LExp
_, LExp
_, LExp
_, LExp
act, (LLam [Name
arg] LExp
k)])
| Name
n forall a. Eq a => a -> a -> Bool
== String -> Name
sUN String
"io_bind"
= do Name
w <- State Int Name
nextN
let env' :: [(Name, LExp)]
env' = (Name
w, Name -> LExp
LV Name
w) forall a. a -> [a] -> [a]
: [(Name, LExp)]
env
LExp
act' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env' [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
act [Name -> LExp
LV Name
w])
Name
argn <- State Int Name
nextN
LExp
k' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] ((Name
arg, Name -> LExp
LV Name
argn) forall a. a -> [a] -> [a]
: [(Name, LExp)]
env') [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
k [Name -> LExp
LV Name
w])
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Name] -> LExp -> LExp
LLam [Name
w] (Name -> LExp -> LExp -> LExp
LLet Name
argn LExp
act' LExp
k')
eval (LExp
world : [LExp]
stk) [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LApp Bool
t (LV Name
n) [LExp
_, LExp
_, LExp
_, LExp
act, (LLam [Name
arg] LExp
k)])
| Name
n forall a. Eq a => a -> a -> Bool
== String -> Name
sUN String
"io_bind"
= do LExp
act' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
act [LExp
world])
Name
argn <- State Int Name
nextN
LExp
k' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk ((Name
arg, Name -> LExp
LV Name
argn) forall a. a -> [a] -> [a]
: [(Name, LExp)]
env) [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
k [LExp
world])
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Name -> LExp -> LExp -> LExp
LLet Name
argn LExp
act' LExp
k'
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LApp Bool
t LExp
f [LExp]
es)
= do [LExp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es
[LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval ([LExp]
es' forall a. [a] -> [a] -> [a]
++ [LExp]
stk) [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
f
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LLet Name
n LExp
val LExp
sc)
= do Name
n' <- State Int Name
nextN
Name -> LExp -> LExp -> LExp
LLet Name
n' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
val
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk ((Name
n, Name -> LExp
LV Name
n') forall a. a -> [a] -> [a]
: [(Name, LExp)]
env) [Name]
rec Ctxt LDecl
defs LExp
sc
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LProj LExp
exp Int
i)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (LExp -> Int -> LExp
LProj forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
exp forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Monad m => a -> m a
return Int
i)
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LCon Maybe Name
loc Int
i Name
n [LExp]
es)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon Maybe Name
loc Int
i Name
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es)
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LCase CaseType
ty LExp
e [])
= forall (f :: * -> *) a. Applicative f => a -> f a
pure LExp
LNothing
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LCase CaseType
ty LExp
e [LAlt]
alts)
= do LExp
e' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
case forall {b}. LExp -> [LAlt' b] -> Maybe ([(Name, LExp)], b)
evalAlts LExp
e' [LAlt]
alts of
Just ([(Name, LExp)]
env', LExp
tm) -> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env' [Name]
rec Ctxt LDecl
defs LExp
tm
Maybe ([(Name, LExp)], LExp)
Nothing ->
do [LAlt]
alts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([LExp]
-> [(Name, LExp)]
-> [Name]
-> Ctxt LDecl
-> LAlt
-> StateT Int Identity LAlt
evalAlt [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LAlt]
alts
let prefix :: [Name]
prefix = [LExp] -> [Name]
getLams (forall a b. (a -> b) -> [a] -> [b]
map forall {e}. LAlt' e -> e
getRHS [LAlt]
alts')
case [Name]
prefix of
[] -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ty LExp
e' (LExp -> [LAlt] -> [LAlt]
replaceInAlts LExp
e' [LAlt]
alts')
[Name]
args -> do [LAlt]
alts_red <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Name] -> LAlt -> StateT Int Identity LAlt
dropArgs [Name]
args) [LAlt]
alts'
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Name] -> LExp -> LExp
LLam [Name]
args
(CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ty LExp
e' (LExp -> [LAlt] -> [LAlt]
replaceInAlts LExp
e' [LAlt]
alts_red))
where
evalAlts :: LExp -> [LAlt' b] -> Maybe ([(Name, LExp)], b)
evalAlts LExp
e' [] = forall a. Maybe a
Nothing
evalAlts (LCon Maybe Name
_ Int
t Name
n [LExp]
args) (LConCase Int
i Name
n' [Name]
es b
rhs : [LAlt' b]
as)
| Name
n forall a. Eq a => a -> a -> Bool
== Name
n' = forall a. a -> Maybe a
Just (forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
es [LExp]
args forall a. [a] -> [a] -> [a]
++ [(Name, LExp)]
env, b
rhs)
evalAlts (LConst Const
c) (LConstCase Const
c' b
rhs : [LAlt' b]
as)
| Const
c forall a. Eq a => a -> a -> Bool
== Const
c' = forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
evalAlts (LCon Maybe Name
_ Int
_ Name
_ [LExp]
_) (LDefaultCase b
rhs : [LAlt' b]
as) = forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
evalAlts (LConst Const
_) (LDefaultCase b
rhs : [LAlt' b]
as) = forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
evalAlts LExp
tm (LAlt' b
_ : [LAlt' b]
as) = LExp -> [LAlt' b] -> Maybe ([(Name, LExp)], b)
evalAlts LExp
tm [LAlt' b]
as
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LOp PrimFn
f [LExp]
es)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimFn -> [LExp] -> LExp
LOp PrimFn
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LForeign FDesc
t FDesc
s [(FDesc, LExp)]
args)
= [LExp] -> LExp -> LExp
unload [LExp]
stk forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
FDesc -> FDesc -> [(FDesc, LExp)] -> LExp
LForeign FDesc
t FDesc
s forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(FDesc
t, LExp
e) -> do LExp
e' <- [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
forall (m :: * -> *) a. Monad m => a -> m a
return (FDesc
t, LExp
e')) [(FDesc, LExp)]
args
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LLam [Name]
args LExp
sc)
| ([(Name, LExp)]
env', [Name]
args', [LExp]
stk') <- [(Name, LExp)]
-> [Name] -> [LExp] -> ([(Name, LExp)], [Name], [LExp])
takeStk [(Name, LExp)]
env [Name]
args [LExp]
stk
= case [Name]
args' of
[] -> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk' [(Name, LExp)]
env' [Name]
rec Ctxt LDecl
defs LExp
sc
[Name]
as -> do [(Name, Name)]
ns' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Name
n -> do Name
n' <- State Int Name
nextN
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n, Name
n')) [Name]
args'
[Name] -> LExp -> LExp
LLam (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, Name)]
ns') forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
[LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk' (forall a b. (a -> b) -> [a] -> [b]
map (\ (Name
n, Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
ns' forall a. [a] -> [a] -> [a]
++ [(Name, LExp)]
env')
[Name]
rec Ctxt LDecl
defs LExp
sc
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs var :: LExp
var@(LV Name
n)
= case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
n [(Name, LExp)]
env of
Just LExp
t
| LExp
t forall a. Eq a => a -> a -> Bool
/= Name -> LExp
LV Name
n Bool -> Bool -> Bool
&& Name
n forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
rec ->
[LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env (Name
n forall a. a -> [a] -> [a]
: [Name]
rec) Ctxt LDecl
defs LExp
t
| Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
t)
Maybe LExp
Nothing
| Name
n forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
rec,
Just (LFun [LOpt]
opts Name
_ [Name]
args LExp
body) <- forall a. Name -> Ctxt a -> Maybe a
lookupCtxtExact Name
n Ctxt LDecl
defs,
LOpt
Inline forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [LOpt]
opts ->
[LExp]
-> [(Name, LExp)]
-> [Name]
-> Ctxt LDecl
-> LExp
-> [Name]
-> LExp
-> State Int LExp
apply [LExp]
stk [(Name, LExp)]
env (Name
n forall a. a -> [a] -> [a]
: [Name]
rec) Ctxt LDecl
defs LExp
var [Name]
args LExp
body
| Just (LConstructor Name
n Int
t Int
a) <- forall a. Name -> Ctxt a -> Maybe a
lookupCtxtExact Name
n Ctxt LDecl
defs ->
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon forall a. Maybe a
Nothing Int
t Name
n [LExp]
stk)
| Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
var)
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
t = forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
t)
evalAlt :: [LExp]
-> [(Name, LExp)]
-> [Name]
-> Ctxt LDecl
-> LAlt
-> StateT Int Identity LAlt
evalAlt [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LConCase Int
i Name
n [Name]
es LExp
rhs)
= do [(Name, Name)]
ns' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Name
n -> do Name
n' <- State Int Name
nextN
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n, Name
n')) [Name]
es
forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
n (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, Name)]
ns') forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
[LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk (forall a b. (a -> b) -> [a] -> [b]
map (\ (Name
n, Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
ns' forall a. [a] -> [a] -> [a]
++ [(Name, LExp)]
env) [Name]
rec Ctxt LDecl
defs LExp
rhs
evalAlt [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LConstCase Const
c LExp
e)
= forall e. Const -> e -> LAlt' e
LConstCase Const
c forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
evalAlt [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs (LDefaultCase LExp
e)
= forall e. e -> LAlt' e
LDefaultCase forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
e
apply :: [LExp] -> [(Name, LExp)] -> [Name] -> LDefs -> LExp ->
[Name] -> LExp -> State Int LExp
apply :: [LExp]
-> [(Name, LExp)]
-> [Name]
-> Ctxt LDecl
-> LExp
-> [Name]
-> LExp
-> State Int LExp
apply [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs LExp
var [Name]
args LExp
body
= [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [LExp]
stk [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs ([Name] -> LExp -> LExp
LLam [Name]
args LExp
body)
dropArgs :: [Name] -> LAlt -> State Int LAlt
dropArgs :: [Name] -> LAlt -> StateT Int Identity LAlt
dropArgs [Name]
as (LConCase Int
i Name
n [Name]
es LExp
t)
= do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
forall (m :: * -> *) a. Monad m => a -> m a
return (forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
n [Name]
es LExp
rhs')
dropArgs [Name]
as (LConstCase Const
c LExp
t)
= do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
forall (m :: * -> *) a. Monad m => a -> m a
return (forall e. Const -> e -> LAlt' e
LConstCase Const
c LExp
rhs')
dropArgs [Name]
as (LDefaultCase LExp
t)
= do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
forall (m :: * -> *) a. Monad m => a -> m a
return (forall e. e -> LAlt' e
LDefaultCase LExp
rhs')
dropArgsTm :: [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as (LLam [Name]
args LExp
rhs)
= do let old :: [Name]
old = forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
as) [Name]
args
[LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\ Name
o Name
n -> (Name
o, Name -> LExp
LV Name
n)) [Name]
old [Name]
as) [] forall {k} {a}. Map k a
emptyContext LExp
rhs
dropArgsTm [Name]
as (LLet Name
n LExp
val LExp
rhs)
= do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
rhs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> LExp -> LExp -> LExp
LLet Name
n LExp
val LExp
rhs')
dropArgsTm [Name]
as LExp
tm = forall (m :: * -> *) a. Monad m => a -> m a
return LExp
tm
caseFloat :: LExp -> LExp
caseFloat :: LExp -> LExp
caseFloat (LApp Bool
tc LExp
e [LExp]
es) = Bool -> LExp -> [LExp] -> LExp
LApp Bool
tc (LExp -> LExp
caseFloat LExp
e) (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
caseFloat [LExp]
es)
caseFloat (LLazyExp LExp
e) = LExp -> LExp
LLazyExp (LExp -> LExp
caseFloat LExp
e)
caseFloat (LForce LExp
e) = LExp -> LExp
LForce (LExp -> LExp
caseFloat LExp
e)
caseFloat (LCon Maybe Name
up Int
i Name
n [LExp]
es) = Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon Maybe Name
up Int
i Name
n (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
caseFloat [LExp]
es)
caseFloat (LOp PrimFn
f [LExp]
es) = PrimFn -> [LExp] -> LExp
LOp PrimFn
f (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
caseFloat [LExp]
es)
caseFloat (LLam [Name]
ns LExp
sc) = [Name] -> LExp -> LExp
LLam [Name]
ns (LExp -> LExp
caseFloat LExp
sc)
caseFloat (LLet Name
v LExp
val LExp
sc) = Name -> LExp -> LExp -> LExp
LLet Name
v (LExp -> LExp
caseFloat LExp
val) (LExp -> LExp
caseFloat LExp
sc)
caseFloat (LCase CaseType
_ (LCase CaseType
ct LExp
exp [LAlt]
alts) [LAlt]
alts')
| forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all LAlt -> Bool
conRHS [LAlt]
alts Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Int
length [LAlt]
alts forall a. Eq a => a -> a -> Bool
== Int
1
= LExp -> LExp
conOpt forall a b. (a -> b) -> a -> b
$ LExp -> LExp
replaceInCase (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ct (LExp -> LExp
caseFloat LExp
exp) (forall a b. (a -> b) -> [a] -> [b]
map ([LAlt] -> LAlt -> LAlt
updateWith [LAlt]
alts') [LAlt]
alts))
where
conRHS :: LAlt -> Bool
conRHS (LConCase Int
_ Name
_ [Name]
_ (LCon Maybe Name
_ Int
_ Name
_ [LExp]
_)) = Bool
True
conRHS (LConstCase Const
_ (LCon Maybe Name
_ Int
_ Name
_ [LExp]
_)) = Bool
True
conRHS (LDefaultCase (LCon Maybe Name
_ Int
_ Name
_ [LExp]
_)) = Bool
True
conRHS LAlt
_ = Bool
False
updateWith :: [LAlt] -> LAlt -> LAlt
updateWith [LAlt]
alts (LConCase Int
i Name
n [Name]
es LExp
rhs) =
forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
n [Name]
es (LExp -> LExp
caseFloat (LExp -> LExp
conOpt (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
Shared (LExp -> LExp
caseFloat LExp
rhs) [LAlt]
alts)))
updateWith [LAlt]
alts (LConstCase Const
c LExp
rhs) =
forall e. Const -> e -> LAlt' e
LConstCase Const
c (LExp -> LExp
caseFloat (LExp -> LExp
conOpt (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
Shared (LExp -> LExp
caseFloat LExp
rhs) [LAlt]
alts)))
updateWith [LAlt]
alts (LDefaultCase LExp
rhs) =
forall e. e -> LAlt' e
LDefaultCase (LExp -> LExp
caseFloat (LExp -> LExp
conOpt (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
Shared (LExp -> LExp
caseFloat LExp
rhs) [LAlt]
alts)))
caseFloat (LCase CaseType
ct LExp
exp [LAlt]
alts')
= LExp -> LExp
conOpt forall a b. (a -> b) -> a -> b
$ LExp -> LExp
replaceInCase (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ct (LExp -> LExp
caseFloat LExp
exp) (forall a b. (a -> b) -> [a] -> [b]
map LAlt -> LAlt
cfAlt [LAlt]
alts'))
where
cfAlt :: LAlt -> LAlt
cfAlt (LConCase Int
i Name
n [Name]
es LExp
rhs) = forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
n [Name]
es (LExp -> LExp
caseFloat LExp
rhs)
cfAlt (LConstCase Const
c LExp
rhs) = forall e. Const -> e -> LAlt' e
LConstCase Const
c (LExp -> LExp
caseFloat LExp
rhs)
cfAlt (LDefaultCase LExp
rhs) = forall e. e -> LAlt' e
LDefaultCase (LExp -> LExp
caseFloat LExp
rhs)
caseFloat LExp
exp = LExp
exp
conOpt :: LExp -> LExp
conOpt :: LExp -> LExp
conOpt (LCase CaseType
ct (LCon Maybe Name
_ Int
t Name
n [LExp]
args) [LAlt]
alts)
= Name -> [LExp] -> [LAlt] -> LExp
pickAlt Name
n [LExp]
args [LAlt]
alts
where
pickAlt :: Name -> [LExp] -> [LAlt] -> LExp
pickAlt Name
n [LExp]
args (LConCase Int
i Name
n' [Name]
es LExp
rhs : [LAlt]
as) | Name
n forall a. Eq a => a -> a -> Bool
== Name
n'
= [(Name, LExp)] -> LExp -> LExp
substAll (forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
es [LExp]
args) LExp
rhs
pickAlt Name
_ [LExp]
_ (LDefaultCase LExp
rhs : [LAlt]
as) = LExp
rhs
pickAlt Name
n [LExp]
args (LAlt
_ : [LAlt]
as) = Name -> [LExp] -> [LAlt] -> LExp
pickAlt Name
n [LExp]
args [LAlt]
as
pickAlt Name
n [LExp]
args [] = forall a. HasCallStack => String -> a
error String
"Can't happen pickAlt - impossible case found"
substAll :: [(Name, LExp)] -> LExp -> LExp
substAll [] LExp
rhs = LExp
rhs
substAll ((Name
n, LExp
tm) : [(Name, LExp)]
ss) LExp
rhs = Name -> LExp -> LExp -> LExp
lsubst Name
n LExp
tm ([(Name, LExp)] -> LExp -> LExp
substAll [(Name, LExp)]
ss LExp
rhs)
conOpt LExp
tm = LExp
tm
replaceInCase :: LExp -> LExp
replaceInCase :: LExp -> LExp
replaceInCase (LCase CaseType
ty LExp
e [LAlt]
alts)
= CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ty LExp
e (LExp -> [LAlt] -> [LAlt]
replaceInAlts LExp
e [LAlt]
alts)
replaceInCase LExp
exp = LExp
exp
replaceInAlts :: LExp -> [LAlt] -> [LAlt]
replaceInAlts :: LExp -> [LAlt] -> [LAlt]
replaceInAlts LExp
exp [LAlt]
alts = forall {e}. [LAlt' e] -> [LAlt' e]
dropDups forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (LExp -> LAlt -> [LAlt]
replaceInAlt LExp
exp) [LAlt]
alts
dropDups :: [LAlt' e] -> [LAlt' e]
dropDups (alt :: LAlt' e
alt@(LConCase Int
_ Name
i [Name]
n e
ns) : [LAlt' e]
alts)
= LAlt' e
alt forall a. a -> [a] -> [a]
: [LAlt' e] -> [LAlt' e]
dropDups (forall a. (a -> Bool) -> [a] -> [a]
filter (forall {e}. Name -> LAlt' e -> Bool
notTag Name
i) [LAlt' e]
alts)
where
notTag :: Name -> LAlt' e -> Bool
notTag Name
i (LConCase Int
_ Name
j [Name]
n e
ns) = Name
i forall a. Eq a => a -> a -> Bool
/= Name
j
notTag Name
_ LAlt' e
_ = Bool
True
dropDups (LAlt' e
c : [LAlt' e]
alts) = LAlt' e
c forall a. a -> [a] -> [a]
: [LAlt' e] -> [LAlt' e]
dropDups [LAlt' e]
alts
dropDups [] = []
replaceInAlt :: LExp -> LAlt -> [LAlt]
replaceInAlt :: LExp -> LAlt -> [LAlt]
replaceInAlt exp :: LExp
exp@(LV Name
_) (LConCase Int
i Name
con [Name]
args LExp
rhs)
= [forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
con [Name]
args forall a b. (a -> b) -> a -> b
$
LExp -> LExp -> LExp -> LExp
replaceExp (Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon forall a. Maybe a
Nothing Int
i Name
con (forall a b. (a -> b) -> [a] -> [b]
map Name -> LExp
LV [Name]
args)) LExp
exp LExp
rhs]
replaceInAlt exp :: LExp
exp@(LV Name
var) (LDefaultCase (LCase CaseType
ty (LV Name
var') [LAlt]
alts))
| Name
var forall a. Eq a => a -> a -> Bool
== Name
var' = [LAlt]
alts
replaceInAlt LExp
exp LAlt
a = [LAlt
a]
replaceExp :: LExp -> LExp -> LExp -> LExp
replaceExp :: LExp -> LExp -> LExp -> LExp
replaceExp (LCon Maybe Name
_ Int
t Name
n [LExp]
args) LExp
new (LCon Maybe Name
_ Int
t' Name
n' [LExp]
args')
| Name
n forall a. Eq a => a -> a -> Bool
== Name
n' Bool -> Bool -> Bool
&& [LExp]
args forall a. Eq a => a -> a -> Bool
== [LExp]
args' = LExp
new
replaceExp (LCon Maybe Name
_ Int
t Name
n [LExp]
args) LExp
new (LApp Bool
_ (LV Name
n') [LExp]
args')
| Name
n forall a. Eq a => a -> a -> Bool
== Name
n' Bool -> Bool -> Bool
&& [LExp]
args forall a. Eq a => a -> a -> Bool
== [LExp]
args' = LExp
new
replaceExp LExp
old LExp
new LExp
tm = LExp
tm
getRHS :: LAlt' e -> e
getRHS (LConCase Int
i Name
n [Name]
es e
rhs) = e
rhs
getRHS (LConstCase Const
_ e
rhs) = e
rhs
getRHS (LDefaultCase e
rhs) = e
rhs
getLams :: [LExp] -> [Name]
getLams [] = []
getLams (LLam [Name]
args LExp
tm : [LExp]
cs) = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
args [LExp]
cs
getLams (LLet Name
n LExp
val LExp
exp : [LExp]
cs) = [LExp] -> [Name]
getLams (LExp
exp forall a. a -> [a] -> [a]
: [LExp]
cs)
getLams [LExp]
_ = []
getLamPrefix :: [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
as [] = [Name]
as
getLamPrefix [Name]
as (LLam [Name]
args LExp
tm : [LExp]
cs)
| forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
args forall a. Ord a => a -> a -> Bool
< forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
as = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
args [LExp]
cs
| Bool
otherwise = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
as [LExp]
cs
getLamPrefix [Name]
as (LLet Name
n LExp
val LExp
exp : [LExp]
cs) = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
as (LExp
exp forall a. a -> [a] -> [a]
: [LExp]
cs)
getLamPrefix [Name]
as (LExp
_ : [LExp]
cs) = []
eta :: LExp -> LExp
eta :: LExp -> LExp
eta (LApp Bool
tc LExp
a [LExp]
es) = Bool -> LExp -> [LExp] -> LExp
LApp Bool
tc (LExp -> LExp
eta LExp
a) (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
eta [LExp]
es)
eta (LLazyApp Name
n [LExp]
es) = Name -> [LExp] -> LExp
LLazyApp Name
n (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
eta [LExp]
es)
eta (LLazyExp LExp
e) = LExp -> LExp
LLazyExp (LExp -> LExp
eta LExp
e)
eta (LForce LExp
e) = LExp -> LExp
LForce (LExp -> LExp
eta LExp
e)
eta (LLet Name
n LExp
val LExp
sc) = Name -> LExp -> LExp -> LExp
LLet Name
n (LExp -> LExp
eta LExp
val) (LExp -> LExp
eta LExp
sc)
eta (LLam [Name]
args (LApp Bool
tc LExp
f [LExp]
args'))
| [LExp]
args' forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map Name -> LExp
LV [Name]
args = LExp -> LExp
eta LExp
f
eta (LLam [Name]
args LExp
e) = [Name] -> LExp -> LExp
LLam [Name]
args (LExp -> LExp
eta LExp
e)
eta (LProj LExp
e Int
i) = LExp -> Int -> LExp
LProj (LExp -> LExp
eta LExp
e) Int
i
eta (LCon Maybe Name
a Int
t Name
n [LExp]
es) = Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon Maybe Name
a Int
t Name
n (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
eta [LExp]
es)
eta (LCase CaseType
ct LExp
e [LAlt]
alts) = CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ct (LExp -> LExp
eta LExp
e) (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LExp -> LExp
eta) [LAlt]
alts)
eta (LOp PrimFn
f [LExp]
es) = PrimFn -> [LExp] -> LExp
LOp PrimFn
f (forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
eta [LExp]
es)
eta LExp
tm = LExp
tm