vm.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}

import Data.Functor ( void )
import Data.Vector ( Vector, (!?), fromList )
import qualified Data.Map.Strict as Map ( Map, lookup, insert, empty )
import Control.Monad ( forever, unless )
import Control.Monad.State.Strict ( MonadState, modify, get, put, evalStateT )
import Control.Monad.IO.Class
import Control.Lens
import System.Exit

newtype Val = Val Int
  deriving stock Show
  deriving newtype (Eq, Num)

newtype PC = PC Int
  deriving stock Show
  deriving newtype Enum

newtype Var = Var String
  deriving stock Show
  deriving newtype (Ord, Eq)

newtype Program = Program (Vector Inst)
  deriving stock Show

data Inst
  = Push Val
  | Pop
  | Jump PC
  | JumpIf PC
  | Add
  | Sub
  | Mul
  | Set Var
  | Get Var
  | Print
  | Halt
  deriving Show

data Env
  = Env
  { _stack  :: [Val]
  , _vars   :: Map.Map Var Val
  , _pc     :: PC
  , _nextPC :: PC
  }
  deriving stock Show

makeLenses ''Env

initEnv :: Env
initEnv = Env [] Map.empty (PC 0) (PC 1)

instAt :: MonadFail m => PC -> Program -> m Inst
instAt (PC idx) (Program insts) =
  case insts !? idx of
    Just inst -> pure inst
    Nothing -> fail ("invalid PC " ++ show idx)


push :: MonadState Env m => Val -> m ()
push v = stack %= (v:)

pop :: (MonadFail m, MonadState Env m) => m Val
pop = do
  (h:t) <- use stack
  stack .= t
  pure h

peek :: (MonadFail m, MonadState Env m) => m Val
peek = f =<< use stack
  where
    f (h : _) = pure h
    f [] = fail "empty stack"

setVar :: MonadState Env m => Var -> Val -> m ()
setVar var val = vars %= Map.insert var val

getVar :: (MonadFail m, MonadState Env m) => Var -> m Val
getVar var = do
  m <- use vars
  case Map.lookup var m of
    Just x -> pure x
    Nothing -> fail ("no such variable " ++ show var)

jump :: MonadState Env m => PC -> m ()
jump c = nextPC .= c

interpret :: (MonadState Env m, MonadIO m, MonadFail m) => Inst -> m ()
interpret (Push v) = push v
interpret Pop = void pop
interpret (Jump c) = jump c
interpret (JumpIf c) = do
  s0 <- peek
  unless (s0 == Val 0) $ jump c
interpret Add = do
  s0 <- pop
  s1 <- pop
  push (s1 + s0)
interpret Sub = do
  s0 <- pop
  s1 <- pop
  push (s1 - s0)
interpret Mul = do
  s0 <- pop
  s1 <- pop
  push (s1 * s0)
interpret (Set var) = do
  s0 <- pop
  setVar var s0
interpret (Get var) = do
  v <- getVar var
  push v
interpret Print = do
  s0 <- peek
  liftIO $ print s0
interpret Halt = liftIO exitSuccess

select :: (MonadFail m, MonadState Env m) => Program -> m Inst
select program = do
  p <- use pc
  instAt p program

run :: Env -> Program -> IO ()
run e p = evalStateT machine e
  where
    machine = forever $ do
      inst <- select p
      nextPC <~ uses pc succ
      interpret inst
      pc <~ use nextPC

program :: Program
program = Program $ fromList
  [ Push 10
  , Push 1
  , Sub
  , Print
  , JumpIf (PC 1)
  , Halt
  ]

main :: IO ()
main = run initEnv program