AoC/2019/day17/Intcode.hs

module Intcode
  ( TuringMachine(TM, tape, pointer, pointerOffset, output, input,
              state)
  , step
  , defaultTM
  , TMOutState(Continue, AwaitInput, Halt)
  , execSteps
  , readIntcode
  ) where

import Control.DeepSeq as DeepSeq
import Data.Char
import Data.List as L
import Data.List.Split
import Data.Maybe
import Debug.Trace
import Data.Vector as V hiding
  ( (++)
  , concatMap
  , elem
  , head
  , last
  , length
  , map
  , reverse
  , splitAt
  )
import qualified Data.Vector as V ((++), last)

data TMOutState
  = Continue
  | AwaitInput
  | Halt
  deriving (Enum, Eq, Show)

data Mode
  = Position
  | Immediate
  | Relative
  deriving (Enum, Eq, Show)

data TuringMachine =
  TM
    { tape :: Vector Integer
    , pointer :: Integer
    , pointerOffset :: Integer
    , output :: [Integer]
    , input :: Maybe Integer
    , state :: TMOutState
    }
  deriving (Show)

readIntcode :: String -> V.Vector Integer
readIntcode = V.fromList . concatMap (map read . splitOn ",") . lines

defaultTM :: Maybe (Int, Integer) -> Vector Integer -> TuringMachine
defaultTM replacement t =
  TM
    { tape = maybe (tapePreprocess t (0,t ! 0)) (tapePreprocess t) replacement
    , pointer = 0
    , pointerOffset = 0
    , output = []
    , input = Nothing
    , state = Continue
    }

tapePreprocess :: Vector Integer -> (Int, Integer) -> Vector Integer
tapePreprocess t (target, repl) =
  flip (//) [(target, repl)] $ (V.++) t $ V.replicate 22000 0

opLength :: Integer -> Integer
opLength x
  | n `elem` "1278" = 4
  | n `elem` "56" = 3
  | n `elem` "349" = 2
  | otherwise = 1
  where
    n = last $ show x

parseModes :: String -> [Mode]
parseModes m = L.replicate (3 - length l) Position ++ l
  where
    l = map (toEnum . digitToInt) m

paramChange ::
     [Mode] -> Integer -> Vector Integer -> Vector Integer -> Vector Integer
paramChange m rbase opvec t = imap f (V.tail opvec)
  where
    f i a =
      case m !! i of
        Immediate -> a
        Position -> t ! fromInteger a
        Relative -> t ! fromInteger (a + rbase)

getOpModes :: Vector Integer -> (String, [Mode])
getOpModes opvec = (op_dedup, parsed_modes)
  where
    (op, modes) = splitAt 2 $ reverse $ show $ opvec ! 0
    parsed_modes = reverse $ parseModes $ reverse modes
    op_dedup =
      if last op == '0'
        then [head op]
        else op

step :: TuringMachine -> TuringMachine
step tm =
  case op of
    "1" -> tmBinop (+)
    "2" -> tmBinop (*)
    "3" ->
      if isJust (input tm)
        then (getNewTM (fromJust $ input tm)) {input = Nothing}
        else tm {state = AwaitInput}
    "4" -> tmn {output = V.last params : output tm}
    "5" ->
      tm
        { pointer =
            if params ! 0 /= 0
              then params ! 1
              else pNew
        }
    "6" ->
      tm
        { pointer =
            if params ! 0 == 0
              then params ! 1
              else pNew
        }
    "7" ->
      tmBinop
        (\x y ->
           if x < y
             then 1
             else 0)
    "8" ->
      tmBinop
        (\x y ->
           if x == y
             then 1
             else 0)
    "9" -> tmn {pointerOffset = pointerOffset tm + (params ! 0)}
    "99" -> tm {state = Halt}
    _ -> error "Illegal Opcode."
  where
    pNew = pointer tm + opLength (tape tm ! fromInteger (pointer tm))
    tmn = tm {pointer = pNew}
    opvec =
      slice
        (fromInteger (pointer tm))
        (fromInteger (pNew - pointer tm))
        (tape tm)
    (op, m) = getOpModes opvec
    params = paramChange m (pointerOffset tm) opvec (tape tm)
    tmBinop x = getNewTM ((params ! 0) `x` (params ! 1))
    {-without the following DeepSeq call, thunks build up eternally and
    the vectors won’t be garbage collected: >4GB RAM usage, god knows
    how much with larger tapes (my laptop crashed), now it’s a cozy ~20mB-}
    getNewTM x =
      tmn {tape = DeepSeq.force (tape tm // [(target, x)]), state = Continue}
    target =
      fromInteger $
      case m !! (length params - 1) of
        Relative -> V.last opvec + pointerOffset tm
        _ -> V.last opvec

execSteps :: TuringMachine -> TuringMachine
execSteps tm =
  case state tmNew of
    Continue -> execSteps tmNew
    _ -> tmNew
  where
    tmNew = step tm