CSK machine for lazy language with let-rec and Shift/Reset

Main.hs

{-# language LambdaCase #-}
{-# language BlockArguments #-}
{-# language GADTs #-}
{-# language DataKinds #-}
{-# language TypeApplications #-}
{-# language DuplicateRecordFields #-}

{-|

  This is a toy virtual machine for a simple call-by-need (Updatable) language.

  The program is encoded as a value of type `Program` and the evaluation result
  is of type `Value`.

-}
import Control.Monad.State (StateT, runStateT, lift)

import Data.Foldable (for_)
import Data.Traversable (for)
import Data.List (intercalate)

import Tuple
import Memory

-- | Source program.
--
data Program where
  Var :: { name :: Name } -> Program        -- | x

  App :: { f :: Program
         , x :: Program } -> Program        -- | f(x)

  Lam :: { arg  :: Name
         , body :: Program } -> Program     -- | (n) => prog

  Let :: { decls :: [(Name, Program)]
         , body  :: Program } -> Program    -- | let f = (x) => 1 + x; f (f x)

  PC  :: { c :: Constant } -> Program       -- | 1 OR (x, y) => x + y

  Rst :: { mark :: Int
         , body :: Program } -> Program     -- | reset(N, prog)

  Sht :: { mark :: Int
         , exit :: Name
         , body :: Program } -> Program     -- | shift(N, (k) => prog)

  Ctr :: { iD    :: Int
         , arity :: Int} -> Program         -- | True, Cons x y

  If  :: { iD      :: Int
         , subject :: Program
         , yes     :: ([Name], Program)
         , no      :: Program } -> Program  -- | if Cons ls ((x, y) => x) 0

infixl 1 `App`

-- | Constant.
--
data Constant where
  Num :: { c   :: Double }           -> Constant  -- | 1
  Bif :: { bif :: Value -> M Value } -> Constant  -- | (x, y) => x + y

-- | Evaluation result.
--
data Value where
  Clo :: { env  :: Context
         , arg  :: Name
         , body :: Program } -> Value         -- a function with context

  VC  :: { c :: Constant } -> Value           -- 1

  Seq :: { k :: [Kont] } -> Value             -- The delimited continuation

  Adt :: { iD    :: Int
         , arity :: Int
         , args  :: [Memory.Node] } -> Value  -- An ADT

-- | What to do next with current value.
--
data Kont where
  Call    :: { arg  :: Memory.Node } -> Kont         -- | current_value(Memory.Node)
  ArgBIF  :: { bif  :: Value -> M Value } -> Kont    -- | bif(current_value)
  Rewrite :: { cell :: Memory.Node } -> Kont         -- | Rewrite the node with current value.
  Prompt  :: { mark :: Int } -> Kont                 -- | The continuation delimiter.
  Select  :: { iD   :: Int
             , yes  :: ([Name], (Context, Program))
             , no   :: Memory.Node } -> Kont         -- | if-construct
  deriving Show

-- | A mutable cell with either value or program/context pair.
--
data Cell where
  OneShot   :: { env :: Context, program :: Program } -> Cell  -- | Non-updatable cell.
  Updatable :: { env :: Context, program :: Program } -> Cell  -- | Updatable cell.
  Ready     :: { result :: Value } -> Cell                  -- | The result.
  deriving Show

type Name = String

-- | Program environment.
--
type Context = [(Name, Memory.Node)]

-- | A monad to work in.
--
type M =
    StateT (Tuple
      [ Memory.Node    -- Current node being evaluated.
      , Memory.T Cell  -- Memory.
      , [Kont]         -- Stack of planned actions.
      , Bool           -- Stop-switch.
      ]
    ) IO

-- | Small-step semantics (evaluate one step).
--
step :: M Bool
step = do
  next \(src, cell, k) -> case (cell, k) of
    (v@Ready {}, Rewrite addr : k) -> do
      writeMem addr v
      return (addr, k)

    (Ready (Clo env name body), Call x : k) -> do
      addr' <- callClos name x env body
      return (addr', k)

    (Ready (Seq dk),       Call x : k) -> return (x, dk ++ k)
    (Ready (VC (Bif bif)), Call x : k) -> return (x, ArgBIF bif : k)

    (Ready v, ArgBIF bif : k) -> do
      v' <- callBif bif v
      return (v', k)

    (OneShot e (Var n), k) -> do
      n' <- findName n e
      return (n', k)

    (OneShot e (App f x), k) -> do
      (f', x') <- preApp e f x
      return (f', Call x' : k)

    (OneShot e (Lam n p), k) -> do
      v' <- mkClos e n p
      return (v', k)

    (OneShot _ (PC c), k) -> do
      v <- mkVal c
      return (v, k)

    (OneShot e (Let ds b), k) -> do
      b' <- pushDecls e ds b
      return (b', k)

    (OneShot e (Rst i p), k) -> do
      p' <- prog (e, p)
      return (p', Prompt i : k)

    (OneShot e (Sht i n p), k) -> do
      (p', k') <- shift e i n p k
      return (p', k')

    (OneShot e (Ctr c a), k) -> do
      p' <- mkAdt c a
      return (p', k)

    (OneShot e (If c x (ns, y) n), k) -> do
      (p', e') <- mkIf e x n
      return (p', Select c (ns, (e, y)) e' : k)

    (Ready (Adt c n s), Call x : k) | n > 0 -> do
      p' <- push c n s x
      return (p', k)

    (Ready (Adt c n s), Select b (ns, y) e : k) | n == 0 -> do
      p' <- match (c == b) ns s y e
      return (p', k)

    (Updatable e p, k) -> do
      p' <- prog (e, p)
      return (p', Rewrite src : k)

    (v, Prompt _ : k) -> do
      return (src, k)

    _ -> do
      stop
      return (src, k)

match :: Bool -> [Name] -> [Memory.Node] -> (Context, Program) -> Node -> M Node
match True ns xs (e, y) _ = prog (zip ns xs ++ e, y)
match False _ _ _ e = return e

push :: Int -> Int -> [Memory.Node] -> Node -> M Node
push c n xs x = val $ Adt c (n - 1) (x : xs)

mkIf :: Context -> Program -> Program -> M (Node, Node)
mkIf e x n = do
  p' <- prog (e, x)
  n' <- prog (e, n)
  return (p', n')

mkAdt :: Int -> Int -> M Node
mkAdt n a = do
  val $ Adt n a []

-- | Capture the continuation.
--
shift :: Context -> Int -> Name -> Program -> [Kont] -> M (Memory.Node, [Kont])
shift e i n p k = do
  let (f, Prompt _ : k') = break (isPrompt i) k
  c <- alloc $ Ready $ Seq f
  let e' = (n, c) : e
  res <- alloc $ OneShot e' p
  return (res, k')

isPrompt :: Int -> Kont -> Bool
isPrompt i (Prompt i') = i == i'
isPrompt _  _          = False

-- | Allocate nodes for let-declarations.
--
pushDecls :: Context -> [(Name, Program)] -> Program -> M Memory.Node
pushDecls e ds b = do
  ds' <- for ds \(n, p) -> do
    c <- alloc (error "boom!")
    return (n, p, c)

  let de = [(n, c) | (n, p, c) <- ds']
  let e' = de ++ e

  for_ ds' \(n, p, c) -> do
    writeMem c $ Updatable e' p

  alloc $ OneShot e' b

-- | Call closure: evaluate function body in captured context with formal arg replaced with concrete one.
--
--   Rewrites an argument cell into a `Rec`-node, unless it is `Ready`.
--
callClos n x e p = do
  c <- readMem x
  let
    c' = case c of
      Updatable e p -> Updatable e p
      OneShot e p -> Updatable e p
      Ready v   -> Ready v
  writeMem x c'
  prog ((n, x) : e, p)

-- | Call builtin: run builtin, allocate a Memory.Node for the result.
--
callBif bif x    = val =<< bif x

findName n e     = maybe (error ("undefined: " ++ n)) return $ lookup n e

-- | Prepare call: allocate cells for both function and argument.
--
preApp e f x = do
  f' <- prog (e, f);
  x' <- prog (e, x)
  return (f', x')

mkClos e n p = val $ Clo e n p
mkVal  f     = val $ VC f

readMem :: Memory.Node -> M Cell
readMem n = do
  pullsM $ Memory.get n

writeMem :: Memory.Node -> Cell -> M ()
writeMem n c = do
  changeM $ Memory.put n c

alloc :: Cell -> M Memory.Node
alloc c = do
  stateM $ Memory.new c

-- | Make a state action out of explicit CK-transformation.
--
--   Provides address of the current cell to the transformation.
--
next
  :: ((Memory.Node, Cell, [Kont]) -> M (Memory.Node, [Kont]))
  -> M Bool
next f = do
  addr <- pullM
  k    <- pullM
  cell <- readMem addr
  (cell', k') <- f (addr, cell, k)
  putM cell'
  putM k'
  pullM

-- | Halts the evaluation.
--
stop :: M ()
stop = do
  putM False

---- Testing area ----

-- | Evaluation loop. Runs `step` up to given gas limit (1 per `step`).
--
--   Returns unspent gas.
--
nStep :: Int -> M Int
nStep n = do
  if n > 0
  then do
    yes <- step
    if yes
    then do
      nStep (n - 1)
    else do
      return n
  else do
    return n

-- | Allocate a Memory.Node for program/context pair.
--
prog :: (Context, Program) -> M Memory.Node
prog (e, p) = alloc (OneShot e p)

-- | Allocate a Memory.Node for a value.
--
val :: Value -> M Memory.Node
val v = alloc (Ready v)

main = do
  let
    -- let three = f => x => f(f(f(x)))
    -- let inc = x => x + 1
    -- three(three(inc)) (1)
    --
    -- test  =
    --   Let
    --     [ ("three", Lam "f" (Lam "x" (Var "f" `App` (Var "f" `App` Var "x"))))
    --     , ("inc",   Lam "x" (PC (Bif (\(VC (Num f)) -> return $ VC (Num (f + 1)))) `App` Var "x"))
    --     , ("inc'",  Lam "x" (Var "inc" `App` Var "x"))
    --     ]
    --   $ App (App (Var "three") (Var "three")) (Var "inc'") `App` PC (Num 1)

    -- (reset 0 (1 + (shift 0 k -> (k (k 1)))))
    --
    test =
      Let
        [ ("+", PC          $ Bif \(VC (Num x)) ->
                return $ VC $ Bif \(VC (Num y)) ->
                return $ VC $ Num (x + y))
        , ("<", PC          $ Bif \(VC (Num x)) ->
                return $ VC $ Bif \(VC (Num y)) ->
                return $ Adt (if x < y then 1 else 0) 0 [])
        ]
      $ If 2 (Ctr 2 2 `App` PC (Num 1) `App` PC (Num 2)) (["a", "b"], Var "+" `App` Var "a" `App` Var "b") (PC (Num 0))

  let (mem, p) = Memory.new (OneShot [] test) Memory.empty
  (n, (a :> mem :> k :> s :> Nil)) <- nStep 40000 `runStateT` (p :> mem :> [] :> True :> Nil)
  print $ Memory.get a mem :> 40000 - n :> Nil

-- | Program.toString()
--
instance Show Program where
  show p = case p of
    Var n   -> n
    App f x -> "[" ++ show f ++ " " ++ show x ++ "]"
    Lam n p -> "[fun " ++ n ++ " -> " ++ show p ++ "]"
    Let d b -> "let " ++ intercalate "; " (map show' d) ++ " in " ++ show b
    PC  c   -> "[" ++ show c ++ "]"
    where
      show' (n, p) = "val " ++ n ++ " = " ++ show p

-- | Value.toString()
--
instance Show Value where
  show v = case v of
    Clo _ n p -> "(fun " ++ n ++ " -> " ++ show p ++ ")"
    VC c      -> show c


-- | Constant.toString()
--
instance Show Constant where
  show c = case c of
    Num d -> show d
    Bif b -> "<bif>"

-- | Function.toString()
--
instance Show (a -> b) where
  show _ = "<fun>"

Memory.hs

{-# language DataKinds #-}
{-# language GADTs #-}

{-|
  An implementation for memory.
-}

module Memory
  ( -- * Interface
    T
  , Node
  , empty
  , new
  , get
  , put
  ) where

import qualified Data.IntMap as Map

import Tuple

-- | An abstract datatype for memory.
--
data T x = Memory
  { pointer :: Int
  , grid    :: Map.IntMap x
  }

-- | Memory address.
--
type Node = Int

-- | Create an empty memory block.
--
empty :: T x
empty = Memory 0 Map.empty

-- | Allocate an item in the memory, return allocation address.
--
new :: x -> T x -> (T x, Node)
new x (Memory ptr m ) =
  (Memory (ptr + 1) (Map.insert ptr x m), ptr)

-- | Read memory at address.
--
get :: Node -> T x -> x
get ptr (Memory _ m) = m Map.! ptr

-- | Write memory at address.
put :: Node -> x -> T x -> T x
put c x (Memory ptr m) = Memory ptr $ Map.insert c x m

Tuple.hs

{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language GADTs #-}
{-# language MultiParamTypeClasses #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language ExplicitForAll #-}
{-# language BlockArguments #-}

{-|
  An extensible tuple.
-}
module Tuple
  ( -- * Basic interface
    Tuple ((:>), Nil)
  , Contains (pull, change)

    -- * MonadState interface
  , pullM
  , pullsM
  , changeM
  , putM
  , stateM
  , zoomM

    -- * Injection
  , Subset (select, assign)
  ) where

import Control.Monad.State

-- | An implementation.
--
data Tuple xs where
  (:>) :: x -> Tuple xs -> Tuple (x : xs)
  Nil  :: Tuple '[]

infixr 1 :>

-- | The `x` is the element of `Tuple xs`.
--
class Contains x xs where
  -- | Get `x` from `Tuple xs`.
  --
  pull :: Tuple xs -> x

  -- | Modify `x` inside `Tuple xs`.
  change :: (x -> x) -> Tuple xs -> Tuple xs

-- | `get`
--
pullM :: forall x xs m. (MonadState (Tuple xs) m, Contains x xs) => m x
pullM = gets pull

-- | `gets`
--
pullsM :: forall x xs m a. (MonadState (Tuple xs) m, Contains x xs) => (x -> a) -> m a
pullsM f = gets (f . pull)

-- | `modify`
--
changeM :: forall x xs m. (MonadState (Tuple xs) m, Contains x xs) => (x -> x) -> m ()
changeM = modify . change

-- | `put`
--
putM :: forall x xs m. (MonadState (Tuple xs) m, Contains x xs) => x -> m ()
putM = changeM . const

-- | `state`
--
stateM :: forall x xs m a. (MonadState (Tuple xs) m, Contains x xs) => (x -> (x, a)) -> m a
stateM f = do
  x <- pullM
  let (x', a) = f x
  putM x'
  return a

-- | `Control.Lens.zoom`
--
zoomM
  :: forall xs ys m a
  . ( Monad m
    , Subset xs ys
    )
  => StateT (Tuple xs) m a -> StateT (Tuple ys) m a
zoomM act = StateT \ys -> do
  let xs = select ys
  (a, xs') <- runStateT act xs
  return (a, assign xs' ys)

instance {-# overlaps #-} Contains x (x : xs) where
  pull (x :> _) = x
  change f (x :> xs) = f x :> xs

instance Contains x xs => Contains x (y : xs) where
  pull (_ :> xs) = pull xs
  change f (x :> xs) = x :> change f xs

instance ShowTuple xs => Show (Tuple xs) where
  show xs = "{" ++ showTuple xs ++ "}"

class ShowTuple xs where
  showTuple :: Tuple xs -> String

instance ShowTuple '[] where
  showTuple _ = ""

instance {-# overlaps #-} Show x => ShowTuple '[x] where
  showTuple (x :> Nil) = show x

instance (Show x, ShowTuple xs) => ShowTuple (x : xs) where
  showTuple (x :> xs) = show x ++ ", " ++ showTuple xs

-- | An injection of `Tuple xs` into `Tuple ys`.
--
class Subset xs ys where
  select :: Tuple ys -> Tuple xs
  assign :: Tuple xs -> (Tuple ys -> Tuple ys)

instance Subset '[] ys where
  select _ = Nil
  assign _ = id

instance (Subset xs ys, Contains x ys) => Subset (x : xs) ys where
  select ys = pull ys :> select ys
  assign (x :> xs) ys = change (const x) $ assign xs ys