X-Git-Url: https://git.njae.me.uk/?p=advent-of-code-20.git;a=blobdiff_plain;f=advent14%2Fsrc%2Fadvent14.hs;h=b6343a4b54be52281c05e235a635b670a0e66a2c;hp=87aaa3de8d6da06f518faa5ac18fef6da987572a;hb=refs%2Fheads%2Fmain;hpb=0438ea0328365f1060941424156cd1b9bc0fae29

diff --git a/advent14/src/advent14.hs b/advent14/src/advent14.hs
index 87aaa3d..b6343a4 100644
--- a/advent14/src/advent14.hs
+++ b/advent14/src/advent14.hs
@@ -10,64 +10,97 @@ import Control.Applicative
 
 import Data.Int
 import Data.Bits
-import Data.Char
 import Data.List
-
 import qualified Data.Map.Strict as M
-import Data.Map.Strict ((!))
 
 
-type MaskMap = M.Map Int Int
+data MaskValue = Zero | One | Wild deriving (Show, Eq)
+type MaskMap = M.Map Int MaskValue
 data Instruction = Mask MaskMap | Assignment Int64 Int64
   deriving (Show, Eq)
 
 type Memory = M.Map Int64 Int64
 data Machine = Machine { mMemory :: Memory
+                       , mMask :: MaskMap
                        , mMask0 :: Int64
                        , mMask1 :: Int64
                        } deriving (Show, Eq)
 
-emtpyMachine = Machine M.empty (complement 0) 0
+emptyMachine = Machine M.empty M.empty (complement 0) 0
 
 
 main :: IO ()
 main = 
   do  text <- TIO.readFile "data/advent14.txt"
       let program = successfulParse text
-      print $ take 6 program
       print $ part1 program
+      print $ part2 program
 
 
 part1 program = sum $ M.elems $ mMemory finalMachine
-  where finalMachine = executeInstructions program
+  where finalMachine = executeInstructions1 program
 
-executeInstructions instructions = 
-  foldl' executeInstruction emtpyMachine instructions
+part2 program = sum $ M.elems $ mMemory finalMachine
+  where finalMachine = executeInstructions2 program
 
-executeInstruction machine (Mask mask) = makeMask machine mask
-executeInstruction machine (Assignment loc value) = assignValue machine loc value
+executeInstructions1 instructions = 
+  foldl' executeInstruction1 emptyMachine instructions
 
+executeInstruction1 :: Machine -> Instruction -> Machine
+executeInstruction1 machine (Mask mask) = makeMask machine mask
+executeInstruction1 machine (Assignment loc value) = 
+  assignValue machine loc value
 
 makeMask machine mask = 
   machine {mMask0 = maskZeroes mask, mMask1 = maskOnes mask}
 
-maskValue machine value = 
-  (value .|. (mMask1 machine)) .&. (mMask0 machine)
-
 assignValue machine loc value = 
   machine {mMemory = M.insert loc value' mem}
   where value' = maskValue machine value
         mem = mMemory machine
 
+maskValue machine value = 
+  (value .|. (mMask1 machine)) .&. (mMask0 machine)
 
 maskOnes :: MaskMap -> Int64
 maskOnes mask = foldl' setBit zeroBits ones
-  where ones = M.keys $ M.filter (== 1) mask
+  where ones = M.keys $ M.filter (== One) mask
 
 maskZeroes :: MaskMap -> Int64
-maskZeroes mask = complement $ foldl' setBit zeroBits ones
-  where ones = M.keys $ M.filter (== 0) mask
+maskZeroes mask = foldl' clearBit (complement zeroBits) zeroes
+  where zeroes = M.keys $ M.filter (== Zero) mask
+
+
+executeInstructions2 instructions = 
+  foldl' executeInstruction2 emptyMachine instructions
+
+executeInstruction2 :: Machine -> Instruction -> Machine
+executeInstruction2 machine (Mask mask) = machine {mMask = mask}
+executeInstruction2 machine (Assignment loc value) = machine {mMemory = mem'}
+  where locs = map encodeMask $ applyAddressMask (mMask machine) $ decodeMask loc
+        mem = mMemory machine
+        mem' = foldl' (\m l -> M.insert l value m) mem locs
+
+
+encodeMask :: MaskMap -> Int64
+encodeMask mask = M.foldrWithKey' setBitValue zeroBits mask
+  where setBitValue _ Zero n = n
+        setBitValue i One n = setBit n i
+
+decodeMask :: Int64 -> MaskMap
+decodeMask val = M.fromList [ (i, decodeBit $ testBit val i) 
+                            | i <- [0..(finiteBitSize val)] 
+                            ]
+  where decodeBit True = One
+        decodeBit False = Zero
+
+applyAddressMask :: MaskMap -> MaskMap -> [MaskMap]
+applyAddressMask mask address = M.foldrWithKey' applyBit [address] mask
 
+applyBit :: Int -> MaskValue -> [MaskMap] -> [MaskMap]
+applyBit _ Zero ms = ms
+applyBit k One  ms = [ M.insert k One m | m <- ms ]
+applyBit k Wild ms = [ M.insert k b m | m <- ms, b <- [Zero, One] ]
 
 -- Parse the input file
 
@@ -77,11 +110,13 @@ maskP = maskify <$> ("mask = " *> (many (digit <|> letter)))
 assignmentP = Assignment <$> ("mem[" *> decimal) <* "] = " <*> decimal
 
 maskify :: String -> Instruction
-maskify chars = Mask (M.fromList locNums)
-  where locChars = zip [0..] $ reverse chars
-        locDigits = filter (isDigit . snd) locChars
-        locNums = map (\(i, n) -> (i, read @Int [n])) locDigits
+maskify chars = Mask (M.fromList locValues)
+  where mValues = map readMaskChar chars
+        locValues = zip [0..] $ reverse mValues
 
+readMaskChar '0' = Zero
+readMaskChar '1' = One
+readMaskChar 'X' = Wild
 
 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 successfulParse input =