aoc2020/day8.hs

{-# LANGUAGE FlexibleContexts #-}

import qualified Data.Vector as V
import Text.Parsec

data Op = Acc Int
        | Jmp Int
        | Nop Int
        deriving Show

main :: IO ()
main = do
  raw <- readFile "day8.txt"
  let ops = V.fromList $ zip (repeat False)
              $ validate $ map (parse lineP []) $ lines raw
      ansA = solveA ops
      ansB = solveB ops
   in do
     putStrLn $ "day8a: " ++ (show ansA)
     putStrLn $ "day8b: " ++ (show ansB)

solveA :: V.Vector (Bool, Op) -> Int
solveA ops = runMachineA ops 0 0

solveB :: V.Vector (Bool, Op) -> Int
solveB ops = case testSwaps ops 0 of
               Nothing -> error "unsolvable?"
               Just x -> x

runMachineA :: V.Vector (Bool, Op) -> Int -> Int -> Int
runMachineA ops pc ac = do
  let (check,op) = ops V.! pc
   in case check of
    True -> ac
    _ -> do
      let nac = case op of
                  (Acc x) -> ac + x
                  _ -> ac
          npc = case op of
                  (Jmp y) -> pc + y
                  _ -> succ pc
          nops = ops V.// [(pc, (True, op))]
        in runMachineA nops npc nac

swapOp :: V.Vector (Bool, Op) -> Int -> Maybe (V.Vector (Bool, Op))
swapOp ops ind =
  case op of
    (Acc _) -> Nothing
    (Nop x) -> Just $ ops V.// [(ind, (check, Jmp x))]
    (Jmp x) -> Just $ ops V.// [(ind, (check, Nop x))]
  where
    (check,op) = ops V.! ind

testSwaps :: V.Vector (Bool, Op) -> Int -> Maybe Int
testSwaps ops ind
  | ind >= V.length ops = Nothing
  | otherwise =
    case res of
      Nothing -> testSwaps ops $ succ ind
      (Just nops) -> case (runMachineB nops 0 0) of
                       Nothing -> testSwaps ops $ succ ind
                       x -> x

  where
    res = swapOp ops ind

runMachineB :: V.Vector (Bool, Op) -> Int -> Int -> Maybe Int
runMachineB ops pc ac =
  if pc == V.length ops
  then Just ac
  else do
  let (check,op) = ops V.! pc
   in case check of
    True -> Nothing
    _ -> do
      let nac = case op of
                  (Acc x) -> ac + x
                  _ -> ac
          npc = case op of
                  (Jmp y) -> pc + y
                  _ -> succ pc
          nops = ops V.// [(pc, (True, op))]
        in runMachineB nops npc nac

validate :: [Either ParseError Op] -> [Op]
validate [] = []
validate ((Left _):_) = error "invalid input"
validate ((Right b):bs) = b:(validate bs)

lineP :: Parsec String () Op
lineP = do
  o <- opP
  whitespaces
  i <- intP
  whitespaces
  eof
  return $ o i

intP :: Parsec String () Int
intP = do
  s <- oneOf ['+','-']
  v <- many1 digit
  return $ (if s == '-' then -1 else 1) * (read v)

opP :: Parsec String () (Int -> Op)
opP = do
  o <- (string "acc") <|> (string "jmp") <|> (string "nop")
  case o of
    "acc" -> return Acc
    "jmp" -> return Jmp
    _ -> return Nop

whitespaces :: Parsec String () String
whitespaces = many $ char ' '