{-# LANGUAGE LambdaCase, NamedFieldPuns, PatternGuards, RecordWildCards, TypeApplications, ViewPatterns #-}
{-# OPTIONS_HADDOCK ignore-exports #-}
module Day23 (day23a, day23b) where
import Data.Array.IArray (IArray, (!), (//), listArray)
import Data.Array.Unboxed (UArray)
import Data.Ix (Ix)
import Data.List (findIndices)
import Math.NumberTheory.Primes (isPrime)
import Text.Read (readMaybe)
data Ins reg imm
= Set {reg :: reg, val :: Either reg imm}
| Sub {reg :: reg, val :: Either reg imm}
| Mul {reg :: reg, val :: Either reg imm}
| Jnz {cnd :: Either reg imm, jmp :: Either reg imm}
data State a reg imm = State {pc :: Int, regs :: a reg imm}
parseVal :: (Read imm) => String -> Either Char imm
parseVal (readMaybe -> Just imm) = Right imm
parseVal [reg] = Left reg
parse :: (Read imm) => String -> [Ins Char imm]
parse = map (parseIns . words) . lines where
parseIns ["set", [reg], parseVal -> val] = Set {..}
parseIns ["sub", [reg], parseVal -> val] = Sub {..}
parseIns ["mul", [reg], parseVal -> val] = Mul {..}
parseIns ["jnz", parseVal -> cnd, parseVal -> jmp] = Jnz {..}
step :: (IArray a imm, Ix reg, Integral imm) =>
[Ins reg imm] -> State a reg imm -> Maybe (State a reg imm)
step ins State {pc} | pc < 0 || null (drop pc ins) = Nothing
step ins state@State {..} = Just $ case ins !! pc of
Set {..} -> State (pc + 1) (regs // [(reg, load val)])
Sub {..} -> State (pc + 1) (regs // [(reg, regs ! reg - load val)])
Mul {..} -> State (pc + 1) (regs // [(reg, regs ! reg * load val)])
Jnz {..}
| load cnd == 0 -> state {pc = pc + 1}
| otherwise -> state {pc = pc + fromIntegral (load jmp)}
where
load = either (regs !) id
stepOptimized :: (IArray a imm, Ix reg, Integral imm) =>
[Ins reg imm] -> State a reg imm -> Maybe (State a reg imm)
stepOptimized ins State {pc} | pc < 0 = Nothing
stepOptimized ins state@State {..} = case drop pc ins of
( Set d (Right 2) :
Set e (Right 2) :
Set g (Left ((== d) -> True)) :
Mul ((== g) -> True) (Left ((== e) -> True)) :
Sub ((== g) -> True) b :
Jnz (Left ((== g) -> True)) (Right 2) :
Set f f0 :
Sub ((== e) -> True) (Right (-1)) :
Set ((== g) -> True) (Left ((== e) -> True)) :
Sub ((== g) -> True) ((== b) -> True) :
Jnz (Left ((== g) -> True)) (Right (-8)) :
Sub ((== d) -> True) (Right (-1)) :
Set ((== g) -> True) (Left ((== d) -> True)) :
Sub ((== g) -> True) ((== b) -> True) :
Jnz (Left ((== g) -> True)) (Right (-13)) : _)
| d /= e, d /= g, d /= f, e /= g, e /= f, g /= f
, either (not . (`elem` [d, e, g, f])) (const True) b
, either (not . (`elem` [d, e, g, f])) (const True) f0
, either Just (const Nothing) b /= either Just (const Nothing) f0 ->
let b' = either (regs !) id b in Just . State (pc + 15) $ regs //
( [(d, b'), (e, b'), (g, 0)] ++
[(f, either (regs !) id f0) | not . isPrime $ fromIntegral b'] )
_ -> step ins state
iterateMaybe :: (a -> Maybe a) -> Maybe a -> [a]
iterateMaybe f = maybe [] $ (:) <*> iterateMaybe f . f
day23a :: String -> Int
day23a input = countMuls $ iterateMaybe (step ins) $ Just state0 where
ins = parse input
muls = findIndices (\case Mul {} -> True; _ -> False) ins
countMuls = length . filter (`elem` muls) . map pc
state0 = State {pc = 0, regs = listArray @UArray @Int ('a', 'h') $ repeat 0}
day23b :: String -> Int
day23b input =
regs (last . iterateMaybe (stepOptimized ins) $ Just state0) ! 'h' where
ins = parse input
state0 =
State {pc = 0, regs = listArray @UArray @Int ('a', 'h') $ 1 : repeat 0}