Latest Cesk kappa calculus interpreter

kappacesk.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Main (main) where

import Control.Applicative
import Control.Monad
import Control.Monad.State.Strict hiding (State (..), state)
import qualified Control.Monad.State.Strict as State
import Data.Array.Unboxed (UArray)
import qualified Data.Array.Unboxed as Array
import Data.Kind
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Type.Equality
import Type.Reflection
import Prelude hiding (fst, id, snd)

-- Strict everything is used mostly to show that semantics don't rely
-- on laziness. Performance isn't really a worry. A similar thing with
-- data families applies. Data families are mostly used to enforce we
-- don't rely on discriminating on a GADT.

sample :: Hoas k => k Unit Z
sample =
  push (z 5)
    >>> pop
      ( \x ->
          push x
            >>> push (z 4)
            >>> add
      )

main :: IO ()
main = do
  loop (inject sample) 0

  let y = eval (inject sample)
  display y

display :: KnownT a => State a -> IO ()
display (State control env store val kont) = do
  putStrLn $ "\tip\t" ++ show control
  putStrLn $ "\tenv\t" ++ show env
  putStrLn $ "\theap\t" ++ show store
  putStrLn $ "\tstack"
  putStrLn $ "\t\t" ++ show val ++ " ⊢"
  prettyPrint kont

-- fixme.. find a better method of halting
loop :: KnownT a => State a -> Int -> IO ()
loop x n = do
  putStrLn $ show n ++ ":"
  display x

  case step x of
    Just y -> loop y (n + 1)
    Nothing -> do
      putStrLn "STUCK"

prettyPrint :: (KnownT a, KnownT b) => Path a b -> IO ()
prettyPrint IdPath = pure ()
prettyPrint (ComposePath h t) = do
  putStrLn $ "\t\t" ++ show h
  prettyPrint t

----------

data T = Unit | T :*: T | Z

infixr 0 :*:

class KnownT a where
  inferTag :: Tagged t => t a

class Tagged t where
  zTag :: t Z
  unitTag :: t Unit
  pairTag :: (KnownT a, KnownT b) => t a -> t b -> t (a :*: b)

instance Tagged TypeRep where
  unitTag = typeRep
  zTag = typeRep
  pairTag a b = a `withTypeable` (b `withTypeable` typeRep)

instance KnownT Unit where
  inferTag = unitTag

instance KnownT Z where
  inferTag = zTag

instance (KnownT a, KnownT b) => KnownT (a :*: b) where
  inferTag = pairTag inferTag inferTag

----------

class Category k where
  id :: KnownT a => k a a
  (>>>) :: (KnownT a, KnownT b, KnownT c) => k a b -> k b c -> k a c

class Category k => Common k where
  bang :: KnownT a => k a Unit

  push :: (KnownT a, KnownT b) => k Unit a -> k b (a :*: b)

  z :: Int -> k Unit Z
  add :: k (Z :*: Z :*: Unit) Z

class Common k => Hoas k where
  pop :: (KnownT a, KnownT b, KnownT c) => (k Unit a -> k b c) -> k (a :*: b) c

class Common k => Named k where
  load :: KnownT a => Var a -> k Unit a
  popVar :: (KnownT a, KnownT b, KnownT c) => Var a -> k b c -> k (a :*: b) c

-----------------

type StateMonad x = State.State x

newtype Namer k a b = P {name :: StateMonad Layout (k a b)}

instance Category k => Category (Namer k) where
  id = P (pure id)
  P f >>> P g = P (pure (>>>) <*> f <*> g)

instance Common k => Common (Namer k) where
  bang = P (pure bang)
  push (P x) = P (pure push <*> x)

  z n = P (pure (z n))
  add = P (pure add)

instance Named k => Hoas (Namer k) where
  pop f = P $ do
    v <- fresh
    y <- name (f (P $ pure (load v)))
    pure (popVar v y)

fresh :: KnownT a => StateMonad Layout (Var a)
fresh = fresh' inferTag

