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