Haskell Quiz/Chip Eight/Solution Jethr0

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