{-# LANGUAGE OverloadedStrings #-}

module A.E ( M, nN, eta ) where

import           A
import           C
import           Control.Monad                    ((<=<))
import           Control.Monad.Trans.State.Strict (State, state)
import qualified Data.Text                        as T
import           Nm
import           U

type M = State Int

nN :: T.Text -> a -> M (Nm a)
nN :: forall a. Text -> a -> M (Nm a)
nN Text
n a
l = (Int -> (Nm a, Int)) -> StateT Int Identity (Nm a)
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state (\Int
i -> (Text -> U -> a -> Nm a
forall a. Text -> U -> a -> Nm a
Nm Text
n (Int -> U
U(Int -> U) -> Int -> U
forall a b. (a -> b) -> a -> b
$Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) a
l, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))

doms :: T -> [T]
doms :: T -> [T]
doms (TyArr T
t T
t') = T
tT -> [T] -> [T]
forall a. a -> [a] -> [a]
:T -> [T]
doms T
t'; doms T
_ = []

cLam :: E a -> Int
cLam :: forall a. E a -> Int
cLam (Lam a
_ Nm a
_ E a
e) = Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+E a -> Int
forall a. E a -> Int
cLam E a
e; cLam E a
_ = Int
0

tuck :: E a -> (E a -> E a, E a)
tuck :: forall a. E a -> (E a -> E a, E a)
tuck (Lam a
l Nm a
n E a
e) = let (E a -> E a
f, E a
e') = E a -> (E a -> E a, E a)
forall a. E a -> (E a -> E a, E a)
tuck E a
e in (a -> Nm a -> E a -> E a
forall a. a -> Nm a -> E a -> E a
Lam a
l Nm a
n(E a -> E a) -> (E a -> E a) -> E a -> E a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.E a -> E a
f, E a
e'); tuck E a
e = (E a -> E a
forall a. a -> a
id, E a
e)

unseam :: [T] -> M (E T -> E T, E T -> E T)
unseam :: [T] -> M (E T -> E T, E T -> E T)
unseam [T]
ts = do
    lApps <- (T -> M (E T -> E T, E T -> E T))
-> [T] -> StateT Int Identity [(E T -> E T, E T -> E T)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\T
t -> do {n <- Text -> T -> M (Nm T)
forall a. Text -> a -> M (Nm a)
nN Text
"x" T
t; pure (\E T
e' -> let t' :: T
t' = E T -> T
forall a. E a -> a
eLoc E T
e' in T -> Nm T -> E T -> E T
forall a. a -> Nm a -> E a -> E a
Lam (T -> T -> T
TyArr T
t T
t') Nm T
n E T
e', \E T
e' -> let TyArr T
_ T
cod = E T -> T
forall a. E a -> a
eLoc E T
e' in T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
cod E T
e' (T -> Nm T -> E T
forall a. a -> Nm a -> E a
Var T
t Nm T
n))}) [T]
ts
    let (ls, eApps) = unzip lApps
    pure (thread ls, thread (reverse eApps))
    where thread :: [b -> b] -> b -> b
thread = ((b -> b) -> (b -> b) -> b -> b) -> (b -> b) -> [b -> b] -> b -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (b -> b) -> (b -> b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) b -> b
forall a. a -> a
id

mkLam :: [T] -> E T -> M (E T)
mkLam :: [T] -> E T -> M (E T)
mkLam [T]
ts E T
e = do
    (lam, app) <- [T] -> M (E T -> E T, E T -> E T)
unseam [T]
ts
    pure $ lam (app e)
-- fmap (($e).(uncurry (.))) (unseam ts)

eta :: E T -> M (E T)
eta = E T -> M (E T)
eM (E T -> M (E T)) -> (E T -> M (E T)) -> E T -> M (E T)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< E T -> M (E T)
eO

eM :: E T -> M (E T)
eM :: E T -> M (E T)
eM (EApp T
t ho :: E T
ho@(BB T
_ BBin
Map) E T
op)     = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(BB T
_ BBin
Filter) E T
op)  = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(BB T
_ BBin
Prior) E T
op)   = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(BB T
_ BBin
DedupOn) E T
op) = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(BB T
_ BBin
Fold1) E T
op)   = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(TB T
_ BTer
Fold) E T
op)    = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(TB T
_ BTer
Scan) E T
op)    = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t ho :: E T
ho@(TB T
_ BTer
ZipW) E T
op)    = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t E T
ho (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eta E T
op
eM (EApp T
t E T
e0 E T
e1)                = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
EApp T
t (E T -> E T -> E T) -> M (E T) -> StateT Int Identity (E T -> E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eM E T
e0 StateT Int Identity (E T -> E T) -> M (E T) -> M (E T)
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
<*> E T -> M (E T)
eM E T
e1
eM (Cond T
t E T
p E T
e0 E T
e1)              = T -> E T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a -> E a
Cond T
t (E T -> E T -> E T -> E T)
-> M (E T) -> StateT Int Identity (E T -> E T -> E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eM E T
p StateT Int Identity (E T -> E T -> E T)
-> M (E T) -> StateT Int Identity (E T -> E T)
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
<*> E T -> M (E T)
eM E T
e0 StateT Int Identity (E T -> E T) -> M (E T) -> M (E T)
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
<*> E T -> M (E T)
eM E T
e1
eM (OptionVal T
t Maybe (E T)
e)               = T -> Maybe (E T) -> E T
forall a. a -> Maybe (E a) -> E a
OptionVal T
t (Maybe (E T) -> E T)
-> StateT Int Identity (Maybe (E T)) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (E T -> M (E T))
-> Maybe (E T) -> StateT Int Identity (Maybe (E T))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse E T -> M (E T)
eM Maybe (E T)
e
eM (Implicit T
t E T
e)                = T -> E T -> E T
forall a. a -> E a -> E a
Implicit T
t (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eM E T
e
eM (Lam T
t Nm T
n E T
e)                   = T -> Nm T -> E T -> E T
forall a. a -> Nm a -> E a -> E a
Lam T
t Nm T
n (E T -> E T) -> M (E T) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eM E T
e
eM (Guarded T
t E T
p E T
e)               = T -> E T -> E T -> E T
forall a. a -> E a -> E a -> E a
Guarded T
t (E T -> E T -> E T) -> M (E T) -> StateT Int Identity (E T -> E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> E T -> M (E T)
eM E T
p StateT Int Identity (E T -> E T) -> M (E T) -> M (E T)
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
<*> E T -> M (E T)
eM E T
e
eM (Tup T
t [E T]
es)                    = T -> [E T] -> E T
forall a. a -> [E a] -> E a
Tup T
t ([E T] -> E T) -> StateT Int Identity [E T] -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (E T -> M (E T)) -> [E T] -> StateT Int Identity [E T]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse E T -> M (E T)
eM [E T]
es
eM (Rec T
t [(Nm T, E T)]
es)                    = T -> [(Nm T, E T)] -> E T
forall a. a -> [(Nm a, E a)] -> E a
Rec T
t ([(Nm T, E T)] -> E T)
-> StateT Int Identity [(Nm T, E T)] -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Nm T, E T) -> StateT Int Identity (Nm T, E T))
-> [(Nm T, E T)] -> StateT Int Identity [(Nm T, E T)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((E T -> M (E T)) -> (Nm T, E T) -> StateT Int Identity (Nm T, E T)
forall (m :: * -> *) b b' a.
Functor m =>
(b -> m b') -> (a, b) -> m (a, b')
secondM E T -> M (E T)
eM) [(Nm T, E T)]
es
eM (Anchor T
t [E T]
es)                 = T -> [E T] -> E T
forall a. a -> [E a] -> E a
Anchor T
t ([E T] -> E T) -> StateT Int Identity [E T] -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (E T -> M (E T)) -> [E T] -> StateT Int Identity [E T]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse E T -> M (E T)
eM [E T]
es
eM (Arr T
t Vector (E T)
es)                    = T -> Vector (E T) -> E T
forall a. a -> Vector (E a) -> E a
Arr T
t (Vector (E T) -> E T)
-> StateT Int Identity (Vector (E T)) -> M (E T)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (E T -> M (E T))
-> Vector (E T) -> StateT Int Identity (Vector (E T))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Vector a -> f (Vector b)
traverse E T -> M (E T)
eM Vector (E T)
es
eM (Let T
t (Nm T
n, E T
e') E T
e)             = do {e'𝜂 <- E T -> M (E T)
eM E T
e'; e𝜂 <- eM e; pure (Let t (n, e'𝜂) e𝜂)}
eM E T
e                             = E T -> M (E T)
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure E T
e

-- outermost
eO :: E T -> M (E T)
eO :: E T -> M (E T)
eO e :: E T
e@(Var t :: T
t@TyArr{} Nm T
_)    = [T] -> E T -> M (E T)
mkLam (T -> [T]
doms T
t) E T
e
eO e :: E T
e@(UB T
t BUn
_)             = [T] -> E T -> M (E T)
mkLam (T -> [T]
doms T
t) E T
e
eO e :: E T
e@(BB T
t BBin
_)             = [T] -> E T -> M (E T)
mkLam (T -> [T]
doms T
t) E T
e
eO e :: E T
e@(TB T
t BTer
_)             = [T] -> E T -> M (E T)
mkLam (T -> [T]
doms T
t) E T
e
eO e :: E T
e@(EApp t :: T
t@TyArr{} E T
_ E T
_) = [T] -> E T -> M (E T)
mkLam (T -> [T]
doms T
t) E T
e
eO e :: E T
e@(Lam t :: T
t@TyArr{} Nm T
_ E T
_)  = do
    let l :: Int
l = [T] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (T -> [T]
doms T
t)
        (E T -> E T
preL, E T
e') = E T -> (E T -> E T, E T)
forall a. E a -> (E a -> E a, E a)
tuck E T
e
    (lam, app) <- [T] -> M (E T -> E T, E T -> E T)
unseam (Int -> [T] -> [T]
forall a. Int -> [a] -> [a]
take (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-E T -> Int
forall a. E a -> Int
cLam E T
e) ([T] -> [T]) -> [T] -> [T]
forall a b. (a -> b) -> a -> b
$ T -> [T]
doms T
t)
    pure (lam (preL (app e')))
eO E T
e                      = E T -> M (E T)
forall a. a -> StateT Int Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure E T
e