Day5.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
module Day5 where

import Prelude hiding (read)
import qualified Data.Text as T
import qualified Data.Vector as V
import Control.Monad.State
import Control.Monad.Fail
import Data.Sequence hiding (replicate, reverse)
import Data.List (scanl')

import Util


data Parameter
  = Immediate Int
  | Position Int
  deriving (Show, Eq)


data Opcode
  = Add Parameter Parameter Parameter
  | Mul Parameter Parameter Parameter
  | Read Parameter
  | Write Parameter
  | JT Parameter Parameter
  | JF Parameter Parameter
  | LessThan Parameter Parameter Parameter
  | Equal Parameter Parameter Parameter
  | Halt

deriving instance Show Opcode
deriving instance Eq Opcode

data VM
  = VM
  { memory  :: !(V.Vector Int)
  , ip      :: !Int
  , input   :: ![Int]
  , output  :: !(Seq Int)
  }
  deriving (Show, Eq)

data DecoderState
  = DecoderState
  { modes   :: ![Int -> Parameter]
  , instr   :: ![Int]
  }

size :: Opcode -> Int
size Halt        = 1
size Read     {} = 2
size Write    {} = 2
size JT       {} = 3
size JF       {} = 3
size LessThan {} = 4
size Equal    {} = 4
size Add      {} = 4
size Mul      {} = 4


initVM :: [Int] -> [Int] -> VM
initVM program inp =
  VM
  { memory = V.fromList program
  , ip     = 0
  , input  = inp
  , output = empty
  }

loop :: MonadState VM m => m ()
loop =
  fetch >>= decode >>= maybe (return ()) handle
  -- undefined
 where
  handle i = do
    old <- getIp
    execute i
    new <- getIp
    when (new == old) $ incIp i
    when (i /= Halt) loop

incIp :: MonadState VM m => Opcode -> m ()
incIp op = getIp >>= setIp . (+ size op)

getIp :: MonadState VM m => m Int
getIp = gets ip

setIp :: MonadState VM m => Int -> m ()
setIp i = modify $ \s -> s { ip = i }

fetch :: MonadState VM m => m Int
fetch = do
  i <- gets ip
  m <- gets memory
  return $ m V.! i


decode :: MonadState VM m => Int -> m (Maybe Opcode)
decode op = do
  next3 <- V.toList <$> gets (\s -> V.slice (ip s + 1) 3 (memory s))
  let decoderState = DecoderState modes' next3
  return $ evalStateT (parseOpcode opcode) decoderState
 where
   (rawModes, opcode) = op `divMod` 100
   modes' =
     reverse
     . map (parseMode . fst)
     . tail
     $ scanl' (\(_, x) d -> x `divMod` d) (0, rawModes) [100, 10, 1]
   parseMode :: Int -> Int -> Parameter
   parseMode 1 = Immediate
   parseMode 0 = Position
   parseMode m = error $ "Undefined mode: " ++ show m

   parseOpcode :: (MonadFail m, MonadState DecoderState m) => Int -> m Opcode
   parseOpcode 1  = Add      <$> popParam <*> popParam <*> popParam
   parseOpcode 2  = Mul      <$> popParam <*> popParam <*> popParam
   parseOpcode 3  = Read     <$> popParam
   parseOpcode 4  = Write    <$> popParam
   parseOpcode 5  = JT       <$> popParam <*> popParam
   parseOpcode 6  = JF       <$> popParam <*> popParam
   parseOpcode 7  = LessThan <$> popParam <*> popParam <*> popParam
   parseOpcode 8  = Equal    <$> popParam <*> popParam <*> popParam
   parseOpcode 99 = return Halt
   parseOpcode o = error $ "Unknown opcode: " ++ show o


popParam :: (MonadFail m, MonadState DecoderState m) => m Parameter
popParam = do
  (i:is) <- gets instr
  (m:ms) <- gets modes
  modify $ \s -> s { instr = is, modes = ms }
  return $ m i

readMem :: MonadState VM m => Parameter -> m Int
readMem (Position p) =
  gets ((V.! p) . memory)
readMem _ = error "Can't read from immediate"

writeMem :: MonadState VM m => Parameter -> Int -> m ()
writeMem (Position p) x =
  modify (\s -> s { memory = memory s V.// [(p, x)]  })
writeMem _ _ = error "Can't write to immediate"

resolve :: MonadState VM m => Parameter -> m Int
resolve (Immediate x)  = return x
resolve p@(Position _) = readMem p

execute :: MonadState VM m => Opcode -> m ()
execute (Read x)  =
  popInput >>= writeMem x
execute (Write p) =
  resolve p >>= \v -> modify $ \s -> s { output = output s |> v }
execute (Add x y o) =
  binop (+) x y o
execute (Mul x y o) =
  binop (*) x y o
execute (LessThan x y o) =
  binop (\x' y' -> fromEnum $ x' < y') x y o
execute (Equal x y o) =
  binop (\x' y' -> fromEnum $ x' == y') x y o
execute (JT b d) =
  resolve b >>= \v -> when (v /= 0) (resolve d >>= setIp)
execute (JF b d) =
  resolve b >>= \v -> when (v == 0) (resolve d >>= setIp)
execute Halt =
  return ()

binop
  :: MonadState VM m
  => (Int -> Int -> Int)
  -> Parameter
  -> Parameter
  -> Parameter
  -> m ()
binop fn x y o =
  liftM2 fn (resolve x) (resolve y) >>= writeMem o

popInput :: MonadState VM m => m Int
popInput = do
  i <- gets input
  modify $ \s -> s { input = tail i }
  return $ head i


setup :: Int -> IO VM
setup i = map readText <$> readInput "day5" (T.split (== ',')) >>= \p -> return $ initVM p [i]

run :: IO ()
run = do
  setup 1 >>= print . viewr . output . execState loop
  setup 5 >>= print . viewr . output . execState loop
  return ()