{-|
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)
import Data.List

import Debug.Trace

inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
inlineAll :: [(Name, LDecl)] -> [(Name, LDecl)]
inlineAll lds :: [(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 (\ (n :: Name
n, def :: 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
+ 1)
           Name -> State Int Name
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 "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' 1

doInline' :: Int -> LDefs -> LDecl -> LDecl
doInline' :: Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' 0 defs :: Ctxt LDecl
defs d :: LDecl
d = LDecl
d
doInline' i :: Int
i defs :: Ctxt LDecl
defs d :: LDecl
d@(LConstructor _ _ _) = LDecl
d
doInline' i :: Int
i defs :: Ctxt LDecl
defs (LFun opts :: [LOpt]
opts topn :: Name
topn args :: [Name]
args exp :: 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 (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 10 LExp
inl in
            case LExp
res of
                 LLam args' :: [Name]
args' body :: LExp
body ->
                   Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- 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
                 _ -> Int -> Ctxt LDecl -> LDecl -> LDecl
doInline' (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- 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 0 tm :: LExp
tm = LExp
tm
    caseFloats n :: t
n tm :: 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
-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 (\n :: Name
n i :: Int
i -> (Name
n, Name -> Int -> Name
newn Name
n Int
i)) [Name]
args [0..]
    initEnv :: [(Name, LExp)]
initEnv = ((Name, Name) -> (Name, LExp)) -> [(Name, Name)] -> [(Name, LExp)]
forall a b. (a -> b) -> [a] -> [b]
map (\(n :: Name
n, n' :: Name
n') -> (Name
n, Name -> LExp
LV Name
n')) [(Name, Name)]
initNames
    newn :: Name -> Int -> Name
newn (UN n :: Text
n) i :: Int
i = Int -> Text -> Name
MN Int
i Text
n
    newn _ i :: Int
i = Int -> String -> Name
sMN Int
i "arg"

unload :: [LExp] -> LExp -> LExp
unload :: [LExp] -> LExp -> LExp
unload [] e :: LExp
e = LExp
e
unload stk :: [LExp]
stk (LApp tc :: Bool
tc e :: LExp
e args :: [LExp]
args) = Bool -> LExp -> [LExp] -> LExp
LApp Bool
tc LExp
e ([LExp]
args [LExp] -> [LExp] -> [LExp]
forall a. [a] -> [a] -> [a]
++ [LExp]
stk)
unload stk :: [LExp]
stk e :: 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 env :: [(Name, LExp)]
env (a :: Name
a : args :: [Name]
args) (v :: LExp
v : stk :: [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 env :: [(Name, LExp)]
env args :: [Name]
args stk :: [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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LLazyApp n :: Name
n es :: [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)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es)
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LForce e :: 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 forced :: 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 n :: Name
n es :: [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 -> State Int LExp
forall (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk (LExp -> LExp
LForce LExp
e'))
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LLazyExp e :: 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 [] env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LApp t :: Bool
t (LV n :: Name
n) [_, _, _, act :: LExp
act, (LLam [arg :: Name
arg] k :: LExp
k)])
    | Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Name
sUN "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 (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 (world :: LExp
world : stk :: [LExp]
stk) env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LApp t :: Bool
t (LV n :: Name
n) [_, _, _, act :: LExp
act, (LLam [arg :: Name
arg] k :: LExp
k)])
    | Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Name
sUN "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 (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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LApp t :: Bool
t f :: LExp
f es :: [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)
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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LLet n :: Name
n val :: LExp
val sc :: 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 (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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LProj exp :: LExp
exp i :: 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 (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> StateT Int Identity Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i)
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LCon loc :: Maybe Name
loc i :: Int
i n :: Name
n es :: [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)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es)
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LCase ty :: CaseType
ty e :: LExp
e [])
    = LExp -> State Int LExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure LExp
LNothing
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LCase ty :: CaseType
ty e :: LExp
e alts :: [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 (env' :: [(Name, LExp)]
env', tm :: 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
              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)
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 p. LAlt' p -> p
getRHS [LAlt]
alts')
                   case [Name]
prefix of
                        [] -> LExp -> State Int LExp
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')
                        args :: [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)
mapM ([Name] -> LAlt -> StateT Int Identity LAlt
dropArgs [Name]
args) [LAlt]
alts'
                                   LExp -> State Int LExp
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 e' :: LExp
e' [] = Maybe ([(Name, LExp)], b)
forall a. Maybe a
Nothing
    evalAlts (LCon _ t :: Int
t n :: Name
n args :: [LExp]
args) (LConCase i :: Int
i n' :: Name
n' es :: [Name]
es rhs :: b
rhs : as :: [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 c :: Const
c) (LConstCase c' :: Const
c' rhs :: b
rhs : as :: [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 _ _ _ _) (LDefaultCase rhs :: b
rhs : as :: [LAlt' b]
as) = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
    evalAlts (LConst _) (LDefaultCase rhs :: b
rhs : as :: [LAlt' b]
as) = ([(Name, LExp)], b) -> Maybe ([(Name, LExp)], b)
forall a. a -> Maybe a
Just ([(Name, LExp)]
env, b
rhs)
    evalAlts tm :: LExp
tm (_ : as :: [LAlt' b]
as) = LExp -> [LAlt' b] -> Maybe ([(Name, LExp)], b)
evalAlts LExp
tm [LAlt' b]
as
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LOp f :: PrimFn
f es :: [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)
mapM ([LExp]
-> [(Name, LExp)] -> [Name] -> Ctxt LDecl -> LExp -> State Int LExp
eval [] [(Name, LExp)]
env [Name]
rec Ctxt LDecl
defs) [LExp]
es
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LForeign t :: FDesc
t s :: FDesc
s args :: [(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)
mapM (\(t :: FDesc
t, e :: 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 (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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LLam args :: [Name]
args sc :: LExp
sc)
    | (env' :: [(Name, LExp)]
env', args' :: [Name]
args', stk' :: [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
               as :: [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)
mapM (\n :: Name
n -> do Name
n' <- State Int Name
nextN
                                              (Name, Name) -> StateT Int Identity (Name, Name)
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 (\ (n :: Name
n, 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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs var :: LExp
var@(LV n :: 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 t :: 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 (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
t)
           Nothing
               | Name
n Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
rec,
                 Just (LFun opts :: [LOpt]
opts _ args :: [Name]
args body :: 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 (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 n :: Name
n t :: Int
t a :: 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 (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 (m :: * -> *) a. Monad m => a -> m a
return ([LExp] -> LExp -> LExp
unload [LExp]
stk LExp
var)
eval stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs t :: LExp
t = LExp -> State Int LExp
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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LConCase i :: Int
i n :: Name
n es :: [Name]
es rhs :: 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)
mapM (\n :: Name
n -> do Name
n' <- State Int Name
nextN
                               (Name, Name) -> StateT Int Identity (Name, Name)
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 (\ (n :: Name
n, 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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LConstCase c :: Const
c e :: 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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs (LDefaultCase e :: 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 stk :: [LExp]
stk env :: [(Name, LExp)]
env rec :: [Name]
rec defs :: Ctxt LDecl
defs var :: LExp
var args :: [Name]
args body :: 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 as :: [Name]
as (LConCase i :: Int
i n :: Name
n es :: [Name]
es t :: LExp
t)
    = do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
         LAlt -> StateT Int Identity LAlt
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 as :: [Name]
as (LConstCase c :: Const
c t :: LExp
t)
    = do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
         LAlt -> StateT Int Identity LAlt
forall (m :: * -> *) a. Monad m => a -> m a
return (Const -> LExp -> LAlt
forall e. Const -> e -> LAlt' e
LConstCase Const
c LExp
rhs')
dropArgs as :: [Name]
as (LDefaultCase t :: LExp
t)
    = do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
t
         LAlt -> StateT Int Identity LAlt
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 as :: [Name]
as (LLam args :: [Name]
args rhs :: LExp
rhs)
    = do let old :: [Name]
old = Int -> [Name] -> [Name]
forall a. Int -> [a] -> [a]
take ([Name] -> 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 (\ o :: Name
o n :: 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 as :: [Name]
as (LLet n :: Name
n val :: LExp
val rhs :: LExp
rhs)
    = do LExp
rhs' <- [Name] -> LExp -> State Int LExp
dropArgsTm [Name]
as LExp
rhs
         LExp -> State Int LExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> LExp -> LExp -> LExp
LLet Name
n LExp
val LExp
rhs')
dropArgsTm as :: [Name]
as tm :: LExp
tm = LExp -> State Int LExp
forall (m :: * -> *) a. Monad m => a -> m a
return LExp
tm

caseFloat :: LExp -> LExp
caseFloat :: LExp -> LExp
caseFloat (LApp tc :: Bool
tc e :: LExp
e es :: [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 e :: LExp
e) = LExp -> LExp
LLazyExp (LExp -> LExp
caseFloat LExp
e)
caseFloat (LForce e :: LExp
e) = LExp -> LExp
LForce (LExp -> LExp
caseFloat LExp
e)
caseFloat (LCon up :: Maybe Name
up i :: Int
i n :: Name
n es :: [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 f :: PrimFn
f es :: [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 ns :: [Name]
ns sc :: LExp
sc) = [Name] -> LExp -> LExp
LLam [Name]
ns (LExp -> LExp
caseFloat LExp
sc)
caseFloat (LLet v :: Name
v val :: LExp
val sc :: LExp
sc) = Name -> LExp -> LExp -> LExp
LLet Name
v (LExp -> LExp
caseFloat LExp
val) (LExp -> LExp
caseFloat LExp
sc)
caseFloat (LCase _ (LCase ct :: CaseType
ct exp :: LExp
exp alts :: [LAlt]
alts) 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 (t :: * -> *) a. Foldable t => t a -> Int
length [LAlt]
alts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 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 _ _ _ (LCon _ _ _ _)) = Bool
True
    conRHS (LConstCase _ (LCon _ _ _ _)) = Bool
True
    conRHS (LDefaultCase (LCon _ _ _ _)) = Bool
True
    conRHS _ = Bool
False

    updateWith :: [LAlt] -> LAlt -> LAlt
updateWith alts :: [LAlt]
alts (LConCase i :: Int
i n :: Name
n es :: [Name]
es rhs :: 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 alts :: [LAlt]
alts (LConstCase c :: Const
c rhs :: 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 alts :: [LAlt]
alts (LDefaultCase rhs :: 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 ct :: CaseType
ct exp :: LExp
exp alts' :: [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 i :: Int
i n :: Name
n es :: [Name]
es rhs :: 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 c :: Const
c rhs :: LExp
rhs) = Const -> LExp -> LAlt
forall e. Const -> e -> LAlt' e
LConstCase Const
c (LExp -> LExp
caseFloat LExp
rhs)
    cfAlt (LDefaultCase rhs :: LExp
rhs) = LExp -> LAlt
forall e. e -> LAlt' e
LDefaultCase (LExp -> LExp
caseFloat LExp
rhs)
caseFloat exp :: LExp
exp = LExp
exp

-- Case of constructor
conOpt :: LExp -> LExp
conOpt :: LExp -> LExp
conOpt (LCase ct :: CaseType
ct (LCon _ t :: Int
t n :: Name
n args :: [LExp]
args) alts :: [LAlt]
alts)
    = Name -> [LExp] -> [LAlt] -> LExp
pickAlt Name
n [LExp]
args [LAlt]
alts
  where
    pickAlt :: Name -> [LExp] -> [LAlt] -> LExp
pickAlt n :: Name
n args :: [LExp]
args (LConCase i :: Int
i n' :: Name
n' es :: [Name]
es rhs :: LExp
rhs : as :: [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 _ _ (LDefaultCase rhs :: LExp
rhs : as :: [LAlt]
as) = LExp
rhs
    pickAlt n :: Name
n args :: [LExp]
args (_ : as :: [LAlt]
as) = Name -> [LExp] -> [LAlt] -> LExp
pickAlt Name
n [LExp]
args [LAlt]
as
    pickAlt n :: Name
n args :: [LExp]
args [] = String -> LExp
forall a. HasCallStack => String -> a
error "Can't happen pickAlt - impossible case found"

    substAll :: [(Name, LExp)] -> LExp -> LExp
substAll [] rhs :: LExp
rhs = LExp
rhs
    substAll ((n :: Name
n, tm :: LExp
tm) : ss :: [(Name, LExp)]
ss) rhs :: LExp
rhs = Name -> LExp -> LExp -> LExp
lsubst Name
n LExp
tm ([(Name, LExp)] -> LExp -> LExp
substAll [(Name, LExp)]
ss LExp
rhs)
conOpt tm :: LExp
tm = LExp
tm

replaceInCase :: LExp -> LExp
replaceInCase :: LExp -> LExp
replaceInCase (LCase ty :: CaseType
ty e :: LExp
e alts :: [LAlt]
alts)
    = CaseType -> LExp -> [LAlt] -> LExp
LCase CaseType
ty LExp
e (LExp -> [LAlt] -> [LAlt]
replaceInAlts LExp
e [LAlt]
alts)
replaceInCase exp :: LExp
exp = LExp
exp

replaceInAlts :: LExp -> [LAlt] -> [LAlt]
replaceInAlts :: LExp -> [LAlt] -> [LAlt]
replaceInAlts exp :: LExp
exp alts :: [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 _ i :: Name
i n :: [Name]
n ns :: e
ns) : alts :: [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 i :: Name
i (LConCase _ j :: Name
j n :: [Name]
n ns :: e
ns) = Name
i Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
j
    notTag _ _ = Bool
True
dropDups (c :: LAlt' e
c : alts :: [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 _) (LConCase i :: Int
i con :: Name
con args :: [Name]
args rhs :: 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 var :: Name
var) (LDefaultCase (LCase ty :: CaseType
ty (LV var' :: Name
var') alts :: [LAlt]
alts))
    | Name
var Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
var' = [LAlt]
alts
replaceInAlt exp :: LExp
exp a :: LAlt
a = [LAlt
a]

replaceExp :: LExp -> LExp -> LExp -> LExp
replaceExp :: LExp -> LExp -> LExp -> LExp
replaceExp (LCon _ t :: Int
t n :: Name
n args :: [LExp]
args) new :: LExp
new (LCon _ t' :: Int
t' n' :: Name
n' args' :: [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 _ t :: Int
t n :: Name
n args :: [LExp]
args) new :: LExp
new (LApp _ (LV n' :: Name
n') args' :: [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 old :: LExp
old new :: LExp
new tm :: LExp
tm = LExp
tm

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

getRHS :: LAlt' p -> p
getRHS (LConCase i :: Int
i n :: Name
n es :: [Name]
es rhs :: p
rhs) = p
rhs
getRHS (LConstCase _ rhs :: p
rhs) = p
rhs
getRHS (LDefaultCase rhs :: p
rhs) = p
rhs

getLams :: [LExp] -> [Name]
getLams [] = []
getLams (LLam args :: [Name]
args tm :: LExp
tm : cs :: [LExp]
cs) = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
args [LExp]
cs
getLams (LLet n :: Name
n val :: LExp
val exp :: LExp
exp : cs :: [LExp]
cs) = [LExp] -> [Name]
getLams (LExp
exp LExp -> [LExp] -> [LExp]
forall a. a -> [a] -> [a]
: [LExp]
cs)
getLams _ = []

getLamPrefix :: [Name] -> [LExp] -> [Name]
getLamPrefix as :: [Name]
as [] = [Name]
as
getLamPrefix as :: [Name]
as (LLam args :: [Name]
args tm :: LExp
tm : cs :: [LExp]
cs)
    | [Name] -> 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 (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 as :: [Name]
as (LLet n :: Name
n val :: LExp
val exp :: LExp
exp : cs :: [LExp]
cs) = [Name] -> [LExp] -> [Name]
getLamPrefix [Name]
as (LExp
exp LExp -> [LExp] -> [LExp]
forall a. a -> [a] -> [a]
: [LExp]
cs)
getLamPrefix as :: [Name]
as (_ : cs :: [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 tc :: Bool
tc a :: LExp
a es :: [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 n :: Name
n es :: [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 e :: LExp
e) = LExp -> LExp
LLazyExp (LExp -> LExp
eta LExp
e)
eta (LForce e :: LExp
e) = LExp -> LExp
LForce (LExp -> LExp
eta LExp
e)
eta (LLet n :: Name
n val :: LExp
val sc :: LExp
sc) = Name -> LExp -> LExp -> LExp
LLet Name
n (LExp -> LExp
eta LExp
val) (LExp -> LExp
eta LExp
sc)
eta (LLam args :: [Name]
args (LApp tc :: Bool
tc f :: LExp
f args' :: [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 args :: [Name]
args e :: LExp
e) = [Name] -> LExp -> LExp
LLam [Name]
args (LExp -> LExp
eta LExp
e)
eta (LProj e :: LExp
e i :: Int
i) = LExp -> Int -> LExp
LProj (LExp -> LExp
eta LExp
e) Int
i
eta (LCon a :: Maybe Name
a t :: Int
t n :: Name
n es :: [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 ct :: CaseType
ct e :: LExp
e alts :: [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 (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LExp -> LExp
eta) [LAlt]
alts)
eta (LOp f :: PrimFn
f es :: [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 tm :: LExp
tm = LExp
tm