Tweaking for clarity
[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 emptyMachine = 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 emptyMachine instructions
48
49 executeInstruction1 :: Machine -> Instruction -> Machine
50 executeInstruction1 machine (Mask mask) = makeMask machine mask
51 executeInstruction1 machine (Assignment loc value) =
52 assignValue machine loc value
53
54 makeMask machine mask =
55 machine {mMask0 = maskZeroes mask, mMask1 = maskOnes mask}
56
57 assignValue machine loc value =
58 machine {mMemory = M.insert loc value' mem}
59 where value' = maskValue machine value
60 mem = mMemory machine
61
62 maskValue machine value =
63 (value .|. (mMask1 machine)) .&. (mMask0 machine)
64
65 maskOnes :: MaskMap -> Int64
66 maskOnes mask = foldl' setBit zeroBits ones
67 where ones = M.keys $ M.filter (== One) mask
68
69 maskZeroes :: MaskMap -> Int64
70 maskZeroes mask = foldl' clearBit (complement zeroBits) zeroes
71 where zeroes = M.keys $ M.filter (== Zero) mask
72
73
74 executeInstructions2 instructions =
75 foldl' executeInstruction2 emptyMachine instructions
76
77 executeInstruction2 :: Machine -> Instruction -> Machine
78 executeInstruction2 machine (Mask mask) = machine {mMask = mask}
79 executeInstruction2 machine (Assignment loc value) = machine {mMemory = mem'}
80 where locs = map encodeMask $ applyAddressMask (mMask machine) $ decodeMask loc
81 mem = mMemory machine
82 mem' = foldl' (\m l -> M.insert l value m) mem locs
83
84
85 encodeMask :: MaskMap -> Int64
86 encodeMask mask = M.foldrWithKey' setBitValue zeroBits mask
87 where setBitValue _ Zero n = n
88 setBitValue i One n = setBit n i
89
90 decodeMask :: Int64 -> MaskMap
91 decodeMask val = M.fromList [ (i, decodeBit $ testBit val i)
92 | i <- [0..(finiteBitSize val)]
93 ]
94 where decodeBit True = One
95 decodeBit False = Zero
96
97 applyAddressMask :: MaskMap -> MaskMap -> [MaskMap]
98 applyAddressMask mask address = M.foldrWithKey' applyBit [address] mask
99
100 applyBit :: Int -> MaskValue -> [MaskMap] -> [MaskMap]
101 applyBit _ Zero ms = ms
102 applyBit k One ms = [ M.insert k One m | m <- ms ]
103 applyBit k Wild ms = [ M.insert k b m | m <- ms, b <- [Zero, One] ]
104
105 -- Parse the input file
106
107 programP = (maskP <|> assignmentP) `sepBy` endOfLine
108
109 maskP = maskify <$> ("mask = " *> (many (digit <|> letter)))
110 assignmentP = Assignment <$> ("mem[" *> decimal) <* "] = " <*> decimal
111
112 maskify :: String -> Instruction
113 maskify chars = Mask (M.fromList locValues)
114 where mValues = map readMaskChar chars
115 locValues = zip [0..] $ reverse mValues
116
117 readMaskChar '0' = Zero
118 readMaskChar '1' = One
119 readMaskChar 'X' = Wild
120
121 -- successfulParse :: Text -> (Integer, [Maybe Integer])
122 successfulParse input =
123 case parseOnly programP input of
124 Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
125 Right program -> program