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 st = st{rv      = func . rv      $ st}
modify_rip     func st = st{rip     = func . rip     $ st}
modify_rmemory func st = st{rmemory = func . rmemory $ st}
getReg x     st = let MachineState{rv = regs} = st in regs!x
setReg x val st = modify_rv (// [(fromIntegral x, val)]) st
 
 
eval :: Word16 -> MachineState -> MachineState
eval instr st = case firstDigit instr of
  -- 0x1NNN
  0x1 -> 
    modify_rip (const nnn) st
    where nnn = instr .&. 0x0FFF
 
  -- 0x3XKK
  0x3 ->
    if vx == kk 
      then modify_rip (2+) st 
      else st
 
  -- 0x6XKK
  0x6 -> 
    setReg x kk st
 
  -- 0x7XKK
  0x7 -> 
    rPlus x vx kk st
 
  -- 0x8XY_
  0x8 -> 
    (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
    ) st
    where y  = fromIntegral $ shiftR (instr .&. 0x00F0) 4
          vy = getReg y st
 
  -- 0xCXKK
  0xC -> 
    setReg x (r .&. kk) st'
    where (r, st') = rRandom st
 
  otherwise -> 
    error $ "opcode not implemented " ++ showHex instr ""
 
  where x  = fromIntegral $ shiftR (instr .&. 0x0F00) 8
        vx = getReg x st
        kk = fromIntegral $ instr .&. 0x00FF
        (>.>) = flip ($)
 
        firstDigit w = fromIntegral $ shiftR w 12
 
        rRandom s = (fromIntegral r, s{rand = gen})
          where MachineState{rand = gen} = s
                (r, gen') = randomR (0, 2^8-1) gen
 
        rPlus target a b s =
          s
          >.> setReg 0xF    (if sum >= 2^8 then 1 else 0)
          >.> setReg target ((fromIntegral sum) .&. 0x00FF)
          where sum = (fromIntegral a) + (fromIntegral b) :: Integer
 
        rMinus target a b s =
          s
          >.> setReg 0xF (if (sum < 0) then 0 else 1)
          >.> setReg target (fromIntegral (if (sum < 0) then (sum + 2^8) 
                                                        else sum))
          where sum = (fromIntegral a) - (fromIntegral b) :: Integer
 
        rShiftR target vx n s = 
          s 
          >.> setReg 0xF    (vx .&. 0x01)
          >.> setReg target (shiftR vx n)
 
        rShiftL target vx n s =
          s
          >.> setReg 0xF    (shiftR (vx .&. 0x80) 7) -- FIXME: is this correct
          >.> setReg target (shiftL vx n)
 
---
 
 
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}
 
 
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 -> do
      modify $ eval instr
      modify $ modify_rip (+2)
      step
 
 
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