Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- {-# language RecordWildCards #-}
- import Control.Arrow ((&&&))
- import Control.Foldl (Fold(Fold), fold)
- import Data.Bits (complement, zeroBits, shift, setBit, clearBit, (.|.), (.&.))
- import Data.Char (isDigit)
- import Data.IntMap (IntMap)
- import qualified Data.IntMap as IM
- import Data.Maybe (listToMaybe)
- data Mask = MkMask { clearBits :: Int, setBits :: Int } deriving (Eq, Show)
- data Write = MkWrite { location :: Int, value :: Int } deriving (Eq, Show)
- data Instruction = WriteInstr Write | MaskInstr Mask deriving (Eq, Show)
- readMay :: Read a => String -> Maybe a
- readMay = fmap fst . listToMaybe . reads
- parseMask :: String -> Mask
- parseMask s = MkMask{..}
- where
- (clearBits, setBits) = fold ((,) <$> clearFold <*> setFold) s
- clearFold = Fold updClear (complement zeroBits) id
- setFold = Fold updSet zeroBits id
- updClear m c = shiftIn ('0' /= c) m
- updSet m c = shiftIn ('1' == c) m
- shiftIn True = flip setBit 0 . flip shift 1
- shiftIn False = flip shift 1
- parseLine :: String -> Maybe Instruction
- parseLine ('m':'a':_:_:_:_:_:maskStr) = pure . MaskInstr $ parseMask maskStr
- parseLine ('m':'e':_:_:memStr) = WriteInstr <$> do
- location <- readMay locStr
- value <- readMay valStr
- return MkWrite{..}
- where
- (locStr, rest) = span isDigit memStr
- valStr = drop 4 rest
- parseLine _ = Nothing
- common :: String -> [String]
- common = lines
- data ProgramState = MkProgState { mask :: Mask, mem :: IntMap Int } deriving (Eq, Show)
- initProgState :: ProgramState
- initProgState =
- MkProgState
- { mask = MkMask { clearBits = complement zeroBits, setBits = zeroBits }
- , mem = IM.empty
- }
- setMask :: Mask -> ProgramState -> ProgramState
- setMask m s = s{ mask = m }
- writeMem :: Int -> Int -> IntMap Int -> IntMap Int
- writeMem loc 0 = IM.update (const Nothing) loc
- writeMem loc val = IM.insert loc val
- write :: Write -> ProgramState -> ProgramState
- write MkWrite{..} s =
- s{ mem = writeMem location (setBits curMask .|. (clearBits curMask .&. value)) (mem s) }
- where curMask = mask s
- applyInstruction :: Instruction -> ProgramState -> ProgramState
- applyInstruction (WriteInstr w) = write w
- applyInstruction (MaskInstr m) = setMask m
- runProgram :: [Instruction] -> ProgramState
- runProgram = fold (Fold (flip applyInstruction) initProgState id)
- part1 :: [String] -> Maybe Int
- part1 = fmap (sum . mem . runProgram) . traverse parseLine
- data MaskBit = Unchanged | Overwritten | Floating deriving (Eq, Show)
- type Mask2 = [MaskBit]
- data Instruction2 = WriteInstr2 Write | MaskInstr2 Mask2 deriving (Eq, Show)
- parseMaskBit :: Char -> Maybe MaskBit
- parseMaskBit '0' = pure Unchanged
- parseMaskBit '1' = pure Overwritten
- parseMaskBit 'X' = pure Floating
- parseMaskBit _ = Nothing
- parseMask2 :: String -> Maybe Mask2
- parseMask2 = traverse parseMaskBit
- parseLine2 :: String -> Maybe Instruction2
- parseLine2 ('m':'a':_:_:_:_:_:maskStr) = MaskInstr2 <$> parseMask2 maskStr
- parseLine2 ('m':'e':_:_:memStr) = WriteInstr2 <$> do
- location <- readMay locStr
- value <- readMay valStr
- return MkWrite{..}
- where
- (locStr, rest) = span isDigit memStr
- valStr = drop 4 rest
- parseLine2 _ = Nothing
- data ProgramState2 = MkProgState2 { mask2 :: Mask2, mem2 :: IntMap Int } deriving (Eq, Show)
- initProgState2 :: ProgramState2
- initProgState2 =
- MkProgState2
- { mask2 = []
- , mem2 = IM.empty
- }
- decodeMem :: Mask2 -> Int -> [Int]
- decodeMem mbs i = snd $ foldr upd (0,[i]) mbs
- where
- upd Unchanged (n, is) = (succ n, is)
- upd Overwritten (n, is) = (succ n, map (flip setBit n) is)
- upd Floating (n, is) = (succ n, map (flip setBit n) is ++ map (flip clearBit n) is)
- setMask2 :: Mask2 -> ProgramState2 -> ProgramState2
- setMask2 m s = s{ mask2 = m }
- write2 :: Write -> ProgramState2 -> ProgramState2
- write2 MkWrite{..} s =
- s{ mem2 = foldr (\l -> writeMem l value) (mem2 s) $ decodeMem (mask2 s) location }
- applyInstruction2 :: Instruction2 -> ProgramState2 -> ProgramState2
- applyInstruction2 (WriteInstr2 w) = write2 w
- applyInstruction2 (MaskInstr2 m) = setMask2 m
- runProgram2 :: [Instruction2] -> ProgramState2
- runProgram2 = fold (Fold (flip applyInstruction2) initProgState2 id)
- part2 :: [String] -> Maybe Int
- part2 = fmap (sum . mem2 . runProgram2) . traverse parseLine2
- interactive :: Show a => (String -> a) -> IO ()
- interactive f = getContents >>= print . f
- main :: IO ()
- main = interactive ((part1 &&& part2) . common)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement