{-|
Module      : IRTS.LangOpts
Description : Transformations to apply to Idris' IR.

License     : BSD3
Maintainer  : The Idris Community.
-}
{-# 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 = [(Name, LDecl)] -> Ctxt LDecl -> Ctxt LDecl
forall a. [(Name, a)] -> Ctxt a -> Ctxt a
addAlist [(Name, LDecl)]
lds Ctxt LDecl
forall {k} {a}. Map k a
emptyContext in
                    ((Name, LDecl) -> (Name, LDecl))
-> [(Name, LDecl)] -> [(Name, LDecl)]
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 <- StateT Int Identity Int
forall s (m :: * -> *). MonadState s m => m s
get
           Int -> StateT Int Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
           Name -> State Int Name
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> State Int Name) -> Name -> State Int Name
forall a b. (a -> b) -> a -> b
$ Int -> String -> Name
sMN Int
i String
"in"

-- | Inline inside a declaration.
--
-- Variables are still Name at this stage.  Need to preserve
-- uniqueness of variable names in the resulting definition, so invent
-- a new name for every variable we encounter
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 = State Int LExp -> Int -> LExp
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)
                             ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
args)
            -- do some case floating, which might arise as a result
            -- then, eta contract
            res :: LExp
res = LExp -> LExp
eta (LExp -> LExp) -> LExp -> LExp
forall a b. (a -> b) -> a -> b
$ Integer -> LExp -> LExp
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Ctxt LDecl
defs (LDecl -> LDecl) -> LDecl -> LDecl
forall a b. (a -> b) -> a -> b
$
                     [LOpt] -> Name -> [Name] -> LExp -> LDecl
LFun [LOpt]
opts Name
topn (((Name, Name) -> Name) -> [(Name, Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Name) -> Name
forall a b. (a, b) -> b
snd [(Name, Name)]
initNames [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ [Name]
args') LExp
body
                 LExp
_ -> Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Ctxt LDecl
defs (LDecl -> LDecl) -> LDecl -> LDecl
forall a b. (a -> b) -> a -> b
$
                        [LOpt] -> Name -> [Name] -> LExp -> LDecl
LFun [LOpt]
opts Name
topn (((Name, Name) -> Name) -> [(Name, Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Name) -> Name
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 LExp -> LExp -> Bool
forall a. Eq a => a -> a -> Bool
== LExp
tm
                 then LExp
res
                 else t -> LExp -> LExp
caseFloats (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) LExp
res

    initNames :: [(Name, Name)]
initNames = (Name -> Int -> (Name, Name)) -> [Name] -> [Int] -> [(Name, Name)]
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 = ((Name, Name) -> (Name, LExp)) -> [(Name, Name)] -> [(Name, LExp)]
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 [LExp] -> [LExp] -> [LExp]
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) (Name, LExp) -> [(Name, LExp)] -> [(Name, LExp)]
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 (LExp -> LExp) -> ([LExp] -> LExp) -> [LExp] -> LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> [LExp] -> LExp
LLazyApp Name
n ([LExp] -> LExp) -> StateT Int Identity [LExp] -> State Int LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((LExp -> State Int LExp) -> [LExp] -> StateT Int Identity [LExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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
_ -> LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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 (LExp -> LExp) -> (LExp -> LExp) -> LExp -> LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LExp -> LExp
LLazyExp (LExp -> LExp) -> State Int LExp -> State Int LExp
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
-- Special case for io_bind, because it needs to keep executing the first
-- action, and is worth inlining to avoid the thunk
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 Name -> Name -> Bool
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) (Name, LExp) -> [(Name, LExp)] -> [(Name, LExp)]
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) (Name, LExp) -> [(Name, LExp)] -> [(Name, LExp)]
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])
         LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (LExp -> State Int LExp) -> LExp -> State Int LExp
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 Name -> Name -> Bool
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) (Name, LExp) -> [(Name, LExp)] -> [(Name, LExp)]
forall a. a -> [a] -> [a]
: [(Name, LExp)]
env) [Name]
rec Ctxt LDecl
defs (Bool -> LExp -> [LExp] -> LExp
LApp Bool
False LExp
k [LExp
world])
         -- Needs to be a LLet to make sure the action gets evaluated
         LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (LExp -> State Int LExp) -> LExp -> State Int LExp
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' <- (LExp -> State Int LExp) -> [LExp] -> StateT Int Identity [LExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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' [LExp] -> [LExp] -> [LExp]
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' (LExp -> LExp -> LExp)
-> State Int LExp -> StateT Int Identity (LExp -> LExp)
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
                 StateT Int Identity (LExp -> LExp)
-> State Int LExp -> State Int LExp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
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') (Name, LExp) -> [(Name, LExp)] -> [(Name, LExp)]
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 (LExp -> LExp) -> State Int LExp -> State Int LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (LExp -> Int -> LExp
LProj (LExp -> Int -> LExp)
-> State Int LExp -> StateT Int Identity (Int -> LExp)
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 StateT Int Identity (Int -> LExp)
-> StateT Int Identity Int -> State Int LExp
forall a b.
StateT Int Identity (a -> b)
-> StateT Int Identity a -> StateT Int Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> StateT Int Identity Int
forall a. a -> StateT Int Identity a
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 (LExp -> LExp) -> State Int LExp -> State Int LExp
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 ([LExp] -> LExp) -> StateT Int Identity [LExp] -> State Int LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (LExp -> State Int LExp) -> [LExp] -> StateT Int Identity [LExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 [])
    = LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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 LExp -> [LAlt] -> Maybe ([(Name, LExp)], LExp)
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' <- (LAlt -> StateT Int Identity LAlt)
-> [LAlt] -> StateT Int Identity [LAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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
                   -- If they're all lambdas, bind the lambda at the top
                   let prefix :: [Name]
prefix = [LExp] -> [Name]
getLams ((LAlt -> LExp) -> [LAlt] -> [LExp]
forall a b. (a -> b) -> [a] -> [b]
map LAlt -> LExp
forall {e}. LAlt' e -> e
getRHS [LAlt]
alts')
                   case [Name]
prefix of
                        [] -> LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (LExp -> State Int LExp) -> LExp -> State Int LExp
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 <- (LAlt -> StateT Int Identity LAlt)
-> [LAlt] -> StateT Int Identity [LAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Name] -> LAlt -> StateT Int Identity LAlt
dropArgs [Name]
args) [LAlt]
alts'
                                   LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (LExp -> State Int LExp) -> LExp -> State Int LExp
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' [] = Maybe ([(Name, LExp)], b)
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 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n' = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
forall a. a -> Maybe a
Just ([Name] -> [LExp] -> [(Name, LExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
es [LExp]
args [(Name, LExp)] -> [(Name, LExp)] -> [(Name, LExp)]
forall a. [a] -> [a] -> [a]
++ [(Name, LExp)]
env, b
rhs)
    evalAlts (LConst Const
c) (LConstCase Const
c' b
rhs : [LAlt' b]
as)
        | Const
c Const -> Const -> Bool
forall a. Eq a => a -> a -> Bool
== Const
c' = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
    evalAlts (LCon Maybe Name
_ Int
_ Name
_ [LExp]
_) (LDefaultCase b
rhs : [LAlt' b]
as) = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
    evalAlts (LConst Const
_) (LDefaultCase b
rhs : [LAlt' b]
as) = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
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 (LExp -> LExp) -> ([LExp] -> LExp) -> [LExp] -> LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimFn -> [LExp] -> LExp
LOp PrimFn
f ([LExp] -> LExp) -> StateT Int Identity [LExp] -> State Int LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (LExp -> State Int LExp) -> [LExp] -> StateT Int Identity [LExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 (LExp -> LExp)
-> ([(FDesc, LExp)] -> LExp) -> [(FDesc, LExp)] -> LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        FDesc -> FDesc -> [(FDesc, LExp)] -> LExp
LForeign FDesc
t FDesc
s ([(FDesc, LExp)] -> LExp)
-> StateT Int Identity [(FDesc, LExp)] -> State Int LExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((FDesc, LExp) -> StateT Int Identity (FDesc, LExp))
-> [(FDesc, LExp)] -> StateT Int Identity [(FDesc, LExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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
                                             (FDesc, LExp) -> StateT Int Identity (FDesc, LExp)
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (FDesc
t, LExp
e')) [(FDesc, LExp)]
args
-- save the interesting cases for the end:
-- lambdas, and names to reduce
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' <- (Name -> StateT Int Identity (Name, Name))
-> [Name] -> StateT Int Identity [(Name, Name)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Name
n -> do Name
n' <- State Int Name
nextN
                                              (Name, Name) -> StateT Int Identity (Name, Name)
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n, Name
n')) [Name]
args'
                        [Name] -> LExp -> LExp
LLam (((Name, Name) -> Name) -> [(Name, Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Name) -> Name
forall a b. (a, b) -> b
snd [(Name, Name)]
ns') (LExp -> LExp) -> State Int LExp -> State Int LExp
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, Name) -> (Name, LExp)) -> [(Name, Name)] -> [(Name, LExp)]
forall a b. (a -> b) -> [a] -> [b]
map (\ (Name
n, Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
ns' [(Name, LExp)] -> [(Name, LExp)] -> [(Name, LExp)]
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 Name -> [(Name, LExp)] -> Maybe LExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
n [(Name, LExp)]
env of
           Just LExp
t
               | LExp
t LExp -> LExp -> Bool
forall a. Eq a => a -> a -> Bool
/= Name -> LExp
LV Name
n Bool -> Bool -> Bool
&& Name
n Name -> [Name] -> Bool
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 Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
rec) Ctxt LDecl
defs LExp
t
               | Bool
otherwise -> LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
t)
           Maybe LExp
Nothing
               | Name
n Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
rec,
                 Just (LFun [LOpt]
opts Name
_ [Name]
args LExp
body) <- Name -> Ctxt LDecl -> Maybe LDecl
forall a. Name -> Ctxt a -> Maybe a
lookupCtxtExact Name
n Ctxt LDecl
defs,
                 LOpt
Inline LOpt -> [LOpt] -> Bool
forall a. Eq a => a -> [a] -> Bool
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 Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
rec) Ctxt LDecl
defs LExp
var [Name]
args LExp
body
               | Just (LConstructor Name
n Int
t Int
a) <- Name -> Ctxt LDecl -> Maybe LDecl
forall a. Name -> Ctxt a -> Maybe a
lookupCtxtExact Name
n Ctxt LDecl
defs ->
                         LExp -> State Int LExp
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon Maybe Name
forall a. Maybe a
Nothing Int
t Name
n [LExp]
stk)
               | Bool
otherwise -> LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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 = LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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' <- (Name -> StateT Int Identity (Name, Name))
-> [Name] -> StateT Int Identity [(Name, Name)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Name
n -> do Name
n' <- State Int Name
nextN
                               (Name, Name) -> StateT Int Identity (Name, Name)
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n, Name
n')) [Name]
es
         Int -> Name -> [Name] -> LExp -> LAlt
forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
n (((Name, Name) -> Name) -> [(Name, Name)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Name) -> Name
forall a b. (a, b) -> b
snd [(Name, Name)]
ns') (LExp -> LAlt) -> State Int LExp -> StateT Int Identity LAlt
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, Name) -> (Name, LExp)) -> [(Name, Name)] -> [(Name, LExp)]
forall a b. (a -> b) -> [a] -> [b]
map (\ (Name
n, Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
ns' [(Name, LExp)] -> [(Name, LExp)] -> [(Name, LExp)]
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)
    = Const -> LExp -> LAlt
forall e. Const -> e -> LAlt' e
LConstCase Const
c (LExp -> LAlt) -> State Int LExp -> StateT Int Identity LAlt
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)
    = LExp -> LAlt
forall e. e -> LAlt' e
LDefaultCase (LExp -> LAlt) -> State Int LExp -> StateT Int Identity LAlt
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
         LAlt -> StateT Int Identity LAlt
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Name -> [Name] -> LExp -> LAlt
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
         LAlt -> StateT Int Identity LAlt
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Const -> LExp -> LAlt
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
         LAlt -> StateT Int Identity LAlt
forall a. a -> StateT Int Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (LExp -> LAlt
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 = Int -> [Name] -> [Name]
forall a. Int -> [a] -> [a]
take ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
as) [Name]
args
         [LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] ((Name -> Name -> (Name, LExp))
-> [Name] -> [Name] -> [(Name, LExp)]
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) [] Ctxt LDecl
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
         LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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 = LExp -> State Int LExp
forall a. a -> StateT Int Identity a
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) ((LExp -> LExp) -> [LExp] -> [LExp]
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 ((LExp -> LExp) -> [LExp] -> [LExp]
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 ((LExp -> LExp) -> [LExp] -> [LExp]
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')
    | (LAlt -> Bool) -> [LAlt] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all LAlt -> Bool
conRHS [LAlt]
alts Bool -> Bool -> Bool
|| [LAlt] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LAlt]
alts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
    = LExp -> LExp
conOpt (LExp -> LExp) -> LExp -> LExp
forall a b. (a -> b) -> a -> b
$ LExp -> LExp
replaceInCase (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ct (LExp -> LExp
caseFloat LExp
exp) ((LAlt -> LAlt) -> [LAlt] -> [LAlt]
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) =
        Int -> Name -> [Name] -> LExp -> LAlt
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) =
        Const -> LExp -> LAlt
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) =
        LExp -> LAlt
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 (LExp -> LExp) -> LExp -> LExp
forall a b. (a -> b) -> a -> b
$ LExp -> LExp
replaceInCase (CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ct (LExp -> LExp
caseFloat LExp
exp) ((LAlt -> LAlt) -> [LAlt] -> [LAlt]
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) = Int -> Name -> [Name] -> LExp -> LAlt
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) = Const -> LExp -> LAlt
forall e. Const -> e -> LAlt' e
LConstCase Const
c (LExp -> LExp
caseFloat LExp
rhs)
    cfAlt (LDefaultCase LExp
rhs) = LExp -> LAlt
forall e. e -> LAlt' e
LDefaultCase (LExp -> LExp
caseFloat LExp
rhs)
caseFloat LExp
exp = LExp
exp

-- Case of constructor
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 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n'
        = [(Name, LExp)] -> LExp -> LExp
substAll ([Name] -> [LExp] -> [(Name, LExp)]
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 [] = String -> LExp
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 = [LAlt] -> [LAlt]
forall {e}. [LAlt' e] -> [LAlt' e]
dropDups ([LAlt] -> [LAlt]) -> [LAlt] -> [LAlt]
forall a b. (a -> b) -> a -> b
$ (LAlt -> [LAlt]) -> [LAlt] -> [LAlt]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (LExp -> LAlt -> [LAlt]
replaceInAlt LExp
exp) [LAlt]
alts

-- Drop overlapping case (arising from case merging of overlapping
-- patterns)
dropDups :: [LAlt' e] -> [LAlt' e]
dropDups (alt :: LAlt' e
alt@(LConCase Int
_ Name
i [Name]
n e
ns) : [LAlt' e]
alts)
    = LAlt' e
alt LAlt' e -> [LAlt' e] -> [LAlt' e]
forall a. a -> [a] -> [a]
: [LAlt' e] -> [LAlt' e]
dropDups ((LAlt' e -> Bool) -> [LAlt' e] -> [LAlt' e]
forall a. (a -> Bool) -> [a] -> [a]
filter (Name -> LAlt' e -> Bool
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 Name -> Name -> Bool
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 LAlt' e -> [LAlt' e] -> [LAlt' e]
forall a. a -> [a] -> [a]
: [LAlt' e] -> [LAlt' e]
dropDups [LAlt' e]
alts
dropDups [] = []


replaceInAlt :: LExp -> LAlt -> [LAlt]
-- In an alternative, if the case appears on the right hand side, replace
-- it with the given expression, to preserve sharing
replaceInAlt :: LExp -> LAlt -> [LAlt]
replaceInAlt exp :: LExp
exp@(LV Name
_) (LConCase Int
i Name
con [Name]
args LExp
rhs)
    = [Int -> Name -> [Name] -> LExp -> LAlt
forall e. Int -> Name -> [Name] -> e -> LAlt' e
LConCase Int
i Name
con [Name]
args (LExp -> LAlt) -> LExp -> LAlt
forall a b. (a -> b) -> a -> b
$
          LExp -> LExp -> LExp -> LExp
replaceExp (Maybe Name -> Int -> Name -> [LExp] -> LExp
LCon Maybe Name
forall a. Maybe a
Nothing Int
i Name
con ((Name -> LExp) -> [Name] -> [LExp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> LExp
LV [Name]
args)) LExp
exp LExp
rhs]
-- if a default case inspects the same variable as the case it's in,
-- remove the inspection and replace with the alternatives
-- (i.e. merge the inner case block)
replaceInAlt exp :: LExp
exp@(LV Name
var) (LDefaultCase (LCase CaseType
ty (LV Name
var') [LAlt]
alts))
    | Name
var Name -> Name -> Bool
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 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n' Bool -> Bool -> Bool
&& [LExp]
args [LExp] -> [LExp] -> Bool
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 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n' Bool -> Bool -> Bool
&& [LExp]
args [LExp] -> [LExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [LExp]
args' = LExp
new
replaceExp LExp
old LExp
new LExp
tm = LExp
tm

-- dropArgs as (LConstCase c rhs) = LConstCase c (dropRHS as rhs)
-- dropArgs as (LDefaultCase rhs) = LDefaultCase (dropRHS as rhs)

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 LExp -> [LExp] -> [LExp]
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)
    | [Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Name] -> Int
forall a. [a] -> Int
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 LExp -> [LExp] -> [LExp]
forall a. a -> [a] -> [a]
: [LExp]
cs)
getLamPrefix [Name]
as (LExp
_ : [LExp]
cs) = []

-- eta contract ('\x -> f x' can just be compiled as 'f' when f is local)
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) ((LExp -> LExp) -> [LExp] -> [LExp]
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 ((LExp -> LExp) -> [LExp] -> [LExp]
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' [LExp] -> [LExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (Name -> LExp) -> [Name] -> [LExp]
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 ((LExp -> LExp) -> [LExp] -> [LExp]
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) ((LAlt -> LAlt) -> [LAlt] -> [LAlt]
forall a b. (a -> b) -> [a] -> [b]
map ((LExp -> LExp) -> LAlt -> LAlt
forall a b. (a -> b) -> LAlt' a -> LAlt' b
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 ((LExp -> LExp) -> [LExp] -> [LExp]
forall a b. (a -> b) -> [a] -> [b]
map LExp -> LExp
eta [LExp]
es)
eta LExp
tm = LExp
tm