src/advent16/advent16.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Data.Text (Text)
   4 import qualified Data.Text as T
   5 import qualified Data.Text.IO as TIO
   6
   7 import Data.Void (Void)
   8
   9 import Text.Megaparsec
  10 import Text.Megaparsec.Char
  11 import qualified Text.Megaparsec.Char.Lexer as L
  12 import qualified Control.Applicative as CA
  13
  14 import Data.List
  15 import qualified Data.Set as S
  16
  17 import qualified Data.Map.Strict as M
  18 import Data.Map.Strict ((!))
  19 import Data.Bits ((.&.), (.|.))
  20
  21 type Memory = M.Map Int Int
  22 data Operation =
  23       Addr
  24     | Addi
  25     | Mulr
  26     | Muli
  27     | Banr
  28     | Bani
  29     | Borr
  30     | Bori
  31     | Setr
  32     | Seti
  33     | Gtir
  34     | Gtri
  35     | Gtrr
  36     | Eqir
  37     | Eqri
  38     | Eqrr
  39     deriving (Eq, Show, Enum, Bounded, Ord)
  40 data RawOperation = RawOperation {_opcode :: Int, _a :: Int, _b :: Int, _c :: Int} deriving (Eq, Show)
  41 data Test = Test {_preMemory :: Memory, _testOperation :: RawOperation, _postMemory :: Memory} deriving (Eq, Show)
  42 type Candidates = M.Map Int (S.Set Operation)
  43 type OpcodeAssignments = M.Map Int Operation
  44
  45 main :: IO ()
  46 main = do
  47     text <- TIO.readFile "data/advent16.txt"
  48     let (tests, program) = successfulParse text
  49     print $ part1 tests
  50     print $ part2 tests program
  51
  52 part1 :: [Test] -> Int
  53 part1 = length . filter (>= 3) . map S.size . map matches
  54
  55 part2 :: [Test] -> [RawOperation] -> Int
  56 part2 tests program = finalMemory!0
  57     where candidates = candidateOpcodes tests
  58           opcodes = findOpcodes candidates
  59           finalMemory = foldl' (\m op -> execute op opcodes m) (M.singleton 0 0) program
  60
  61
  62 execute :: RawOperation -> OpcodeAssignments -> Memory -> Memory
  63 execute op codes m = perform o (_a op) (_b op) (_c op) m
  64     where o = codes!(_opcode op)
  65
  66 perform :: Operation -> Int -> Int -> Int -> Memory -> Memory
  67 perform Addr a b c memory = M.insert c (memory!a + memory!b) memory
  68 perform Addi a b c memory = M.insert c (memory!a + b) memory
  69 perform Mulr a b c memory = M.insert c (memory!a * memory!b) memory
  70 perform Muli a b c memory = M.insert c (memory!a * b) memory
  71 perform Banr a b c memory = M.insert c (memory!a .&. memory!b) memory
  72 perform Bani a b c memory = M.insert c (memory!a .&. b) memory
  73 perform Borr a b c memory = M.insert c (memory!a .|. memory!b) memory
  74 perform Bori a b c memory = M.insert c (memory!a .|. b) memory
  75 perform Setr a b c memory = M.insert c (memory!a) memory
  76 perform Seti a b c memory = M.insert c a memory
  77 perform Gtir a b c memory = M.insert c (if a > (memory!b) then 1 else 0) memory
  78 perform Gtri a b c memory = M.insert c (if (memory!a) > b then 1 else 0) memory
  79 perform Gtrr a b c memory = M.insert c (if (memory!a) > (memory!b) then 1 else 0) memory
  80 perform Eqir a b c memory = M.insert c (if a == memory!b then 1 else 0) memory
  81 perform Eqri a b c memory = M.insert c (if (memory!a) == b then 1 else 0) memory
  82 perform Eqrr a b c memory = M.insert c (if (memory!a) == (memory!b) then 1 else 0) memory
  83
  84 doTest :: Test -> Operation -> Bool
  85 doTest test operation = calculatedMemory == (_postMemory test)
  86     where rawOp = _testOperation test
  87           calculatedMemory = perform operation (_a rawOp) (_b rawOp) (_c rawOp) (_preMemory test)
  88
  89 matches :: Test -> S.Set Operation
  90 matches test = S.fromList $ filter (doTest test) [minBound..maxBound]
  91
  92
  93
  94 candidateOpcodes :: [Test] -> Candidates
  95 candidateOpcodes tests = foldl' restrict M.empty tests
  96
  97 restrict :: Candidates -> Test -> Candidates
  98 restrict candidates test = M.insertWith S.intersection opcode possibles candidates
  99     where opcode = _opcode $ _testOperation test
 100           possibles = matches test
 101
 102 findOpcodes :: Candidates -> OpcodeAssignments
 103 findOpcodes candidates = findOpcodes' candidates M.empty
 104
 105 findOpcodes' :: Candidates -> OpcodeAssignments -> OpcodeAssignments
 106 findOpcodes' candidates assignments
 107     | M.null candidates = assignments
 108     | otherwise = findOpcodes' candidates'' assignments'
 109     where singletons = M.map S.findMin $ M.filter (\c -> S.size c == 1) candidates
 110           assignments' = assignments `M.union` singletons
 111           candidates' = candidates `M.difference` singletons
 112           founds = S.fromList $ M.elems assignments
 113           candidates'' = M.map (\cs -> S.filter (\c -> c `S.notMember` founds) cs) candidates'
 114
 115
 116 -- allOps :: S.Set Operation
 117 -- allOps = S.fromList [minBound..maxBound]
 118
 119 -- Parse the input file
 120
 121 type Parser = Parsec Void Text
 122
 123 sc :: Parser ()
 124 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 125
 126 lexeme  = L.lexeme sc
 127 integer = lexeme L.decimal
 128 symb = L.symbol sc
 129
 130 memoryP = (M.fromList . zip [0..] . map fromIntegral) <$> between (symb "[") (symb "]") (integer `sepBy` (symb ","))
 131
 132 beforeP = symb "Before:" *> memoryP
 133 afterP = symb "After:" *> memoryP
 134
 135 rawOpP = opify <$> integer <*> integer <*> integer <*> integer
 136     where opify o a b c = RawOperation { _opcode = fromIntegral o
 137                                        , _a = fromIntegral a
 138                                        , _b = fromIntegral b
 139                                        , _c = fromIntegral c
 140                                        }
 141
 142 testP = testify <$> beforeP <*> rawOpP <*> afterP
 143     where testify b o a = Test {_preMemory = b, _testOperation = o, _postMemory = a}
 144
 145 fileP = (,) <$> (many testP) <*> (many rawOpP)
 146
 147 successfulParse :: Text -> ([Test], [RawOperation])
 148 successfulParse input =
 149         case parse fileP "input" input of
 150                 Left  _error -> ([], []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 151                 Right world -> world