Day 14 part 2
[advent-of-code-20.git] / advent14 / src / advent14.hs
1 -- import Debug.Trace
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.Attoparsec.Text hiding (take)
8 -- import Data.Attoparsec.Combinator
9 import Control.Applicative
10
11 import Data.Int
12 import Data.Bits
13 import Data.List
14 import qualified Data.Map.Strict as M
15
16
17 data MaskValue = Zero | One | Wild deriving (Show, Eq)
18 type MaskMap = M.Map Int MaskValue
19 data Instruction = Mask MaskMap | Assignment Int64 Int64
20 deriving (Show, Eq)
21
22 type Memory = M.Map Int64 Int64
23 data Machine = Machine { mMemory :: Memory
24 , mMask :: MaskMap
25 , mMask0 :: Int64
26 , mMask1 :: Int64
27 } deriving (Show, Eq)
28
29 emtpyMachine = Machine M.empty M.empty (complement 0) 0
30
31
32 main :: IO ()
33 main =
34 do text <- TIO.readFile "data/advent14.txt"
35 let program = successfulParse text
36 print $ part1 program
37 print $ part2 program
38
39
40 part1 program = sum $ M.elems $ mMemory finalMachine
41 where finalMachine = executeInstructions1 program
42
43 part2 program = sum $ M.elems $ mMemory finalMachine
44 where finalMachine = executeInstructions2 program
45
46 executeInstructions1 instructions =
47 foldl' executeInstruction1 emtpyMachine instructions
48
49 executeInstruction1 machine (Mask mask) = makeMask machine mask
50 executeInstruction1 machine (Assignment loc value) =
51 assignValue machine loc value
52
53 makeMask machine mask =
54 machine {mMask0 = maskZeroes mask, mMask1 = maskOnes mask}
55
56 assignValue machine loc value =
57 machine {mMemory = M.insert loc value' mem}
58 where value' = maskValue machine value
59 mem = mMemory machine
60
61 maskValue machine value =
62 (value .|. (mMask1 machine)) .&. (mMask0 machine)
63
64 maskOnes :: MaskMap -> Int64
65 maskOnes mask = foldl' setBit zeroBits ones
66 where ones = M.keys $ M.filter (== One) mask
67
68 maskZeroes :: MaskMap -> Int64
69 maskZeroes mask = complement $ foldl' setBit zeroBits ones
70 where ones = M.keys $ M.filter (== Zero) mask
71
72
73 executeInstructions2 instructions =
74 foldl' executeInstruction2 emtpyMachine instructions
75
76 executeInstruction2 machine (Mask mask) = machine {mMask = mask}
77 executeInstruction2 machine (Assignment loc value) = machine {mMemory = mem'}
78 where locs = map encodeMask $ applyAddressMask (mMask machine) $ decodeMask loc
79 mem = mMemory machine
80 mem' = foldl' (\m l -> M.insert l value m) mem locs
81
82
83 encodeMask :: MaskMap -> Int64
84 encodeMask mask = M.foldrWithKey' setBitValue zeroBits mask
85 where setBitValue _ Zero n = n
86 setBitValue i One n = setBit n i
87
88 decodeMask :: Int64 -> MaskMap
89 decodeMask val = M.fromList [ (i, decodeBit $ testBit val i)
90 | i <- [0..(finiteBitSize val)]
91 ]
92 where decodeBit True = One
93 decodeBit False = Zero
94
95 applyAddressMask :: MaskMap -> MaskMap -> [MaskMap]
96 applyAddressMask mask address = M.foldrWithKey' applyBit [address] mask
97
98 applyBit :: Int -> MaskValue -> [MaskMap] -> [MaskMap]
99 applyBit _ Zero ms = ms
100 applyBit k One ms = [ M.insert k One m | m <- ms ]
101 applyBit k Wild ms = [ M.insert k b m | m <- ms, b <- [Zero, One] ]
102
103 -- Parse the input file
104
105 programP = (maskP <|> assignmentP) `sepBy` endOfLine
106
107 maskP = maskify <$> ("mask = " *> (many (digit <|> letter)))
108 assignmentP = Assignment <$> ("mem[" *> decimal) <* "] = " <*> decimal
109
110 maskify :: String -> Instruction
111 maskify chars = Mask (M.fromList locValues)
112 where mValues = map readMaskChar chars
113 locValues = zip [0..] $ reverse mValues
114
115 readMaskChar '0' = Zero
116 readMaskChar '1' = One
117 readMaskChar 'X' = Wild
118
119 -- successfulParse :: Text -> (Integer, [Maybe Integer])
120 successfulParse input =
121 case parseOnly programP input of
122 Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
123 Right program -> program