Bytecode interpreter

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedLists #-}

module Main
    ( main
    ) where

import Control.Monad.ST (ST, runST)
import Data.Bits ((.|.), shiftL)
import Data.ByteString (ByteString)
import Data.Function ((&))
import Data.Semigroup ((<>))
import Data.Vector.Mutable (MVector)
import Data.Word (Word16)

import qualified Data.ByteString as ByteString
import qualified Data.Vector as Vector
import qualified Data.Vector.Mutable as MVector

data Value
    = BoolValue !Bool
    | IntValue {-# UNPACK #-} !Int
    deriving stock (Show)

interpret :: ByteString -> Int -> MVector s Value -> ST s Value
interpret code pc mem =
    case ByteString.index code pc of
        -- CONTINUE
        0x00 ->
            interpret code (pc + 1) mem

        -- GOTO
        0x01 -> do
            let target = decodeWord16 code (pc + 1) & fromIntegral
            interpret code target mem

        -- IF
        0x02 -> do
            let target = decodeWord16 code (pc + 1) & fromIntegral
            let reg = decodeWord16 code (pc + 3) & fromIntegral
            BoolValue condition <- MVector.read mem reg
            if condition
                then interpret code target mem
                else interpret code (pc + 5) mem

        -- RETURN
        0x04 -> do
            let reg = decodeWord16 code (pc + 1) & fromIntegral
            MVector.read mem reg

        -- COPY
        0x10 -> do
            let sourceReg = decodeWord16 code (pc + 1) & fromIntegral
            let targetReg = decodeWord16 code (pc + 3) & fromIntegral
            value <- MVector.read mem sourceReg
            MVector.write mem targetReg value
            interpret code (pc + 5) mem

        other -> fail $ "Invalid opcode: " <> show other

decodeWord16 :: ByteString -> Int -> Word16
decodeWord16 bs i =
    let b1 = fromIntegral $ ByteString.index bs (i + 0) in
    let b2 = fromIntegral $ ByteString.index bs (i + 1) in
    b2 `shiftL` 8 .|. b1
{-# INLINE decodeWord16 #-}

main :: IO ()
main =
    let
        code = ByteString.pack 
            [ {-  0 -} 0x00                 -- CONTINUE
            , {-  1 -} 0x01, 0x05, 0x00     -- GOTO 5
            , {-  4 -} 0xFF                 -- (invalid opcode)
            , {-  5 -} 0x02, 0x0F, 0x00     -- IF 15 0
                           , 0x00, 0x00
            , {- 10 -} 0x10, 0x02, 0x00     -- COPY 2 3
                           , 0x03, 0x00
            , {- 15 -} 0x02, 0x17, 0x00     -- IF 23 1
                           , 0x01, 0x00
            , {- 20 -} 0x04, 0x02, 0x00     -- RETURN 2
            , {- 23 -} 0x04, 0x03, 0x00     -- RETURN 3
            ]
        pc = 0
        mem1 = Vector.thaw [BoolValue True,  BoolValue True,  IntValue 1, IntValue 2]
        mem2 = Vector.thaw [BoolValue True,  BoolValue False, IntValue 1, IntValue 2]
        mem3 = Vector.thaw [BoolValue False, BoolValue True,  IntValue 1, IntValue 2]
        mem4 = Vector.thaw [BoolValue False, BoolValue False, IntValue 1, IntValue 2]
    in
        do { print $ runST (interpret code pc =<< mem1)
           ; print $ runST (interpret code pc =<< mem2)
           ; print $ runST (interpret code pc =<< mem3)
           ; print $ runST (interpret code pc =<< mem4) }