advent-of-code/hs/src/Aoc/Y2020/D14.hs
2020-12-14 11:02:07 +00:00

85 lines
2.1 KiB
Haskell

{-# LANGUAGE OverloadedStrings #-}
module Aoc.Y2020.D14
( day
) where
import Control.Monad
import Data.List
import qualified Data.Map as Map
import Aoc.Day
import Aoc.Parse
data Instr
= Mask [Maybe Bool]
| Set Int Int
parser :: Parser [Instr]
parser = manyLines (pMask <|> pSet)
where
pMask = do
void $ string "mask = "
bits <- sequenceA $ replicate 36 $ (Nothing <$ char 'X') <|> (Just False <$ char '0') <|> (Just True <$ char '1')
pure $ Mask $ reverse bits -- Most significant first
pSet = do
void $ string "mem["
addr <- decimal
void $ string "] = "
Set addr <$> decimal
digits :: Int -> [Bool]
digits = map odd . iterate (`div` 2)
undigits :: [Bool] -> Int
undigits = sum . map fst . filter snd . zip (iterate (*2) 1)
mask1 :: [Maybe Bool] -> Int -> Int
mask1 m = undigits . zipWith maskBit m . digits
where
maskBit Nothing = id
maskBit (Just a) = const a
mask2 :: [Maybe Bool] -> Int -> [Int]
mask2 m = map undigits . sequenceA . zipWith maskBit m . digits
where
maskBit (Just False) b = [b]
maskBit (Just True) _ = [True]
maskBit Nothing _ = [False, True]
data Mem = Mem
{ mMask :: [Maybe Bool]
, mMem :: Map.Map Int Int
} deriving (Show)
newMem :: Mem
newMem = Mem{mMask = replicate 36 Nothing, mMem = Map.empty}
setMask :: [Maybe Bool] -> Mem -> Mem
setMask m mem = mem{mMask = m}
setMem :: Int -> Int -> Mem -> Mem
setMem addr val mem = mem{mMem = Map.insert addr val $ mMem mem}
doInstr1 :: Instr -> Mem -> Mem
doInstr1 (Mask m) mem = setMask m mem
doInstr1 (Set addr val) mem = setMem addr (mask1 (mMask mem) val) mem
doInstr2 :: Instr -> Mem -> Mem
doInstr2 (Mask m) mem = setMask m mem
doInstr2 (Set addr val) mem = foldl' (\m a -> setMem a val m) mem $ mask2 (mMask mem) addr
solver :: [Instr] -> IO ()
solver instrs = do
putStrLn ">> Part 1"
let mem1 = foldl' (flip doInstr1) newMem instrs
print $ sum $ mMem mem1
putStrLn ""
putStrLn ">> Part 2"
let mem2 = foldl' (flip doInstr2) newMem instrs
print $ sum $ mMem mem2
day :: Day
day = dayParse parser solver