Haskell Quiz/Chip Eight/Solution Jethr0

As vincenz pointed out in one of my other solutions, I'm probably using State too heavily now that I've (hopefully) figured out how to use it ;)

At the moment the eval function could be written in pure fashion and maybe I'll do that refactoring some time. Interpreter isn't fully tested, but the sample program seems to be "running" correctly

module Main where

import qualified Data.Array as Array
import qualified Data.Bits as Bits
import qualified Data.Char as Char
import Data.Word            (Word8, Word16)
import Data.Bits            ((.&.), (.|.), shiftL, shiftR)
import Data.Array           ((!), (//), Array, listArray, accumArray)
import Control.Monad        (when)
import Control.Monad.State  (get, put, modify, liftIO, execStateT, StateT)
import Control.Monad.Identity
import Numeric              (showHex, showIntAtBase)
import Text.Printf          (printf)
import System.Random        (randomR, mkStdGen, StdGen, newStdGen)

numRegisters = 16
sizeMemory   = 2^12

data MachineState = MachineState {
  rv      :: Array Int Word8,
  rip     :: Word16,
  rmemory :: Array Int Word8,
  rand    :: StdGen
}

instance Show MachineState where
  show MachineState {rv = regs, rmemory = mem, rip = ip} =
    unlines showRegs ++ "IP: " ++ (printf "%08x" (fromIntegral ip :: Integer))
    where showRegs  = zipWith (\r v -> r ++ ": " ++ v) regNames regValues
          regNames  = [printf "V%x" (fromIntegral x :: Integer) | x <- [0..]]
          --regValues = [printf "%04x" (fromIntegral x :: Integer) | x <- Array.elems regs]
          regValues = [showIntAtBase 2 Char.intToDigit x "" | x <- Array.elems regs]

type StT = StateT MachineState
type St  = StT Identity
type Offset = Int


---

modify_rv      func = modify (\st -> st{rv      = func . rv      $ st})
modify_rip     func = modify (\st -> st{rip     = func . rip     $ st})
modify_rmemory func = modify (\st -> st{rmemory = func . rmemory $ st})
getReg x     = do {MachineState{rv = regs} <- get; return $ regs!x}
setReg x val = modify_rv (// [(fromIntegral x, val)])


step :: St ()
step = do
  MachineState{rip = ip, rmemory = mem} <- get
  let (i1,i2) = (fromIntegral $ mem ! (fromIntegral ip)
                ,fromIntegral $ mem ! (fromIntegral $ ip+1))
      instr   = (shiftL i1 8) + i2 :: Word16
  case instr of
    0x0000    -> return ()
    otherwise -> eval instr >> modify_rip (2+) >> step


eval :: Word16 -> St ()
eval instr = case firstDigit instr of
  0x1 -> do
    let nnn = instr .&. 0x0FFF
    modify_rip (const nnn)

  0x3 -> do
    let (x, kk) = (bitX, bitKK)
    vx <- getReg x
    when (vx == kk) (modify_rip (2+))

  0x6 -> do
    let (x, kk) = (bitX, bitKK)
    setReg x kk
    
  0x7 -> do
    let (x, kk) = (bitX, bitKK)
    vx <- getReg x
    rPlus x vx kk

  0x8 -> do
    let x = bitX
        y = fromIntegral $ shiftR (instr .&. 0x00F0) 4
    vx <- getReg x
    vy <- getReg y
    case (instr .&. 0x000F) of 
      0x0 -> setReg  x $ vy
      0x1 -> setReg  x $ vx .|. vy
      0x2 -> setReg  x $ vx .&. vy
      0x3 -> setReg  x $ vx `Bits.xor` vy
      0x4 -> rPlus   x vx vy
      0x5 -> rMinus  x vx vy
      0x6 -> rShiftR x vx 1
      0x7 -> rMinus  x vy vx
      0xE -> rShiftL x vx 1
    where rShiftR target vx n = do
            setReg 0xF    $ vx .&. 0x01
            setReg target $ shiftR vx n
          rShiftL target vx n = do
            setReg 0xF    $ shiftR (vx .&. 0x80) 7 -- FIXME: is this correct
            setReg target $ shiftL vx n

  0xC -> do
    let (x, kk) = (bitX, bitKK)
    r <- rRandom
    setReg x $ r .&. kk

  otherwise -> 
    error $ "opcode not implemented " ++ showHex instr ""
    
  where bitX  = fromIntegral $ shiftR (instr .&. 0x0F00) 8
        bitKK = fromIntegral $ instr .&. 0x00FF

        firstDigit w = fromIntegral $ shiftR w 12

        rRandom = do
          state@MachineState{rand = gen} <- get
          let (r, gen') = randomR (0, 2^8-1) gen
          put state{rand = gen'}
          return $ fromIntegral r

        rPlus :: Integral i => i -> Word8 -> Word8 -> St ()
        rPlus target a b = do
          let sum = (fromIntegral a) + (fromIntegral b) :: Integer
          setReg 0xF $ if sum >= 2^8 then 1 else 0
          setReg target $ (fromIntegral sum) .&. 0x00FF

        rMinus :: Integral i => i -> Word8 -> Word8 -> St ()
        rMinus target a b = do
          let sum = (fromIntegral a) - (fromIntegral b) :: Integer
          setReg 0xF $ if (sum < 0) then 0 else 1
          setReg target . fromIntegral $ if (sum < 0) then (sum + 2^8) else sum
  

---


initialState :: StdGen -> MachineState
initialState gen = MachineState {
  rv      = accumArray const 0 (0, numRegisters-1) [],
  rmemory = accumArray const 0 (0, sizeMemory-1)   [],
  rip     = 0,
  rand    = gen
}


modifyMemory :: Offset -> [Word8] -> MachineState -> MachineState
modifyMemory offset words state = 
  state{rmemory = rmemory state // zip [offset..] words}

modifyRegisters :: [(Int, Word8)] -> MachineState -> MachineState
modifyRegisters pairs state = 
  state{rv = rv state // pairs}

main = do
  g       <- newStdGen
  file    <- readFile "Chip8Test"
  let program = map (fromIntegral . Char.ord) file
      start   = initialState g
      new     = modifyMemory 0 program
              . modifyRegisters []
              $ start
  let res = runIdentity $ execStateT step new
  print res