[hs] Solve 2020_14 part 2

hs/src/Aoc/Y2020/D14.hs

@@ -32,15 +32,25 @@ parser = manyLines (pMask <|> pSet)
 digits :: Int -> [Bool]
 digits = map odd . iterate (`div` 2)
 
-mask :: [Maybe Bool] -> Int -> [Bool]
-mask m n = zipWith maskBit m $ digits n
+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 [Bool]
+  , mMem  :: Map.Map Int Int
   } deriving (Show)
 
 newMem :: Mem
@@ -50,22 +60,26 @@ 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 masked $ mMem mem}
-  where
-    masked = mask (mMask mem) val
+setMem addr val mem = mem{mMem = Map.insert addr val $ mMem mem}
 
-doInstr :: Instr -> Mem -> Mem
-doInstr (Mask m)       = setMask m
-doInstr (Set addr val) = setMem addr val
+doInstr1 :: Instr -> Mem -> Mem
+doInstr1 (Mask m)       mem = setMask m mem
+doInstr1 (Set addr val) mem = setMem addr (mask1 (mMask mem) val) mem
 
-value :: [Bool] -> Int
-value = sum . map fst . filter snd . zip (iterate (*2) 1)
+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 mem = foldl' (flip doInstr) newMem instrs
-  print $ sum $ map value $ Map.elems $ mMem mem
+  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