{- 
-- 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 : map Just list


-- show for functions with listable domain

instance (Show a, Listable a, Show b) => Show (a -> b) where
   show f = show (map (\ 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 (Eq, 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


-- Krivine's machine

-- Again "forall" below is exists.

data State b = forall a . (Show a, Listable a) 
                         => State (LamX a) (a -> 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"
-}

instance Show b => Show (State b) where
  show (State t g us) = "Term: " ++ show t ++ "\nSubst: " 
                 ++ show g ++ "\nStack: " ++ show us ++ "\n"

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

final :: State b -> Bool
final (State (VarX a) g _ ) = case g a of
                                 VarX _     -> True
                                   -- in this case g is actually return
                                 SubX _ _ -> False
final (State (AbsX _) _ []) = 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 print 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 (State (VarX a) g us) = app (g a) us 
                              where app t [] = t
                                    app t (u : us) = app (AppX t u) us
                                   {- in this case g is actually return,
                                      so g a = VarX a,
                                      but we cannot write VarX a,
                                      as this has type LamX a
                                      where a is the Skolem constant for 
                                      the bound type variable of exists
                                      and we need a term of type LamX b
                                    -}
collect (State (AbsX t) g []) = SubX (AbsX t) g

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

whnfXIO :: (Show b, Listable b) => LamX b -> IO (LamX b)
whnfXIO t = do s <- runIO (initialize t)
               putStr "WhnfX: "
               let t' = collect s
               print t'
               return t'

whnfIO :: (Show b, Listable b) => Lam b -> IO (Lam b)
whnfIO t = do t' <- whnfXIO (embed t)
              putStr "Whnf: "
              let t'' = resolve t'
              print t''
              return 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)))))