adventofcode-2020/day14/main.hs

#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple attoparsec])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports #-}
{-# OPTIONS_GHC -Wno-unused-matches #-}
{-# LANGUAGE OverloadedStrings #-}

import Debug.Trace (trace)
import Text.Pretty.Simple

import Data.IntMap (IntMap)
import qualified Data.IntMap as I

import Data.Bits ((.&.),(.|.),FiniteBits,bit,complementBit)
import Data.Maybe (catMaybes)

import Numeric (showIntAtBase)
import Data.Char (intToDigit)

import Data.Attoparsec.Text (Parser)
import qualified Data.Attoparsec.Text as P

import Data.Text (Text)
import qualified Data.Text as T

import Data.String (IsString)

import Data.Monoid (Sum(..))

exampleData :: String
exampleData = unlines $
  [ "mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"
  , "mem[8] = 11"
  , "mem[7] = 101"
  , "mem[8] = 0"
  ]

-- The *special* mask is essentially a successive application of:
-- - (.&.) with 0 (all others are 1), in order to force 0 on some bits, or
-- - (.|.) with 1 (all others are 0), in order to force 1 on some bits.
--
data BW a = BWAnd a
          | BWOr a

instance (Show a, Integral a) => Show (BW a)
  where
    show (BWAnd a) = "And " <> showIntAtBase 2 intToDigit a ""
    show (BWOr a) = "Or " <> showIntAtBase 2 intToDigit a ""

mkBW :: String -> Mask
mkBW xs | length xs == 36 = catMaybes $ map go (zip [35,34..] xs)
  where
    go (p,'1') = Just $ BWOr  (bit p)
    go (p,'0') = Just $ BWAnd (complementBit (2 ^ (36 :: Int) - 1) p)
    go (p,'X') = Nothing
    go (_,c)   = trace ("Illegal char: " <> show c) $ undefined
mkBW xs | otherwise = trace "Mask is not 36 bit long" $ undefined

applyMask :: (FiniteBits a) => [BW a] -> a -> a
applyMask xs a = foldl applyBitwise a xs
  where
    applyBitwise v (BWAnd n) = n .&. v
    applyBitwise v (BWOr n) = n .|. v
--

-- We need a datatype to hold our datastructure
--
type Mask = [BW Int]
type Mem = IntMap Int
type Addr = Int

data InitBlock = InitBlock { unMask :: Mask, unInitLines :: [(Addr,Int)] }
  deriving Show

data Program = Program { unMem :: Mem, unInits :: [InitBlock] }
  deriving Show
--

-- Now, let's parse our input uwu
--
-- mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
parseMask :: Parser Mask
parseMask = do
  maskStr <- "mask = " *> P.take 36
  P.endOfLine
  pure $ mkBW $ T.unpack maskStr

-- mem[x] = y
parseInitLines :: Parser (Addr,Int)
parseInitLines = do
  x <- "mem[" *> P.decimal
  y <- "] = " *> P.decimal
  P.endOfLine
  pure $ (x,y)

parseBlock :: Parser InitBlock
parseBlock = do
  mask <- parseMask
  ix <- P.many1 parseInitLines
  pure $ InitBlock mask ix

parseInput :: Parser [InitBlock]
parseInput = do
  bx <- P.many1 parseBlock
  P.endOfInput
  pure bx
--

-- Now we've got everyting for part1, let's solve the thing
--
runInitBlock :: Mem -> InitBlock -> Mem
runInitBlock mem (InitBlock mask inits) =
  I.union res mem
  where
    res = foldMap (\(addr,v) -> I.insert addr (applyMask mask v) I.empty) $ reverse inits

solvePart1 :: Text -> Either String Int
solvePart1 input = do
  bx <- P.parseOnly parseInput input
  pure . getSum . (foldMap Sum) . I.elems $ fullMem bx
  where
    fullMem :: [InitBlock] -> Mem
    fullMem = foldMap (runInitBlock I.empty) . reverse
--

-- Part 2 starts by redefining the meaning for the bits in the mask
--
data BM = BMUnchanged
        | BMForceOne
        | BMFloating

instance Show (BM)
  where
    show BMUnchanged = "⁰"
    show BMForceOne  = "¹"
    show BMFloating  = "^"

type Mask2 = [BM]

data InitBlock2 = InitBlock2 { unMask2 :: Mask2, unInitLines2 :: [(Addr,Int)] }
  deriving Show

data Program2 = Program2 { unMem2 :: Mem, unInits2 :: [InitBlock2] }
  deriving Show

mkBM :: String -> Mask2
mkBM xs | length xs == 36 = map go xs
  where
    go '1' = BMForceOne
    go '0' = BMUnchanged
    go 'X' = BMFloating
    go c   = trace ("Illegal char: " <> show c) $ undefined
mkBM xs | otherwise = trace "Mask is not 36 bit long" $ undefined


-- We have to redo the parsing side, because we were too specific (gotcha!)
--
parseMask2 :: Parser Mask2
parseMask2 = do
  maskStr <- "mask = " *> P.take 36
  P.endOfLine
  pure $ mkBM $ T.unpack maskStr

-- mem[x] = y
parseInitLines2 :: Parser (Addr,Int)
parseInitLines2 = do
  x <- "mem[" *> P.decimal
  y <- "] = " *> P.decimal
  P.endOfLine
  pure $ (x,y)

parseBlock2 :: Parser InitBlock2
parseBlock2 = do
  mask <- parseMask2
  ix <- P.many1 parseInitLines
  pure $ InitBlock2 mask ix

parseInput2 :: Parser [InitBlock2]
parseInput2 = do
  bx <- P.many1 parseBlock2
  P.endOfInput
  pure bx

applyMask2 :: Mask2 -> Int -> [BM]
applyMask2 xs a = map go (zip xs (showIntAtBase (2::Int) intToDigit a ""))
  where
    go (BMForceOne ,_   )  = BMForceOne
    go (BMUnchanged,'1' )  = BMForceOne
    go (BMUnchanged,'0' )  = BMUnchanged
    go (BMFloating ,_   )  = BMFloating
    go (nb,_) = trace "Invalid char in base 2 number" $ undefined

main :: IO ()
main = do
  input <- T.pack <$> readFile "day14/input"
  putStrLn ":: Tests - Part 1"
  pPrint . lines $ exampleData
  print $ mkBW "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"
  print $ applyMask (mkBW "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X") 11
  print $ applyMask (mkBW "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X") 101
  print $ applyMask (mkBW "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X") 0
  pPrint . P.parseOnly parseBlock . T.pack $ exampleData
  -- pPrint . P.parseOnly parseInput $ input
  putStrLn ":: Day 14 - Part 1"
  print . solvePart1 . T.pack $ exampleData
  print . solvePart1 $ input
  putStrLn ":: Tests - Part 2"
  print $ mkBM "000000000000000000000000000000X1001X"
  print $ applyMask2 (mkBM "000000000000000000000000000000X1001X") (42::Int)
  putStrLn ":: Day 14 - Part 2"