{- 
-- defined in Prelude

class Monad m where
  return :: a -> m a
  (>>=) :: m a -> (a -> m b) -> m b

-}

-- empty type

data Void

instance Show Void


-- listability of Void and Maybe types

class Listable a where
  list :: [a]

instance Listable Void where
  list = []

instance Listable a => Listable (Maybe a) where
  list = Nothing : fmap Just list


-- show for functions with listable domain

instance (Show a, Listable a, Show b) => Show (a -> b) where
   show f = show (fmap (\ a -> (a, f a)) list)


-- a useful distributivity

lift :: Monad m => (a -> m b) -> Maybe a -> m (Maybe b)
lift f Nothing  = return Nothing
lift f (Just x) = f x >>= return . Just 

-- kleisli composition

(<.) :: Monad m => (b -> m c) -> (a -> m b) -> a -> m c
g <. f = \ a -> f a >>= g


-- lambda calculus syntax

data Lam a = Var a  
           | App (Lam a) (Lam a) | Abs (Lam (Maybe a)) deriving Show

instance Monad Lam where
  return a = Var a
  Var a >>= f   =  f a
  App t u >>= f = App (t >>= f) (u >>= f)
  Abs t >>= f   = Abs (t >>= lift f)



-- lambda calculus with explicit substitutions

{- 
The "forall" below is Haskell for exists. The motivation is that the desired 
type 

exists a . (Lam X a,  a -> LamX b) -> LamX b

of SubX is isomorphic to

forall a . LamX a -> (a -> LamX b) -> LamX b
-}

data LamX b = VarX b 
            | AppX (LamX b) (LamX b) | AbsX (LamX (Maybe b))
            | forall a . (Show a, Listable a) 
                            => SubX (LamX a) (a -> LamX b) 

-- Haskell cannot derive Show for datatypes with existially typed 
-- components

instance Show b => Show (LamX b) where
  showsPrec d (VarX b)   = showParen (d > 10) 
                             (showString "VarX " . showsPrec 11 b)
  showsPrec d (AppX t u) = showParen (d > 10) 
                             (showString "AppX " . showsPrec 11 t 
                                . showString " " . showsPrec 11 u)
  showsPrec d (AbsX t)   = showParen (d > 10) 
                             (showString "AbsX " . showsPrec 11 t)
  showsPrec d (SubX t f) = showParen (d > 10) 
                             (showString "SubX " . showsPrec 11 t 
                                . showString " " . showsPrec 11 f)
{-
  show (VarX b) = "VarX (" ++ show b ++ ")"
  show (AppX t u) = "(AppX (" ++ show t ++ " " ++ show u ++ ")"
  show (AbsX t) = "AbsX (" ++ show t ++ ")"
  show (SubX t f) = "SubX (" ++ show t ++ " " ++ show f ++ ")" 
-} 

instance Monad LamX where
  return a = VarX a
  VarX a >>= g     = g a
  AppX t u >>= g   = AppX (t >>= g) (u >>= g)
  AbsX t >>= g     = AbsX (t >>= lift g)
  SubX t f >>= g = SubX t (g <. f)

embed :: Lam a -> LamX a
embed (Var a)   = VarX a
embed (App t u) = AppX (embed t) (embed u)
embed (Abs t)   = AbsX (embed t)

resolve :: LamX a -> Lam a
resolve (VarX a)     = Var a
resolve (AppX t u)   = App (resolve t) (resolve u)
resolve (AbsX t)     = Abs (resolve t)
resolve (SubX t f) = resolve t >>= resolve . f


-- Eager machine with stack markers a la Curien

-- Again "forall" below is exists.

data State b = forall a . (Show a, Listable a) 
                            => State (LamX a) (a -> LamX b) 
                                     [Either (LamX b) (LamX b)]
             | State0 (LamX b) [Either (LamX b) (LamX b)]

printState :: Show b => State b -> IO ()
printState (State t g us) =           do putStr "Term: "
                                         print t
                                         putStr "Subst: " 
                                         print g
                                         putStr "Stack: "
                                         print us
                                         putStr "\n"
printState (State0 t us) =        do putStr "Term: "
                                     print t
                                     putStr "Stack: "
                                     print us
                                     putStr "\n"

trans :: (Show b, Listable b) => State b -> State b
trans (State (VarX a)   g us) = State0 (g a) us
trans (State (AppX t u) g us) = State t g (Left (SubX u g) : us)
trans (State (AbsX t)   g us) = State0 (SubX (AbsX t) g) us
trans (State (SubX t f) g us) = State t (g <. f) us
trans (State0 t (Left  (SubX u f) : us)) 
                              = State u f (Right t : us) 
trans (State0 t (Right (SubX (AbsX u) f) : us))
                              = State u (maybe t f) us

final :: State b -> Bool
final (State0 _ []) = True 
final (State0 _ (Right u : _)) = case u of 
                                          SubX (AbsX _) _ -> False
                                          _ -> True
final _ = False

run :: (Show b, Listable b) => State b -> State b
run state = if final state then state else run (trans state) 

runIO :: (Show b, Listable b) => State b -> IO (State b)
runIO state = do printState state
                 if final state then return state else runIO (trans state) 


-- using the machine for wh normalization 

initialize :: (Show b, Listable b) => LamX b -> State b
initialize t = State t return []

collect :: State b -> LamX b
collect (State0 t us) = app t us 
                              where app t [] = t
                                    app t (Left  u : us) = app (AppX t u) us
                                    app t (Right u : us) = app (AppX u t) us

whnfX :: (Show b, Listable b) => LamX b -> LamX b
whnfX  = collect . run . initialize

whnf :: (Show b, Listable b) => Lam b -> Lam b
whnf = resolve . whnfX . embed


printAndReturn :: Show a => a -> IO a
printAndReturn a = do print a
                      return a

whnfXIO :: (Show b, Listable b) => LamX b -> IO (LamX b)
whnfXIO t = do state <- runIO (initialize t)
               putStr "WhnfX: "
               printAndReturn (collect state)

whnfIO :: (Show b, Listable b) => Lam b -> IO (Lam b)
whnfIO t = do t' <- whnfXIO (embed t)
              putStr "Whnf: "
              printAndReturn (resolve t')



-- examples

instance Listable String where
  list = ["x", "y", "z"]

-- test1 = (\ x -> \ y -> y x) x y

test1 :: Lam String
test1 = App (App (Abs (Abs (App (Var Nothing) (Var (Just Nothing))))) 
        (Var "x")) (Var "y")

-- test2 = (\ x -> \ y -> z x) x y 

test2 :: Lam String
test2 = App (App (Abs (Abs (App (Var (Just (Just "z"))) (Var (Just Nothing))))) 
        (Var "x")) (Var "y")

-- y  = (\ x -> x x) (\ x -> x x)

y :: Lam a
y = App (Abs (App (Var Nothing) (Var Nothing))) (Abs (App (Var Nothing) (Var Nothing)))


-- k  = \ x -> \ y -> x

k :: Lam a
k = Abs (Abs (Var (Just Nothing)))

-- k' = \ x -> \ y -> y

k' :: Lam a
k' = Abs (Abs (Var Nothing))

-- s = \ x -> \ y -> \ z -> (x z) (y z) 

s :: Lam a
s = Abs (Abs (Abs (App  (App (Var (Just (Just Nothing))) (Var Nothing)) 
                 (App (Var (Just Nothing)) (Var Nothing)))))