Advent of Code 2020 Day 8

{-# language TupleSections #-}
{-# language RecordWildCards #-}

import Control.Applicative ((<|>))
import Control.Monad ((<=<))
import Data.Bool (bool)
import Data.List.NonEmpty (NonEmpty((:|)), nonEmpty)
import qualified Data.List.NonEmpty as NE
import Data.Maybe (listToMaybe)

data Instruction
  = Nop !Int
  | Acc !Int
  | Jmp !Int
  deriving (Eq, Show)

parseInstruction :: String -> Maybe Instruction
parseInstruction s =
  case splitAt 3 s of
   (op, _:s:n) -> parseOp op <*> parseNum s n
   _ -> Nothing
 where
  parseOp "nop" = Just Nop
  parseOp "acc" = Just Acc
  parseOp "jmp" = Just Jmp
  parseOp _ = Nothing
  parseNum '+' n = fmap fst . listToMaybe $ reads n
  parseNum '-' n = fmap (negate . fst) . listToMaybe $ reads n
  parseNum _ _ = Nothing

data Console a = MkConsole
  { accumulator :: !Int
  , rPast :: [(Instruction, a)]
  , ip :: !(Instruction, a)
  , future :: [(Instruction, a)]
  } deriving (Eq, Show)

initConsole :: a -> NonEmpty Instruction -> Console a
initConsole x (h:|t) =
  MkConsole
  { accumulator = 0
  , rPast = []
  , ip = (h, x)
  , future = map (,x) t
  }

stepConsole :: (a -> Maybe a) -> Console a -> Either a (Either Int (Console a))
stepConsole f c@MkConsole{..} =
  case f $ snd ip of
   Nothing -> Left $ snd ip
   Just x ->
    Right $ case fst ip of
     Nop _ ->
      case future of
       [] -> Left accumulator
       (ip':future') -> Right c{ rPast = (fst ip, x) : rPast, ip = ip', future = future' }
     Acc d -> let accumulator' = accumulator + d in
      case future of
       [] -> Left accumulator'
       (ip':future') ->
        Right c
        { accumulator = accumulator'
        , rPast = (fst ip, x) : rPast
        , ip = ip'
        , future = future'
        }
     Jmp o ->
      case compare o 0 of
       LT ->
        let
          o' = abs o
          (iPast, tPast) = splitAt o' rPast
          (ip':future') = reverse iPast ++ (fst ip, x) : future
         in
        case reverse iPast ++ (fst ip, x) : future of
         [] -> Left accumulator
         (ip':future') -> Right c{ rPast = tPast, ip = ip', future = future' }
       EQ -> Right c{ ip = (fst ip, x) }
       GT -> let (iFuture, tFuture) = splitAt o ((fst ip, x) : future) in
        case tFuture of
         [] -> Left accumulator
         (ip':future') -> Right c{ rPast = reverse iFuture ++ rPast, ip = ip', future = future' }

runConsole :: Console Bool -> Int
runConsole c =
  case stepConsole (bool (Just True) Nothing) c of
   Left _ -> accumulator c
   Right (Left x) -> x
   Right (Right c') -> runConsole c'

testConsole :: Console Bool -> Maybe Int
testConsole = testConsole' False
 where
  testConsole' False c =
    case ip c of
     (Nop x, y) -> go False c <|> go True c{ ip = (Jmp x, y) }
     (Acc _, _) -> go False c
     (Jmp x, y) -> go False c <|> go True c{ ip = (Nop x, y) }
  testConsole' True c = go True c
  go v c = case stepConsole (bool (Just True) Nothing) c of
   Left _ -> Nothing
   Right (Left x) -> Just x
   Right (Right c') -> testConsole' v c'

part1 = interact (show . fmap (runConsole . initConsole False) . (nonEmpty <=< traverse parseInstruction) . lines)

main = interact (show . (testConsole . initConsole False <=< nonEmpty <=< traverse parseInstruction) . lines)