{-# LANGUAGE FlexibleInstances, ImpredicativeTypes #-}

module NBE where

import Data.Maybe

data Term' u v = Free u
               | Bound v
               | Lam (v -> Term' u v)
               | Ap (Term' u v) (Term' u v)

type Term u = forall v. Term' u v

data Val' u v = VLam (Val' u v -> Val' u v)
              | N (Neut u v)

type Val u = forall v. Val' u v

data Neut u v = NAp (Neut u v) (Val' u v)
              | NFree u
              | Place v

data WVal' u v = WLam (WVal' u v -> WVal' u v)
               | WFree u
               | WAp (WVal' u v) (WVal' u v)
               | WPlace v

type WVal u = forall v. WVal' u v

eval' :: Term' u (Val' u v) -> Val' u v
eval' (Free u)  = N $ NFree u
eval' (Bound v) = v
eval' (Lam e)   = VLam (\v -> eval' $ e v)
eval' (Ap f x)  = eval' f @@ eval' x

eval :: Term u -> Val u
eval t = eval' t

(@@) :: Val' u v -> Val' u v -> Val' u v
VLam f  @@ x = f x
N f     @@ x = N $ NAp f x

trans :: Term' u (WVal' u v) -> WVal' u v
trans (Free  u) = WFree u
trans (Bound v) = v
trans (Lam e)   = WLam $ \wv -> trans . e $ wv
trans (Ap f x)  = WAp (trans f) (trans x)

reduce :: WVal' u v -> WVal' u v
reduce (WAp f x) = reduce f @@@ x
reduce t         = t

(@@@) :: WVal' u v -> WVal' u v -> WVal' u v
WLam f @@@ x = reduce $ f x
f      @@@ x = WAp f x

weval' :: Term' u (WVal' u v) -> WVal' u v
weval' = reduce . trans

weval :: Term u -> WVal u
weval t = weval' t

quote' :: Val' u (Term' u v) -> Term' u v
quote' (VLam f)  = Lam (quote' . f . N .Place . Bound)
quote' (N t)     = quoteN t

quoteN :: Neut u (Term' u v) -> Term' u v
quoteN (NAp f x) = quoteN f `Ap` quote' x
quoteN (NFree u) = Free u
quoteN (Place t) = t

quote :: Val u -> Term u
quote v = quote' v

normalize :: Term u -> Term u
normalize t = quote (eval t)

ns :: [String]
ns = ["x","y","z"] ++ [ c : show n | c <- "xyz" , n <- [1..] ]

display :: [String] -> Term' String String -> String
display _         (Free s)  = s
display _         (Bound s) = s
display (n:names) (Lam e)   = "\\" ++ n ++ " -> " ++ display names (e n)
display names     (Ap f x)  = "(" ++ display names f ++ ") (" 
                                  ++ display names x ++ ")"

displayV :: [String] -> Val' String String -> String
displayV (n:names) (VLam f) = "\\" ++ n ++ " -> " 
                                   ++ displayV names (f . N . Place $ n)
displayV names     (N t)    = displayN names t

displayN :: [String] -> Neut String String -> String
displayN names (NAp f x) = "(" ++ displayN names f ++ ") ("
                               ++ displayV names x ++ ")"
displayN _     (NFree u) = u
displayN _     (Place s) = s

displayW :: [String] -> WVal' String String -> String
displayW (n:names) (WLam f)   = "\\" ++ n ++ " -> "
                                     ++ displayW names (f . WPlace $ n)
displayW names     (WAp f x)  = "(" ++ displayW names f ++ ") ("
                                    ++ displayW names x ++ ")"
displayW names     (WFree u)  = u
displayW names     (WPlace s) = s

-- \x -> x
-- => \x -> x
ident1 :: Term u
ident1 = Lam (\x -> Bound x)

-- (\x -> x) (\x -> x)
-- => \x -> x
ident2 :: Term u
ident2 = Ap (Lam (\x -> Bound x)) (Lam (\x -> Bound x))

-- (\x f -> f x) (\x -> x)
-- => \f -> f (\x -> x)
neut :: Term u
neut = Ap (Lam (\x -> Lam (\f -> Ap (Bound f) (Bound x)))) (Lam (\x -> Bound x))

-- (\x -> x x) (\x -> x x)
-- => loop
w :: Term u
w = Ap (Lam (\x -> Ap (Bound x) (Bound x))) (Lam (\x -> Ap (Bound x) (Bound x)))

-- \f -> (\x -> f (x x)) (\x -> f (x x))
y :: Term u
y = Lam (\f -> Ap (Lam (\x -> Ap (Bound f) (Ap (Bound x) (Bound x))))
                  (Lam (\x -> Ap (Bound f) (Ap (Bound x) (Bound x)))))

-- \x y -> x
k :: Term u
k = Lam (\x -> Lam (\y -> Bound x))

-- \x -> y (k x)
-- =?> \x -> x
ident3 :: Term u
ident3 = Lam (\x -> Ap y (Ap k (Bound x)))

-- (\x f -> f x) ((\x -> x) (\x -> x))
-- => \f -> f (\x -> x)
-- ~> \f -> f ((\x -> x) (\x -> x))
neut2 = Ap (Lam $ \x -> Lam $ \f -> Bound f `Ap` Bound x) (ident1 `Ap` ident1)