85 lines
2.1 KiB
Haskell
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
|