{-# LANGUAGE LambdaCase #-} import Data.List (tails, foldl') import qualified Data.Map.Strict as Map import Data.Map.Strict (Map) data Regex = L String -- ^ literal | A [Regex] -- ^ alternation (or) | S [Regex] -- ^ sequence (concatenation) | M Regex -- ^ many (star) | Set [Char] -- ^ character class | Grp Int Regex -- ^ capture group deriving (Show) instance Semigroup Regex where a <> S bs = S (a : bs) a <> b = S [a, b] instance Monoid Regex where mempty = L "" some :: Regex -> Regex some r = r <> M r -- (full match, groups, rest) match :: Regex -> String -> [(String, Map Int String, String)] match (L r) s | take (length r) s == r = pure (r, mempty, drop (length r) s) | otherwise = [] match (A rs) s = concatMap (`match` s) rs match (S []) s = [("", mempty, s)] match (S (r:rs)) s = do (f1, m1, rest1) <- match r s (f2, m2, rest2) <- match (S rs) rest1 pure (f1 ++ f2, m1 <> m2, rest2) match (M r) s = case match r s of [] -> [("", mempty, s)] ms -> do (f1, m1, rest1) <- ms case match (M r) rest1 of [] -> pure (f1, m1, rest1) ms' -> do (f2, m2, rest2) <- ms' pure (f1 ++ f2, m1 <> m2, rest2) match (Set chs) s | c : s' <- s, c `elem` chs = [([c], mempty, s')] | otherwise = [] match (Grp i r) s = do (f, m, rest) <- match r s pure (f, Map.insert i f m, rest) matchall :: Regex -> String -> [(String, Map Int String, String)] matchall r s = concatMap (match r) (tails s) main :: IO () main = do let num_reg = some (Set ['0'..'9']) mul_reg = L "mul(" <> Grp 1 num_reg <> L "," <> Grp 2 num_reg <> L ")" input <- getContents print $ sum [read @Int (m Map.! 1) * read (m Map.! 2) | (_, m, _) <- matchall mul_reg input] let instr_reg = A [mul_reg, L "do()", L "don't()"] let instrs = [case (,) <$> Map.lookup 1 m <*> Map.lookup 2 m of Just (a, b) -> Right (read @Int a * read b) Nothing -> Left (s == "do()") | (s, m, _) <- matchall instr_reg input] print $ fst $ foldl' (\(acc, en) -> \case Right v | en -> (acc + v, en) | otherwise -> (acc, en) Left en' -> (acc, en')) (0, True) instrs