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 emtpyMachine = 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 emtpyMachine instructions
  48
  49 executeInstruction1 machine (Mask mask) = makeMask machine mask
  50 executeInstruction1 machine (Assignment loc value) =
  51   assignValue machine loc value
  52
  53 makeMask machine mask =
  54   machine {mMask0 = maskZeroes mask, mMask1 = maskOnes mask}
  55
  56 assignValue machine loc value =
  57   machine {mMemory = M.insert loc value' mem}
  58   where value' = maskValue machine value
  59         mem = mMemory machine
  60
  61 maskValue machine value =
  62   (value .|. (mMask1 machine)) .&. (mMask0 machine)
  63
  64 maskOnes :: MaskMap -> Int64
  65 maskOnes mask = foldl' setBit zeroBits ones
  66   where ones = M.keys $ M.filter (== One) mask
  67
  68 maskZeroes :: MaskMap -> Int64
  69 maskZeroes mask = complement $ foldl' setBit zeroBits ones
  70   where ones = M.keys $ M.filter (== Zero) mask
  71
  72
  73 executeInstructions2 instructions =
  74   foldl' executeInstruction2 emtpyMachine instructions
  75
  76 executeInstruction2 machine (Mask mask) = machine {mMask = mask}
  77 executeInstruction2 machine (Assignment loc value) = machine {mMemory = mem'}
  78   where locs = map encodeMask $ applyAddressMask (mMask machine) $ decodeMask loc
  79         mem = mMemory machine
  80         mem' = foldl' (\m l -> M.insert l value m) mem locs
  81
  82
  83 encodeMask :: MaskMap -> Int64
  84 encodeMask mask = M.foldrWithKey' setBitValue zeroBits mask
  85   where setBitValue _ Zero n = n
  86         setBitValue i One n = setBit n i
  87
  88 decodeMask :: Int64 -> MaskMap
  89 decodeMask val = M.fromList [ (i, decodeBit $ testBit val i)
  90                             | i <- [0..(finiteBitSize val)]
  91                             ]
  92   where decodeBit True = One
  93         decodeBit False = Zero
  94
  95 applyAddressMask :: MaskMap -> MaskMap -> [MaskMap]
  96 applyAddressMask mask address = M.foldrWithKey' applyBit [address] mask
  97
  98 applyBit :: Int -> MaskValue -> [MaskMap] -> [MaskMap]
  99 applyBit _ Zero ms = ms
 100 applyBit k One  ms = [ M.insert k One m | m <- ms ]
 101 applyBit k Wild ms = [ M.insert k b m | m <- ms, b <- [Zero, One] ]
 102
 103 -- Parse the input file
 104
 105 programP = (maskP <|> assignmentP) `sepBy` endOfLine
 106
 107 maskP = maskify <$> ("mask = " *> (many (digit <|> letter)))
 108 assignmentP = Assignment <$> ("mem[" *> decimal) <* "] = " <*> decimal
 109
 110 maskify :: String -> Instruction
 111 maskify chars = Mask (M.fromList locValues)
 112   where mValues = map readMaskChar chars
 113         locValues = zip [0..] $ reverse mValues
 114
 115 readMaskChar '0' = Zero
 116 readMaskChar '1' = One
 117 readMaskChar 'X' = Wild
 118
 119 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 120 successfulParse input =
 121   case parseOnly programP input of
 122     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 123     Right program -> program