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
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 =