{-# LANGUAGE FlexibleContexts, RecordWildCards, TypeApplications #-}
module Day16 (day16a, day16b) where
import Control.Arrow ((&&&))
import Control.Monad (foldM)
import Data.Array.Unboxed (IArray, UArray, Ix, (!), (//), listArray)
import Data.Bits (Bits, (.&.), (.|.))
import Data.Bool (bool)
import Data.List (genericLength)
import Data.Map.Lazy ((!?), elems, fromListWith, union)
import qualified Data.Map.Lazy as M (empty, null, partition)
import Data.Set (difference, findMin, fromList, intersection, size, unions)
import Text.Megaparsec (MonadParsec, between, count, many, parseMaybe, sepBy, sepEndBy1)
import Text.Megaparsec.Char (char, newline, string)
import Text.Megaparsec.Char.Lexer (decimal)
data Op
= ADDR | ADDI | MULR | MULI | BANR | BANI | BORR | BORI
| SETR | SETI | GTIR | GTRI | GTRR | EQIR | EQRI | EQRR
deriving (Bounded, Enum, Eq, Ord)
data Sample a op i = Sample
{ sampleOp :: op
, sampleA :: i
, sampleB :: i
, sampleC :: i
, sampleR0 :: a i i
, sampleR1 :: a i i
}
unknownInstructionParser :: (MonadParsec e String m, Integral op, Integral i) => m (op, i, i, i)
unknownInstructionParser = do
op <- decimal
[a, b, c] <- count 3 $ char ' ' *> decimal
return (op, a, b, c)
sampleParser :: (MonadParsec e String m, Integral op, Integral i, Ix i, IArray a i) => m (Sample a op i)
sampleParser = do
r0 <- between (string "Before: [") (string "]") (sepBy decimal (string ", ")) <* newline
(sampleOp, sampleA, sampleB, sampleC) <- unknownInstructionParser <* newline
r1 <- between (string "After: [") (string "]") (sepBy decimal (string ", ")) <* newline
return Sample
{ sampleR0 = listArray (0, genericLength r0 - 1) r0
, sampleR1 = listArray (0, genericLength r1 - 1) r1
, ..
}
parser :: (MonadParsec e String m, Integral op, Integral i, Ix i, IArray a i) => m ([Sample a op i], [(op, i, i, i)])
parser = (,)
<$> (sepEndBy1 sampleParser newline <* many newline)
<*> sepEndBy1 unknownInstructionParser newline
doOp :: (IArray a i, Bits i, Ix i, Num i) => Op -> i -> i -> i -> a i i -> a i i
doOp ADDR a b c r = r // [(c, r ! a + r ! b)]
doOp ADDI a b c r = r // [(c, r ! a + b)]
doOp MULR a b c r = r // [(c, r ! a * r ! b)]
doOp MULI a b c r = r // [(c, r ! a * b)]
doOp BANR a b c r = r // [(c, r ! a .&. r ! b)]
doOp BANI a b c r = r // [(c, r ! a .&. b)]
doOp BORR a b c r = r // [(c, r ! a .|. r ! b)]
doOp BORI a b c r = r // [(c, r ! a .|. b)]
doOp SETR a _ c r = r // [(c, r ! a)]
doOp SETI a _ c r = r // [(c, a)]
doOp GTIR a b c r = r // [(c, bool 0 1 $ a > r ! b)]
doOp GTRI a b c r = r // [(c, bool 0 1 $ r ! a > b)]
doOp GTRR a b c r = r // [(c, bool 0 1 $ r ! a > r ! b)]
doOp EQIR a b c r = r // [(c, bool 0 1 $ a == r ! b)]
doOp EQRI a b c r = r // [(c, bool 0 1 $ r ! a == b)]
doOp EQRR a b c r = r // [(c, bool 0 1 $ r ! a == r ! b)]
validOps :: (IArray a i, Eq (a i i), Bits i, Ix i, Num i) => Sample a op i -> [Op]
validOps Sample {..} =
[op | op <- [minBound..maxBound], doOp op sampleA sampleB sampleC sampleR0 == sampleR1]
day16a :: String -> Maybe Int
day16a input = do
(samples, _) <- parseMaybe @() parser input
return $ length $ filter ambiguous samples
where ambiguous sample = case validOps @UArray @Int @Int sample of
_:_:_:_ -> True
_ -> False
day16b :: String -> Maybe Int
day16b input = do
(samples, instructions) <- parseMaybe @() parser input
codings <- re M.empty $ fromListWith intersection $
(sampleOp &&& fromList . validOps @UArray @Int @Int) <$> samples
(! 0) <$> foldM (doOp' codings) (listArray @UArray @Int @Int (0, 3) [0, 0, 0, 0]) instructions
where re done pending
| M.null pending = return done
| (done', pending') <- M.partition ((== 1) . size) pending, not $ M.null done' =
re (union done $ findMin <$> done') $
flip difference (unions $ elems done') <$> pending'
| otherwise = fail "can not reverse engineer instruction encoding"
doOp' codings r (op, a, b, c) = do
op' <- codings !? op
return $ doOp op' a b c r