Nes.hs

{-# LANGUAGE DuplicateRecordFields      #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Emulator.Nes (
    Nes(..)
  , CPU(..)
  , PPU(..)
  , Emulator(..)
  , Sprite(..)
  , Coords
  , Color
  , Flag(..)
  , IncrementMode(..)
  , SpriteSize(..)
  , Interrupt(..)
  , new
  , runEmulator
  , debug
  , loadCpu
  , storeCpu
  , modifyCpu
  , loadPpu
  , storePpu
  , modifyPpu
  , readCpuMemory8
  , readCpuMemory16
  , writeCpuMemory8
  , writeCpuMemory16
  , readPpuMemory
  , writeControl
  , writeMask
  , writeOAMAddress
  , readOAMData
  , readPalette
  , storeKeys
  , loadKeys
  , writeScreen
  , loadScreen
  , toggleNmi
) where

import           Control.Monad
import           Control.Monad.IO.Class       (liftIO)
import           Control.Monad.Reader         (MonadReader, ReaderT, ask, runReaderT)
import           Control.Monad.Trans          (MonadIO)
import           Data.Bits                    (shiftL, shiftR, testBit, (.&.), (.|.))
import qualified Data.ByteString              as BS
import           Data.IORef
import           Data.Set                     as Set
import qualified Data.Vector                  as V
import qualified Data.Vector.Storable.Mutable as VUM
import           Data.Word
import           Emulator.Cartridge           as Cartridge
import qualified Emulator.Controller          as Controller
import qualified Emulator.Mapper              as Mapper
import           Emulator.Util
import           Prelude                      hiding (read, replicate)

data Sprite = Sprite {
  sIndex         :: Int,
  sCoords        :: Coords,
  sTileIndexByte :: Word8,
  sAttributeByte :: Word8,
  sPattern       :: Word32,
  sPriority      :: Word8
} deriving (Show, Eq)

type Coords = (Int, Int)

type Color = (Word8, Word8, Word8)

data IncrementMode = Horizontal | Vertical

data SpriteSize = Normal | Double

data ColorMode = Color | Grayscale

data Visibility = Hidden | Shown

data Interrupt
  = IRQ
  | NMI
  deriving (Eq, Show)

data Nes = Nes {
  cpu        :: CPU,
  ppu        :: PPU,
  cart       :: Cartridge,
  mapper     :: Mapper.Mapper,
  controller :: Controller.Controller
}

data CPU = CPU {
  pc        :: IORef Word16,
  sp        :: IORef Word8,
  a         :: IORef Word8,
  x         :: IORef Word8,
  y         :: IORef Word8,
  p         :: IORef Word8,
  cpuCycles :: IORef Int,
  interrupt :: IORef (Maybe Interrupt),
  ram       :: VUM.IOVector Word8
}

data PPU = PPU {
  -- Misc
  ppuCycles             :: IORef Int,
  scanline              :: IORef Int,
  frameCount            :: IORef Int,
  writeToggle           :: IORef Bool,
  ppuRegister           :: IORef Word8,
  oddFrame              :: IORef Bool,
  -- Data
  oamData               :: VUM.IOVector Word8,
  nameTableData         :: VUM.IOVector Word8,
  paletteData           :: VUM.IOVector Word8,
  screen                :: VUM.IOVector Word8,
  -- Addresses
  currentVramAddress    :: IORef Word16,
  tempVramAddress       :: IORef Word16,
  oamAddress            :: IORef Word8,
  -- NMI
  nmiOutput             :: IORef Bool,
  nmiOccurred           :: IORef Bool,
  nmiDelay              :: IORef Word8,
  nmiPrevious           :: IORef Bool,
  -- Control register bits
  nameTable             :: IORef Word16,
  incrementMode         :: IORef IncrementMode,
  spriteTable           :: IORef Word16,
  bgTable               :: IORef Word16,
  spriteSize            :: IORef SpriteSize,
  -- Mask register bits
  colorMode             :: IORef ColorMode,
  leftBgVisibility      :: IORef Visibility,
  leftSpritesVisibility :: IORef Visibility,
  bgVisibility          :: IORef Bool,
  spriteVisibility      :: IORef Bool,
  intensifyReds         :: IORef Bool,
  intensifyGreens       :: IORef Bool,
  intensifyBlues        :: IORef Bool,
  -- Status register bits
  spriteOverflow        :: IORef Bool,
  spriteZeroHit         :: IORef Bool,
  -- verticalBlank         :: IORef Bool,
  -- Scroll register
  fineX                 :: IORef Word8,
  -- Data register
  dataV                 :: IORef Word8,
  -- Temp vars
  nameTableByte         :: IORef Word8,
  attrTableByte         :: IORef Word8,
  loTileByte            :: IORef Word8,
  hiTileByte            :: IORef Word8,
  tileData              :: IORef Word64,
  sprites               :: IORef (V.Vector Sprite)
}

