1 files changed, 61 insertions, 0 deletions
diff --git a/2020/14.hs b/2020/14.hs
new file mode 100644
index 0000000..4a4e24d
--- /dev/null
+++ b/2020/14.hs
@@ -0,0 +1,61 @@
+module Main where
+
+import Data.Bits
+import qualified Data.IntMap.Strict as IM
+import Data.IntMap.Strict (IntMap)
+import Data.List (foldl')
+
+import Input
+
+
+subsets :: [a] -> [[a]]
+subsets [] = [[]]
+subsets (x:xs) = subsets xs ++ map (x:) (subsets xs)
+
+data Mask = Mask Int Int
+  deriving (Show)
+
+applyMask :: Mask -> Int -> Int
+applyMask (Mask set reset) n = (n .|. set) .&. complement reset
+
+data Instr = SetMask Mask | Store Int Int
+  deriving (Show)
+
+parseInstr :: String -> Instr
+parseInstr str
+  | str !! 1 == 'a' = SetMask (parseMask (drop 7 str))
+  | [(idx, rest)] <- reads (drop 4 str) = Store idx (read (drop 4 rest))
+  | otherwise = error "Can't parse Instr"
+
+parseMask :: String -> Mask
+parseMask s =
+    let pairs = zip s (map bit [35,34..0])
+        setbits   = sum . map snd . filter ((== '1') . fst) $ pairs
+        resetbits = sum . map snd . filter ((== '0') . fst) $ pairs
+    in Mask setbits resetbits
+
+data State = State { sMask :: Mask, sMem :: IntMap Int }
+  deriving (Show)
+
+exec1 :: State -> Instr -> State
+exec1 (State _ mem) (SetMask mask) = State mask mem
+exec1 (State mask mem) (Store idx val) =
+    State mask (IM.insert idx (applyMask mask val) mem)
+
+expandWithMask :: Mask -> Int -> [Int]
+expandWithMask (Mask ones zeros) n =
+    let floating = (bit 36 - 1) .&. complement (ones .|. zeros)
+        n' = (n .|. ones) .&. complement floating
+        singlebits = [bit i | i <- [0..35], testBit floating i]
+    in map ((n' .|.) . sum) (subsets singlebits)
+
+exec2 :: State -> Instr -> State
+exec2 (State _ mem) (SetMask mask) = State mask mem
+exec2 (State mask mem) (Store idx val) =
+    State mask (foldr (\idx' -> IM.insert idx' val) mem (expandWithMask mask idx))
+
+main :: IO ()
+main = do
+    input <- map parseInstr <$> getInput 14
+    print (sum . IM.elems . sMem $ foldl' exec1 (State (Mask 0 0) mempty) input)
+    print (sum . IM.elems . sMem $ foldl' exec2 (State (Mask 0 0) mempty) input)