Haskell Nes Cpu

module Nes where

import Data.Word
import Data.Int
import Data.Bits (Bits, (.|.), (.&.), shiftL, shiftR, testBit, xor, clearBit, setBit)
import Data.List (foldl')
import Control.Monad (ap)
import Data.Maybe (fromMaybe)
import Numeric (showHex)

import qualified Data.Vector as V
import qualified Data.Map as M
import Data.Vector ((!), (!?), (//))

data Cpu = Cpu
  { rom :: V.Vector Word8
  , memory :: V.Vector Word8
  , a :: Word8
  , x :: Word8
  , y :: Word8
  , pc :: Word16
  , sp :: Word8
  , nReg :: Bool
  , vReg :: Bool
  -- EMPTY
  , bReg :: Bool
  , dReg :: Bool
  , iReg :: Bool
  , zReg :: Bool
  , cReg :: Bool
  } deriving (Eq)

bool2int False = 0
bool2int True = 1

fromListBE :: (Num b, Bits b) => [Bool] -> b
fromListBE = foldl' f 0
  where
    f i b = (i `shiftL` 1) .|. bool2int b

statusRegisters :: Cpu -> Word8
statusRegisters = fromListBE . (ap registers) . pure
  where registers = [nReg, vReg, const True, bReg, dReg, iReg, zReg, cReg]


instance Show Cpu where
  show cpu =
    "A=$" ++ (showHex . a) cpu " X=$" ++ (showHex . x) cpu " Y=$" ++ (showHex . y) cpu "\n" ++
    "SP=$" ++ (showHex . sp) cpu " PC=$" ++ (showHex . pc) cpu "\n" ++
    "NV-BDIZC" ++ "\n" ++
    (show . bool2int . nReg) cpu ++
    (show . bool2int . vReg) cpu ++ "--" ++
    (show . bool2int . dReg) cpu ++
    (show . bool2int . iReg) cpu ++
    (show . bool2int . zReg) cpu ++
    (show . bool2int . cReg) cpu



data AddressMode
  = Immediate -- Immediate
  | ZeroPage  -- ZeroPage
  | ZeroPageX -- ZeroPage Indexed
  | ZeroPageY -- ZeroPage Indexed
  | Absolute -- Absolute
  | AbsoluteX -- Indexed
  | AbsoluteY -- Indexed
  | Implied -- Implied
  | Relative -- Relative
  | Indirect --  AbsoluteIndirect, Only used for JMP
  | IndirectX -- PreIndexed Indirect
  | IndirectY -- PostIndexed Indirect
  deriving (Show, Eq)

mkCpuFromMemory :: [Word8] -> Cpu
mkCpuFromMemory m = Cpu
  { rom = V.fromList m
  , memory = V.fromList m
  , a = 0x00, x = 0x00, y = 0x00
  , pc = 0x0000, sp = 0x00
  , nReg = False, vReg = False, dReg = False, iReg = False, zReg = False, cReg = False, bReg = True
  }

mkCpu :: Int -> Cpu
mkCpu b = Cpu
  { rom = (V.replicate b 0x00)
  , memory = (V.replicate b 0x00)
  , a = 0x00, x = 0x00, y = 0x00
  , pc = 0x0000, sp = 0x00
  , nReg = False, vReg = False, dReg = False, iReg = False, zReg = False, cReg = False, bReg = True
  }

writeWord8 :: V.Vector Word8 -> Word16 -> Word8 -> V.Vector Word8
writeWord8 mem address byte = mem // [(fromIntegral address, byte)]

readWord8 :: Cpu -> Word16 -> Word8
readWord8 cpu address = (memory cpu) ! (fromIntegral address)

readWord16 :: Cpu -> Word16 -> Word16
readWord16 cpu address = fromWord8s $ V.slice (fromIntegral address) 2 (memory cpu)

fromWord8s :: V.Vector Word8 -> Word16
fromWord8s ws = foldr accum 0x0000 ws
  where accum o a = (a `shiftL` 8) .|. fromIntegral o

getArgument :: Cpu -> AddressMode -> Maybe Word16
-- TODO: change the type of this to be *ARGUMENT*, either an address or nothing
getArgument cpu Implied = Nothing
getArgument cpu Immediate = Just $ pc cpu
getArgument cpu Absolute = Just $ readWord16 cpu (pc cpu)
getArgument cpu ZeroPage = Just $ fromIntegral $ readWord8 cpu (pc cpu)
getArgument cpu ZeroPageX = Just $ fromIntegral operand
  where
    operand = readWord8 cpu $ (pc cpu) + (fromIntegral . x) cpu
getArgument cpu ZeroPageY = Just $ fromIntegral operand
  where
    operand = readWord8 cpu $ (pc cpu) + (fromIntegral . y) cpu
getArgument cpu AbsoluteX = Just $ argVal + (fromIntegral . x) cpu
  where
    argVal = readWord16 cpu (pc cpu)
getArgument cpu AbsoluteY = Just $ argVal + (fromIntegral . y) cpu
  where
    argVal = readWord16 cpu (pc cpu)
getArgument cpu Relative = Just $ fromIntegral operand
  where
    argVal = readWord8 cpu (pc cpu)
    argVal' = fromIntegral argVal :: Int8
    operand = 1 + (fromIntegral (pc cpu)) + (fromIntegral argVal')
getArgument cpu Indirect = Just $ readWord16 cpu argVal
  where
    argVal = readWord16 cpu (pc cpu)
getArgument cpu IndirectX = Just $ readWord16 cpu xIndex
  where
    argVal = readWord8 cpu (pc cpu )
    xIndex = fromIntegral argVal + (fromIntegral . x) cpu
getArgument cpu IndirectY = Just $ preAddress + (fromIntegral . y) cpu
  where
    argVal = fromIntegral $ readWord8 cpu (pc cpu)
    preAddress = readWord16 cpu argVal


--instructions ::
instructions = M.fromList
  [
  -- ADC
    (0x69, (Immediate, adc))
  , (0x65, (ZeroPage, adc))
  , (0x75, (ZeroPageX, adc))
  , (0x6d, (Absolute, adc))
  , (0x7d, (AbsoluteX, adc))
  , (0x79, (AbsoluteY, adc))
  , (0x61, (IndirectX, adc))
  , (0x71, (IndirectY, adc))
  -- AND
  , (0x29, (Immediate, andOp))
  , (0x25, (ZeroPage, andOp))
  , (0x35, (ZeroPageX, andOp))
  , (0x2d, (Absolute, andOp))
  , (0x3d, (AbsoluteX, andOp))
  , (0x39, (AbsoluteY, andOp))
  , (0x21, (IndirectX, andOp))
  , (0x31, (IndirectY, andOp))
  -- ASL
  , (0x0a, (Implied, asl))
  , (0x06, (ZeroPage, asl))
  , (0x16, (ZeroPageX, asl))
  , (0x0e, (Absolute, asl))
  , (0x1e, (AbsoluteX, asl))
  -- BCC
  , (0x90, (Relative, bcc))
  -- BCS
  , (0xb0, (Relative, bcs))
  -- BEQ
  , (0xf0, (Relative, beq))
  -- BIT
  , (0x24, (ZeroPage, bit))
  , (0x2c, (Absolute, bit))
  ]

incPc :: Cpu -> Cpu
incPc c = c { pc = 1 + pc c }

incPcByMode :: Cpu -> AddressMode -> Cpu
incPcByMode c Implied = c
incPcByMode c Immediate = c { pc = 1 + pc c }
incPcByMode c ZeroPage = c { pc = 1 + pc c }
incPcByMode c ZeroPageX = c { pc = 1 + pc c }
incPcByMode c ZeroPageY = c { pc = 1 + pc c }
incPcByMode c Absolute = c { pc = 2 + pc c }
incPcByMode c AbsoluteX = c { pc = 2 + pc c }
incPcByMode c AbsoluteY = c { pc = 2 + pc c }
incPcByMode c Indirect = c { pc = 2 + pc c }
incPcByMode c IndirectX = c { pc = 1 + pc c }
incPcByMode c IndirectY = c { pc = 1 + pc c }
incPcByMode c Relative = c { pc = 1 + pc c }

data CpuError
  = OpcodeNotImplemented Word8
  | InvalidArgument String
  deriving (Eq)

instance Show CpuError where
   show (OpcodeNotImplemented opcode) = "OpcodeNotImplemented 0x" ++ (showHex opcode) ""
   show (InvalidArgument inst) = "InvalidArgument for " ++ inst

lookupOpcode :: Word8 -> Either CpuError (AddressMode, Cpu -> Maybe Word16 -> Either CpuError Cpu)
lookupOpcode o =
  case M.lookup o instructions of
    Nothing -> Left $ OpcodeNotImplemented o
    Just mf -> Right mf

updateCpuM :: Cpu -> Either CpuError Cpu
updateCpuM c = do
  opcode <- return $ readWord8 c (pc c)
  (mode, function) <- lookupOpcode opcode
  c' <- return $ incPc c
  argument <- return $ getArgument c' mode
  c'' <- return $ incPcByMode c' mode
  function c'' argument


setZeroFlag :: Word8 -> Bool
setZeroFlag 0x00 = True
setZeroFlag _ = False

setCarryFlagAdd :: Word8 -> Word8 -> Bool
setCarryFlagAdd x y =
  testBit (x' + y') 8
  where
    x' = fromIntegral x
    y' = fromIntegral y
    s = (x' + y') :: Word16

setNegativeFlag :: Word8 -> Bool
setNegativeFlag b = testBit b 7

setOverflowFlag :: Word8 -> Word8 -> Word8 -> Bool
setOverflowFlag a b s = 0 /= ((a `xor` s) .&. (b `xor` s)) .&. 0x80

adc :: Cpu -> Maybe Word16 -> Either CpuError Cpu
adc c (Just address) = Right $
  c { a = a'
    , zReg = setZeroFlag a'
    , cReg = setCarryFlagAdd a1 a2
    , nReg = setNegativeFlag a'
    , vReg = setOverflowFlag a1 a2 a'
    }
  where
    a' = a1 + a2 + (bool2int . cReg) c
    a1 = readWord8 c address
    a2 = a c
adc _ Nothing = Left $ InvalidArgument "ADC"

andOp :: Cpu -> Maybe Word16 -> Either CpuError Cpu
andOp c (Just address) = Right $
  c { a = res
    , zReg = setZeroFlag res
    , nReg = setNegativeFlag res
    }
  where
    arg = readWord8 c address
    res = (a c) .&. arg
andOp _ Nothing = Left $ InvalidArgument "AND"

asl :: Cpu -> Maybe Word16 -> Either CpuError Cpu
asl c Nothing = Right $
  c { a = a'
    , cReg = testBit (a c) 7
    , zReg = setZeroFlag a'
    , nReg = setNegativeFlag a'
    }
  where
    a' = (a c) `shiftL` 1
asl c (Just address) = Right $
  c { memory = m
    , cReg = testBit arg 7
    , zReg = setZeroFlag (a c)
    , nReg = setNegativeFlag shifted
    }
  where
    arg = readWord8 c address
    shifted = arg `shiftL` 1
    m = writeWord8 (memory c) address shifted

bcc :: Cpu -> Maybe Word16 -> Either CpuError Cpu
bcc c (Just address) = Right $
  if (cReg c)
  then c
  else c { pc = address }
bcc _ Nothing = Left $ InvalidArgument "BCC"

bcs :: Cpu -> Maybe Word16 -> Either CpuError Cpu
bcs c (Just address) = Right $
  if (cReg c)
  then c { pc = address }
  else c
bcs _ Nothing = Left $ InvalidArgument "BCS"

beq :: Cpu -> Maybe Word16 -> Either CpuError Cpu
beq c (Just address) = Right $
  if (zReg c)
  then c { pc = address }
  else c
beq _ Nothing = Left $ InvalidArgument "BEQ"

bit :: Cpu -> Maybe Word16 -> Either CpuError Cpu
bit c (Just address) = Right $
  c { zReg = setZeroFlag result
    , vReg = testBit arg 6
    , nReg = setNegativeFlag arg
    }
  where
    arg = readWord8 c address
    result = (a c) .&. arg
bit _ Nothing = Left $ InvalidArgument "BIT"