newtype Emulator a = Emulator { unNes :: ReaderT Nes IO a }
  deriving (Monad, Applicative, Functor, MonadIO, MonadReader Nes)

runEmulator :: BS.ByteString -> Emulator a ->  IO a
runEmulator bs (Emulator reader) = do
  cart <- Cartridge.parse bs
  nes <- new cart
  runReaderT reader nes

debug :: String -> Emulator ()
debug = liftIO . putStrLn

{-# INLINE with #-}
with :: (Nes -> b) -> (b -> IO a) -> Emulator a
with field f = do
  nes <- ask
  liftIO $ f (field nes)

{-# INLINE loadCpu #-}
loadCpu :: (CPU -> IORef b) -> Emulator b
loadCpu field = with (field . cpu) readIORef

{-# INLINE storeCpu #-}
storeCpu :: (CPU -> IORef b) -> b -> Emulator ()
storeCpu field v = with (field . cpu) (`modifyIORef'` const v)

{-# INLINE modifyCpu #-}
modifyCpu :: (CPU -> IORef b) -> (b -> b) -> Emulator ()
modifyCpu field v = with (field . cpu) (`modifyIORef'` v)

{-# INLINE loadPpu #-}
loadPpu :: (PPU -> IORef b) -> Emulator b
loadPpu field = with (field . ppu) readIORef

{-# INLINE storePpu #-}
storePpu :: (PPU -> IORef b) -> b -> Emulator ()
storePpu field v = with (field . ppu) (`modifyIORef'` const v)

{-# INLINE modifyPpu #-}
modifyPpu :: (PPU -> IORef b) -> (b -> b) -> Emulator ()
modifyPpu field v = with (field . ppu) (`modifyIORef'` v)

readCpuMemory8 :: Word16 -> Emulator Word8
readCpuMemory8 addr
  | addr < 0x2000 = readCPURam addr
  | addr < 0x4000 = readPPURegister $ 0x2000 + addr `rem` 8
  | addr == 0x4014 = readPPURegister addr
  | addr == 0x4015 = pure 0
  | addr == 0x4016 = readController
  | addr == 0x4017 = pure 0
  | addr < 0x6000 = pure 0
  | addr >= 0x6000 = readMapper addr
  | otherwise = error $ "Erroneous read detected at " ++ show addr ++ "!"

readCpuMemory16 :: Word16 -> Emulator Word16
readCpuMemory16 addr = do
  lo <- readCpuMemory8 addr
  hi <- readCpuMemory8 (addr + 1)
  pure $ makeW16 lo hi

writeCpuMemory8 :: Word16 -> Word8 -> Emulator ()
writeCpuMemory8 addr value
  | addr < 0x2000 = writeCPURam addr value
  | addr < 0x4000 = writePPURegister (0x2000 + addr `rem` 8) value
  | addr == 0x4014 = writePPURegister addr value
  | addr == 0x4016 = writeController value
  | addr >= 0x4000 && addr <= 0x4017 = pure ()
  | addr >= 0x4018 && addr <= 0x401F = pure ()
  | addr >= 0x6000 = writeMapper addr value
  | otherwise = error $ "Erroneous write detected at " ++ show addr ++ "!"

writeCpuMemory16 :: Word16 -> Word16 -> Emulator ()
writeCpuMemory16 addr value = do
  let (lo, hi) = splitW16 value
  writeCpuMemory8 addr lo
  writeCpuMemory8 (addr + 1) hi

readPpuMemory :: Word16 -> Emulator Word8
readPpuMemory addr
  | addr' < 0x2000 = readMapper addr'
  | addr' < 0x3F00 = readNametableData addr'
  | addr' < 0x4000 = readPalette addr'
  | otherwise = error $ "Erroneous read detected at " ++ show addr ++ "!"
  where addr' = addr `rem` 0x4000

writePPUMemory :: Word16 -> Word8 -> Emulator ()
writePPUMemory addr v
  | addr' < 0x2000 = writeMapper addr' v
  | addr' < 0x3F00 = writeNametableData addr' v
  | addr' < 0x4000 = writePalette addr' v
  | otherwise = error $ "Erroneous write detected at " ++ show addr ++ "!"
  where addr' = addr `rem` 0x4000

data Flag
  = Negative
  | Overflow
  | Unused
  | Break
  | Decimal
  | InterruptDisable
  | Zero
  | Carry
  deriving (Enum)

new :: Cartridge -> IO Nes
new cart = do
  cpu <- newCPU
  ppu <- newPPU
  mapper <- Mapper.new cart
  controller <- Controller.new
  pure $ Nes cpu ppu cart mapper controller

newCPU :: IO CPU
newCPU = do
  pc <- newIORef 0x0
  sp <- newIORef 0xFD
  a <- newIORef 0x0
  x <- newIORef 0x0
  y <- newIORef 0x0
  p <- newIORef 0x24 -- should this be 0x34?
  cycles <- newIORef 0
  interrupt <- newIORef Nothing
  ram <- VUM.replicate 65536 0x0
  pure $ CPU pc sp a x y p cycles interrupt ram

readCPURam :: Word16 -> Emulator Word8
readCPURam addr = with cpu $ \cpu ->
  VUM.read (ram cpu) (fromIntegral addr `rem` 0x0800)

writeCPURam :: Word16 -> Word8 -> Emulator ()
writeCPURam addr v = with cpu $ \cpu ->
  VUM.write (ram cpu) (fromIntegral addr `rem` 0x0800) v

readMapper :: Word16 -> Emulator Word8
readMapper addr = with mapper $ \mapper ->
  Mapper.read mapper addr

writeMapper :: Word16 -> Word8 -> Emulator ()
writeMapper addr value = with mapper $ \mapper ->
  Mapper.write mapper addr value

readController :: Emulator Word8
readController = with controller Controller.read

writeController :: Word8 -> Emulator ()
writeController value = with controller (`Controller.write` value)

storeKeys :: Set Controller.Key -> Emulator ()
storeKeys keys = with controller (`Controller.setKeysDown` keys)

loadKeys :: Emulator (Set Controller.Key)
loadKeys = with controller Controller.readKeysDown

writeNametableData :: Word16 -> Word8 -> Emulator ()
writeNametableData addr v = do
  mirror <- with cart $ \cart -> readIORef $ Cartridge.mirror cart
  let addr' = fromIntegral (mirroredNametableAddr addr mirror) `rem` 0x800
  with ppu $ \ppu ->
    VUM.write (nameTableData ppu) addr' v

readNametableData :: Word16 -> Emulator Word8
readNametableData addr = do
  mirror <- with cart $ \cart -> readIORef $ Cartridge.mirror cart
  let addr' = fromIntegral (mirroredNametableAddr addr mirror) `rem` 0x800
  with ppu $ \ppu ->
    VUM.read (nameTableData ppu) addr'

writePalette :: Word16 -> Word8 -> Emulator ()
writePalette addr value = with ppu $ \ppu ->
  VUM.write (paletteData ppu) (fromIntegral $ mirroredPaletteAddr addr) value

readPalette :: Word16 -> Emulator Word8
readPalette addr = with ppu $ \ppu ->
  VUM.read (paletteData ppu) (fromIntegral $ mirroredPaletteAddr addr)

readPPURegister :: Word16 -> Emulator Word8
readPPURegister addr = case addr of
  0x2002 -> readStatus
  0x2004 -> readOAMData'
  0x2007 -> readData
  other  -> pure 0

writePPURegister :: Word16 -> Word8 -> Emulator ()
writePPURegister addr v = do
  storePpu ppuRegister v
  case addr of
    0x2000 -> writeControl v
    0x2001 -> writeMask v
    0x2003 -> writeOAMAddress v
    0x2004 -> writeOAMData v
    0x2005 -> writeScroll v
    0x2006 -> writeAddress v
    0x2007 -> writeData v
    0x4014 -> writeDMA v
    _      -> error $ "Erroneous write detected at " ++ show addr ++ "!"

readStatus :: Emulator Word8
readStatus = do
  registerV <- loadPpu ppuRegister
  spriteOverflowV <- loadPpu spriteOverflow
  spriteZeroHitV <- loadPpu spriteZeroHit
  nmiOccurred' <- loadPpu nmiOccurred
  let r = registerV .&. 0x1F
  let r' = r .|. fromIntegral (fromEnum spriteOverflowV `shiftL` 5)
  let r'' = r' .|. fromIntegral (fromEnum spriteZeroHitV `shiftL` 6)
  let rFinal = r'' .|. fromIntegral (fromEnum nmiOccurred' `shiftL` 7)
  toggleNmi False
  storePpu writeToggle False
  pure $ fromIntegral rFinal

readOAMData' :: Emulator Word8
readOAMData' = do
  addr <- loadPpu oamAddress
  with ppu $ \ppu ->
    VUM.read (oamData ppu) (fromIntegral addr `rem` 0x0800)

readOAMData :: Word16 -> Emulator Word8
readOAMData addr = with ppu $ \ppu ->
  VUM.read (oamData ppu) (fromIntegral addr)

readData :: Emulator Word8
readData = do
  addr <- loadPpu currentVramAddress

  rv <- if (addr `rem` 0x4000) < 0x3F00 then do
    v <- readPpuMemory addr
    buffered <- loadPpu dataV
    storePpu dataV v
    pure buffered
  else do
    v <- readPpuMemory (addr - 0x1000)
    storePpu dataV v
    readPpuMemory addr

  incMode <- loadPpu incrementMode
  let inc = case incMode of
        Horizontal -> 1
        Vertical   -> 32
  modifyPpu currentVramAddress (+ inc)
  pure rv

writeData :: Word8 -> Emulator ()
writeData v = do
  addr <- loadPpu currentVramAddress
  writePPUMemory addr v
  incMode <- loadPpu incrementMode
  let inc = case incMode of
        Horizontal -> 1
        Vertical   -> 32
  modifyPpu currentVramAddress (+ inc)

writeDMA :: Word8 -> Emulator ()
writeDMA v = do
  let startingAddr = toWord16 v `shiftL` 8
  let addresses = fmap (+ startingAddr) [0..255]
  forM_ addresses (\addr -> do
    oamA <- loadPpu oamAddress
    oamV <- readCpuMemory8 addr
    with ppu $ \ppu ->
      VUM.write (oamData ppu) (toInt oamA) oamV
    modifyPpu oamAddress (+ 1))

writeControl :: Word8 -> Emulator ()
writeControl v = do
  storePpu nameTable $ case (v `shiftR` 0) .&. 3 of
    0 -> 0x2000
    1 -> 0x2400
    2 -> 0x2800
    3 -> 0x2C00
  storePpu incrementMode $ if testBit v 2 then Vertical else Horizontal
  storePpu spriteTable $ if testBit v 3 then 0x1000 else 0x0000
  storePpu bgTable $ if testBit v 4 then 0x1000 else 0x0000
  storePpu spriteSize $ if testBit v 5 then Double else Normal

  storePpu nmiOutput $ testBit v 7
  nmiOccurred' <- loadPpu nmiOccurred
  toggleNmi nmiOccurred'
  tv <- loadPpu tempVramAddress
  storePpu tempVramAddress ((tv .&. 0xF3FF) .|. (toWord16 v .&. 0x03) `shiftL` 10)

writeMask :: Word8 -> Emulator ()
writeMask v = do
  storePpu colorMode $ if testBit v 0 then Grayscale else Color
  storePpu leftBgVisibility $ if testBit v 1 then Shown else Hidden
  storePpu leftSpritesVisibility $ if testBit v 2 then Shown else Hidden
  storePpu bgVisibility $ testBit v 3
  storePpu spriteVisibility $ testBit v 4
  storePpu intensifyReds $ testBit v 5
  storePpu intensifyGreens $ testBit v 6
  storePpu intensifyBlues $ testBit v 7

writeOAMAddress :: Word8 -> Emulator ()
writeOAMAddress = storePpu oamAddress

writeOAMData :: Word8 -> Emulator ()
writeOAMData v = do
  addr <- loadPpu oamAddress
  with ppu $ \ppu ->
    VUM.write (oamData ppu) (toInt addr) v
  modifyPpu oamAddress (+ 1)

writeScroll :: Word8 -> Emulator ()
writeScroll v = do
  wv <- loadPpu writeToggle
  tv <- loadPpu tempVramAddress
  if wv then do
    let tv' = (tv .&. 0x8FFF) .|. ((toWord16 v .&. 0x07) `shiftL` 12)
    let tv'' = (tv' .&. 0xFC1F) .|. ((toWord16 v .&. 0xF8) `shiftL` 2)
    storePpu tempVramAddress tv''
    storePpu writeToggle False
  else do
    let tv' = (tv .&. 0xFFE0) .|. (toWord16 v `shiftR` 3)
    storePpu tempVramAddress tv'
    storePpu fineX $ v .&. 0x07
    storePpu writeToggle True

writeAddress :: Word8 -> Emulator ()
writeAddress v = do
  wv <- loadPpu writeToggle
  tv <- loadPpu tempVramAddress
  if wv then do
    let tv' = (tv .&. 0xFF00) .|. toWord16 v
    storePpu tempVramAddress tv'
    storePpu currentVramAddress tv'
    storePpu writeToggle False
  else do
    let tv' = (tv .&. 0x80FF) .|. ((toWord16 v .&. 0x3F) `shiftL` 8)
    storePpu tempVramAddress tv'
    storePpu writeToggle True

newPPU :: IO PPU
newPPU = do
  -- Misc
  cycles <- newIORef 0
  scanline <- newIORef 0
  frameCount <- newIORef 0
  writeToggle <- newIORef False
  ppuRegister <- newIORef 0x0
  oddFrame <- newIORef False
  -- Data
  oamData <- VUM.replicate 0x100 0x0
  nameTableData <- VUM.replicate 0x800 0x0
  paletteData <- VUM.replicate 0x20 0x0
  screen <- VUM.replicate (256 * 240 * 3) 255
  -- Addresses
  currentVramAddress <- newIORef 0x0
  tempVramAddress <- newIORef 0x0
  oamAddress <- newIORef 0x0
  -- NMI
  nmiOutput <- newIORef False
  nmiOccurred <- newIORef False
  nmiDelay <- newIORef 0
  nmiPrevious <- newIORef False
  -- Control register
  nameTable <- newIORef 0x2000
  incrementMode <- newIORef Horizontal
  spriteTable <- newIORef 0x0000
  bgTable <- newIORef 0x0000
  spriteSize <- newIORef Normal
  -- Mask register
  colorMode <- newIORef Color
  leftBgVis <- newIORef Hidden
  leftSpritesVis <- newIORef Hidden
  bgVis <- newIORef False
  spriteVis <- newIORef False
  intensifyReds <- newIORef False
  intensifyGreens <- newIORef False
  intensifyBlues <- newIORef False
  -- Status register
  spriteOverflow <- newIORef False
  spriteZeroHit <- newIORef False
  -- Scroll register
  fineX <- newIORef 0x0
  -- Data register
  dataV <- newIORef 0x0
  -- Temp vars
  nameTableByte <- newIORef 0x0
  attrTableByte <- newIORef 0x0
  loTileByte <- newIORef 0x0
  hiTileByte <- newIORef 0x0
  tileData <- newIORef 0x0
  sprites <- newIORef V.empty

  pure $ PPU
    -- Misc
    cycles scanline frameCount writeToggle ppuRegister oddFrame
    -- Data
    oamData nameTableData paletteData screen
    -- Addresses
    currentVramAddress tempVramAddress oamAddress
    -- NMI
    nmiOutput nmiOccurred nmiDelay nmiPrevious
    -- Control register
    nameTable incrementMode spriteTable bgTable spriteSize
    -- Mask register
    colorMode leftBgVis leftSpritesVis bgVis spriteVis
    intensifyReds intensifyGreens intensifyBlues
    -- Status register
    spriteOverflow spriteZeroHit
    -- Scroll register
    fineX
    -- Data register
    dataV
    -- Temp vars
    nameTableByte attrTableByte loTileByte hiTileByte tileData sprites

mirroredPaletteAddr :: Word16 -> Word16
mirroredPaletteAddr addr = if addr' >= 16 && addr' `rem` 4 == 0 then addr' - 16 else addr'
  where addr' = addr `rem` 32

mirroredNametableAddr :: Word16 -> Mirror -> Word16
mirroredNametableAddr addr mirror = 0x2000 + lookup + offset
  where addr' = (addr - 0x2000) `rem` 0x1000
        tableIndex = fromIntegral $ addr' `div` 0x0400
        lookup = ((nameTableMirrorLookup V.! fromEnum mirror) V.! tableIndex) * 0x0400
        offset = fromIntegral $ addr' `rem` 0x0400
        nameTableMirrorLookup = V.fromList (fmap V.fromList [
            [0, 0, 1, 1],
            [0, 1, 0, 1],
            [0, 0, 0, 0],
            [1, 1, 1, 1],
            [0, 1, 2, 3]
          ])

loadScreen :: Emulator (VUM.IOVector Word8)
loadScreen = with ppu $ \ppu ->
  pure $ screen ppu

writeScreen :: Coords -> Color -> Emulator ()
writeScreen (x, y) (r, g, b) = with ppu $ \ppu -> do
  let offset = (x + (y * 256)) * 3
  VUM.write (screen ppu) (offset + 0) r
  VUM.write (screen ppu) (offset + 1) g
  VUM.write (screen ppu) (offset + 2) b

toggleNmi :: Bool -> Emulator ()
toggleNmi occurred = do
  storePpu nmiOccurred occurred
  output <- loadPpu nmiOutput
  occurred <- loadPpu nmiOccurred
  previous <- loadPpu nmiPrevious
  let nmi = output && occurred

  when (nmi && not previous) $
    storePpu nmiDelay 15

  storePpu nmiPrevious nmi