A more strongly typed CEK machine

CEK.hs

{-# Language StrictData #-}
{-# Language GADTs #-}
{-# Language DeriveTraversable #-}
{-# Language LambdaCase #-}
module CEK where

import Control.Monad (ap)
import Data.Maybe
import Data.Void

-- C   -- Control
-- E   -- Environment
-- (S) -- Store
-- K   -- Continuation

data Exp a
  = Var a
  | Lam (Exp (Maybe a))
  | Ap (Exp a) (Exp a)
  deriving (Show, Functor, Foldable, Traversable)

instance Applicative Exp where
  pure = Var
  (<*>) = ap

instance Monad Exp where
  return = Var
  Var a >>= f = f a
  Ap l r >>= f = Ap (l >>= f) (r >>= f)
  Lam b >>= f = Lam $ b >>= \case
    Nothing -> Var Nothing
    Just a -> Just <$> f a

abstract :: (Functor f, Eq a) => a -> f a -> f (Maybe a)
abstract a = fmap go where
  go b
    | a == b = Nothing
    | otherwise = Just b

lam :: Eq a => a -> Exp a -> Exp a
lam a b = Lam (abstract a b)

closed :: Exp a -> Exp b
closed = fromJust . traverse (const Nothing)

newtype Env a = Env { (!) :: a -> Value }
-- instance Contravariant Env

instance Show (Env a) where
  show _ = "Env"

data Value where
  Closure :: Show a => Exp (Maybe a) -> Env a -> Value

data Kont where
  Top :: Kont
  Arg :: Show a => Exp a -> Env a -> Kont -> Kont
  Fun :: Show a => Exp (Maybe a) -> Env a -> Kont -> Kont

instance Show Kont where
  showsPrec d Top = showString "Top"
  showsPrec d (Arg c e k) = showParen (d > 10) $
    showString "Arg " . showsPrec 11 c . showChar ' ' . showsPrec 11 e . showChar ' ' . showsPrec 11 k
  showsPrec d (Fun b e k) = showParen (d > 10) $
    showString "Fun " . showsPrec 11 b . showChar ' ' . showsPrec 11 e . showChar ' ' . showsPrec 11 k

data State where 
  State :: Show a => Exp a -> Env a -> Kont -> State

instance Show State where
  showsPrec d (State c e k) = showParen (d > 10) $ 
    showString "State " . showsPrec 11 c . showChar ' ' . showsPrec 11 e . showChar ' ' . showsPrec 11 k

start :: Exp Void -> State
start c = State c (Env absurd) Top

id_ :: Exp Void
id_ = closed $ lam "x" $ Var "x"

const_ :: Exp Void
const_ = closed $ lam "x" $ lam "y" $ Var "x"

-- small-step semantics step
step :: State -> State
step s@(State c e k) = case c of
  Var v -> case e ! v of
    Closure b e' -> State (Lam b) e' k
  Ap cf cx -> State cf e (Arg cx e k)
  Lam b -> case k of
    Top -> s
    Arg cx e' k' -> State cx e' (Fun b e k')
    Fun b' e' k' -> State b' (extend (Closure b e) e') k'

extend :: Value -> Env a -> Env (Maybe a)
extend v (Env f) = Env $ maybe v f

final :: State -> Bool
final (State Lam{} _ Top) = True
final _ = False

-- until :: (a -> Bool) -> (a -> a) -> a -> a

eval :: State -> State
eval = until final step 

-- big-step semantics

big :: Show a => Exp a -> Env a -> Kont -> State
big c e k = case c of
  Var v -> case e ! v of
    Closure b e' -> big (Lam b) e' k
  Ap cf cx -> big cf e (Arg cx e k)
  Lam b -> case k of
    Top -> State c e k
    Arg cx e' k' -> big cx e' (Fun b e k')
    Fun b' e' k' -> big b' (extend (Closure b e) e') k'