fresh' :: KnownT a => TypeRep a -> StateMonad Layout (Var a)
fresh' t = do
  l <- get
  let (v, mp') = Map.insertLookupWithKey (\_ -> (+)) (TypeKey t) 1 (vars l)
  put (Layout mp')
  let n = case v of
        Just x -> x
        Nothing -> 0
  pure (Var n)

---

data TypeKey = forall x. KnownT x => TypeKey (TypeRep x)

instance Eq TypeKey where
  TypeKey x == TypeKey y = SomeTypeRep x == SomeTypeRep y

instance Ord TypeKey where
  compare (TypeKey x) (TypeKey y) = compare (SomeTypeRep x) (SomeTypeRep y)

instance Show TypeKey where
  show (TypeKey x) = show x

-----

data State c = forall a b v env.
  (Env env, KnownT a, KnownT b) =>
  State
  { control :: Control env a b,
    env :: env,
    store :: Store,
    arg :: Val a,
    kont :: Path b c
  }

data family Val (a :: T)

data instance Val Unit = BangVal

newtype instance Val Z = ZVal Int

data instance Val (a :*: b) = TupleVal (Val a) (Val b)

-- kind of like pure ?
inject :: KnownT a => (forall k. Hoas k => k Unit a) -> State a
inject c = case runState (name c) emptyLayout of
  (c', layout) -> State c' (allocEnv layout) emptyStore BangVal id

eval :: KnownT a => State a -> State a
eval x = case step x of
  Nothing -> x
  Just y -> eval y

step :: (KnownT a, MonadFail m) => State a -> m (State a)
step (State control env store x k) = step' control env store x k

----

newtype Push a = Push { push' :: forall e b. (Env e, KnownT a, KnownT b) => Control e Unit a -> Control e b (a :*: b) }
instance Tagged Push where
  unitTag = Push {
    push' = \_ -> Control
      { fold = push bang,
        step' = \env store value kont ->
          pure (State bang env store BangVal (pushK value >>> kont))
      }
  }
  zTag = Push {
    push' = pushGeneric
  }
  pairTag _ _ = Push {
    push' = pushGeneric
  }

pushGeneric :: (Env env, KnownT a, KnownT b) => Control env Unit a -> Control env b (a :*: b)
pushGeneric x =
    Control
      { fold = push (fold x),
        step' = \env store value kont ->
          pure (State x env store BangVal (pushK value >>> kont))
      }
instance Category (Control env) where
  id =
    Control
      { fold = id,
        step' = \_ store value kont ->
          applyKont kont value store
      }
  f >>> g =
    Control
      { fold = fold f >>> fold g,
        step' = \env store param kont ->
          pure (State f env store param (composeK env g >>> kont))
      }

instance Env env => Common (Control env) where
  bang =
    Control
      { fold = bang,
        step' = \_ store _ kont ->
          applyKont kont BangVal store
      }
  push = pushGeneric

  z n =
    Control
      { fold = z n,
        step' = \_ store _ kont ->
          applyKont kont (ZVal n) store
      }
  add =
    Control
      { fold = add,
        step' = \_ store (TupleVal (ZVal x) (TupleVal (ZVal y) _)) kont ->
          applyKont kont (ZVal (x + y)) store
      }

instance Env env => Named (Control env) where
  load v =
    Control
      { fold = load v,
        step' = \env store _ kont -> do
          let Just val = getEnv env v
          applyKont kont val store
      }

  popVar v body =
    Control
      { fold = popVar v (fold body),
        step' = \env store (TupleVal h t) kont ->
          pure (State body (putEnv env v h) store t kont)
      }

-----

class Show env => Env env where
  getEnv :: KnownT a => env -> Var a -> Maybe (Val a)
  putEnv :: KnownT a => env -> Var a -> Val a -> env

data EmptyEnv = EmptyEnv
instance Show EmptyEnv where
  show EmptyEnv = "!"
instance Env EmptyEnv where
  getEnv EmptyEnv _ = Nothing
  putEnv EmptyEnv _ _ = EmptyEnv

emptyEnv :: EmptyEnv
emptyEnv = EmptyEnv

data SomeEnv s = SomeEnv
  { zValues :: ValArray Z,
    tupleMap :: Map TypeKey SomeArray
  }

instance Env (SomeEnv s) where
  getEnv = getEnv' inferTag
  putEnv = putEnv' inferTag

data SomeArray = forall a. KnownT a => SomeArray (ValArray a)

newtype GetEnv s a = GetEnv { getEnv' :: KnownT a => SomeEnv s -> Var a -> Maybe (Val a) }

instance Tagged (GetEnv s) where
  unitTag = GetEnv $ \_ _ -> pure BangVal
  zTag = GetEnv $ \env n -> getValArray (zValues env) n
  pairTag _ _ = GetEnv getEnvGeneric

newtype PutEnv s a = PutEnv { putEnv' :: KnownT a => SomeEnv s -> Var a -> Val a -> SomeEnv s }

instance Tagged (PutEnv s) where
  unitTag = PutEnv $ \env _ -> pure env
  zTag = PutEnv $ \env n x -> env { zValues = putValArray (zValues env) n x }
  pairTag _ _ = PutEnv putEnvGeneric

allocEnv :: Layout -> SomeEnv ()
allocEnv (Layout mp0) =
  let zs = case Map.lookup (TypeKey zTag) mp0 of
        Nothing -> 0
        Just n -> n
      mp1 = Map.delete (TypeKey unitTag) mp0
      mp2 = Map.delete (TypeKey zTag) mp1
      mp3 = Map.mapWithKey (\(TypeKey t) n -> SomeArray (doAlloc t n)) mp2

      doAlloc :: KnownT a => TypeRep a -> Int -> ValArray a
      doAlloc _ n = allocArray n
   in SomeEnv { zValues = allocArray zs, tupleMap = mp3 }

getEnvGeneric :: KnownT a => SomeEnv s -> Var a -> Maybe (Val a)
getEnvGeneric = getEnvGeneric' inferTag

-- fixme... handle ... maybes better
putEnvGeneric :: KnownT a =>  SomeEnv s -> Var a -> Val a -> SomeEnv s
putEnvGeneric env v x = case putEnvGeneric' inferTag env v x of
  Nothing -> undefined
  Just y -> y

getEnvGeneric' :: KnownT a => TypeRep a -> SomeEnv s -> Var a -> Maybe (Val a)
getEnvGeneric' t mp n = do
  SomeArray arr <- Map.lookup (TypeKey t) (tupleMap mp)
  arr' <- fromArray arr
  val <- getValArray arr' n
  pure val

-- fixme... lookup then insert is dumb, do one pass
putEnvGeneric' :: KnownT a => TypeRep a -> SomeEnv s -> Var a -> Val a -> Maybe (SomeEnv s)
putEnvGeneric' t env n x = do
  let t' = TypeKey t
  let mp = tupleMap env
  SomeArray arr <- Map.lookup t' mp
  arr' <- fromArray arr
  let arr'' = putValArray arr' n x
  let mp' = Map.insert t' (SomeArray arr'') mp
  pure (env {tupleMap = mp'})

fromArray :: (KnownT a, KnownT b) => ValArray a -> Maybe (ValArray b)
fromArray = fromArray' inferTag inferTag

fromArray' :: TypeRep a -> TypeRep b -> f a -> Maybe (f b)
fromArray' t t' arr = do
  Refl <- testEquality t t'
  pure arr

instance Show (SomeEnv s) where
  showsPrec _ env = ("{" ++) . shows (zValues env) . ("," ++) . shows (Map.assocs (tupleMap env)) . ("}" ++)

data family ValArray (a :: T)

data instance ValArray Unit = EmptyArray

newtype instance ValArray Z = ZArray (UArray Int Int)

data instance ValArray (a :*: b) = PairArray (ValArray a) (ValArray b)

allocArray :: KnownT a => Int -> ValArray a
allocArray = alloc inferTag

newtype Alloc a = Alloc { alloc :: Int -> ValArray a }
instance Tagged Alloc where
  unitTag = Alloc $ \n -> EmptyArray
  zTag = Alloc $ \n -> ZArray (Array.listArray (0, n) (repeat (0)))
  pairTag a b = Alloc $ \n -> PairArray (alloc a n) (alloc b n)

newtype GetVal a = GetVal { getVal :: ValArray a -> Int -> Maybe (Val a) }
instance Tagged GetVal where
  unitTag = GetVal $ \EmptyArray n -> pure BangVal
  zTag = GetVal $ \(ZArray arr) n -> do
          let val = arr Array.! n
          pure (ZVal val)
  pairTag a b = GetVal $ \(PairArray x y) n -> do
          x' <- getVal a x n
          y' <- getVal b y n
          pure (TupleVal x' y')

newtype PutVal a = PutVal { putVal :: ValArray a -> Int -> Val a -> ValArray a }
instance Tagged PutVal where
  unitTag = PutVal $ \EmptyArray _ BangVal ->
                       EmptyArray
  zTag = PutVal $ \(ZArray arr) n (ZVal v) ->
                    ZArray (arr Array.// [(n, v)])
  pairTag a b = PutVal $ \(PairArray x y) n (TupleVal l r) ->
                           PairArray (putVal a x n l) (putVal b y n r)

newtype ShowArray a = ShowArray { showArray :: ValArray a -> String -> String }
instance Tagged ShowArray where
  unitTag = ShowArray $ \EmptyArray -> ("!" ++)
  zTag = ShowArray $ \(ZArray arr) ->
          shows (Array.elems arr)
  pairTag a b = ShowArray $ \(PairArray x y) ->
    ("<" ++) . showArray a x . (", " ++) . showArray b y . (">" ++)

getValArray :: KnownT a => ValArray a -> Var n -> Maybe (Val a)
getValArray arr (Var n) = getVal inferTag arr n

putValArray :: KnownT a => ValArray a -> Var a -> Val a -> ValArray a
putValArray arr (Var n) x = putVal inferTag arr n x

instance KnownT a => Show (ValArray a) where
  showsPrec _ = showArray inferTag

instance Show SomeArray where
  showsPrec _ (SomeArray arr) = showArray inferTag arr

data Store = Store

instance Show Store where
  show Store = "{}"

emptyStore :: Store
emptyStore = Store

data Var (a :: T) = Var Int

instance KnownT a => Show (Var a) where
  showsPrec _ (Var n) = ("#" ++) . shows n

newtype Layout = Layout {vars :: Map TypeKey Int}

emptyLayout :: Layout
emptyLayout = Layout {vars = Map.empty}

class Category k => K k where
  pushK :: (KnownT a, KnownT b) => Val a -> k b (b :*: a)
  composeK :: (Env env, KnownT a, KnownT b) => env -> Control env a b -> k a b

data Kont a b = Kont
  { foldK :: forall k. K k => k a b,
    applyKont' :: forall m. (KnownT a, KnownT b, MonadFail m) => Val a -> Store -> m (State b)
  }

data Path a b where
  IdPath :: KnownT a => Path a a
  ComposePath :: (KnownT a, KnownT b, KnownT c) => Kont a b -> Path b c -> Path a c

instance Category Path where
  id = IdPath
  IdPath >>> f = f
  ComposePath f g >>> h = ComposePath f (g >>> h)

path :: (KnownT a, KnownT b) => Kont a b -> Path a b
path x = ComposePath x IdPath

instance K Path where
  pushK h =
    path $
      Kont
        { foldK = pushK h,
          applyKont' = \value store ->
            pure (State id emptyEnv store (TupleVal value h) id)
        }
  composeK env f =
    path $
      Kont
        { foldK = composeK env f,
          applyKont' = \value store ->
            pure (State f env store value id)
        }

applyKont :: (KnownT a, KnownT b, MonadFail m) => Path a b -> Val a -> Store -> m (State b)
applyKont k value store = case k of
  IdPath -> fail "halt"
  ComposePath h t -> do
    -- this part doesn't seem quite right to me but IDK
    State c e s x k <- applyKont' h value store
    pure (State c e s x (k >>> t))

instance Show (Kont a b) where
  showsPrec _ x = view (foldK x) 0

instance KnownT a => Show (Val a) where
  showsPrec _ = showValue inferTag

newtype ShowValue a = ShowValue {showValue :: Val a -> String -> String}

instance Tagged ShowValue where
  unitTag = ShowValue $ \BangVal -> ("!" ++)
  zTag = ShowValue $ \(ZVal n) -> shows n
  pairTag a b = ShowValue $ \(TupleVal x y) ->
    ("<" ++)
      . showValue a x
      . (", " ++)
      . showValue b y
      . (">" ++)

data Control env a b = Control
  { fold :: forall k. Named k => k a b,
    step' :: forall m c. (Env env, KnownT a, KnownT b, KnownT c, MonadFail m) => env -> Store -> Val a -> Path b c -> m (State c)
  }

instance Show env => Show (Control env a b) where
  showsPrec _ control = view (fold control) 0

newtype View (a :: T) (b :: T) = V {view :: Int -> String -> String}

instance K View where
  pushK x = V $ \p ->
    showParen (p > 10) $ ("push " ++) . shows x
  composeK env f = V $ \p ->
    ("<" ++) . shows f . (", " ++) . shows env . (">" ++)

instance Category View where
  id = V $ \p -> ("id" ++)
  g >>> f = V $ \p ->
    showParen (p > 9) $
      view f 10 . (" ∘ " ++) . view g 10

instance Common View where
  bang = V $ \p -> ("!" ++)
  push x = V $ \p ->
    showParen (p > 10) $ ("push " ++) . view x 11

  z n = V $ \_ -> shows n
  add = V $ \_ -> ("add" ++)

instance Named View where
  load v = V $ \_ -> shows v
  popVar v body = V $ \p ->
    showParen (p > 11) $
      ("κ " ++) . shows v . (". " ++) . view body 0