Discussion on: Advent of Code 2020 Solution Megathread - Day 14: Docking Data

View post

My Haskell soloution, have to save I've never used bitwise operations in Haskell very much, so had to look all that up. And during that search I realised that foldl', which I defined on day 8, is actually implemented in Data.List, oops.. :-)

data Inst = Mask Integer Integer  -- mask with off bits and mask with on bits
          | Mem Integer Integer
  deriving (Show, Eq)

type Program = [Inst]

parse :: [String] -> Program 
parse = map aux
  where
    aux ('m':'a':'s':'k':' ':'=':' ':mask) = 
        let (off, ""):_ = readInt 2 (const True) (fromEnum . (== '1')) mask
            (on, ""):_ = readInt 2 (const True) (fromEnum . (/= '0')) mask
        in Mask on off
    aux xs = let idx = read (takeWhile isDigit (drop 4 xs)) :: Integer
                 n   = read (drop 2 (dropWhile (/= '=') xs)) :: Integer
             in Mem idx n

task1 :: [Inst] -> Integer
task1 = M.foldr (+) 0 . snd . foldl' (uncurry aux) ((0, 0), M.empty) 
  where 
    aux _         mem (Mask on off) = ((on, off), mem) 
    aux (on, off) mem (Mem idx value) = 
      ((on,off), M.insert idx (value .&. on .|. off) mem)

task2 :: [Inst] -> Integer
task2 = M.foldr (+) 0 . snd . foldl' (uncurry aux) ((0, 0), M.empty)
  where 
    aux _ mem (Mask on off) = ((on, off), mem)
    aux (on,off) mem (Mem idx value) = 
      ((on, off), 
       let idxs = map (xor (idx .|. off)) (bitPowerSet (off `xor` on))
       in foldl' (\mem idx -> M.insert idx value mem) mem idxs)

-- function from reddit
bitPowerSet:: Integer -> [Integer]
bitPowerSet bits = chooseBits <$> [0..1 `shiftL` popCount bits - 1] 
  where
        chooseBits i = fst $ foldl' (f i) (0, bits) [0..popCount bits - 1]
        f i (x, k) j =
            (x `xor` bits .&. (k `xor` (k - i `shiftR` j .&. 1)), k .&. (k - 1))

main :: IO ()
main = do
  insts <- readFile "day14_input" <&> lines <&> parse
  print (task1 insts)
  print (task2 insts)