Day 11

src/Days/D11.hs

@@ -0,0 +1,180 @@
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TemplateHaskell #-}
+
+module Days.D11 where
+
+import Common
+import Control.Monad.Combinators.Expr (Operator (..), makeExprParser)
+import Parse
+import Text.Megaparsec (between, choice, sepBy, some)
+import Text.Megaparsec.Char (char, letterChar, newline, space)
+import Text.Megaparsec.Char.Lexer qualified as L
+
+import Control.Monad (foldM_)
+import Control.Monad.State (State, execState, gets, modify)
+import Data.Map (Map)
+import Data.Map qualified as M
+import Data.Maybe (fromJust)
+
+import Control.Lens (makeLenses, over, (<&>))
+import Data.List (sortBy)
+import Data.Ord (Down (..), comparing)
+
+data Op = Mult Op Op | Add Op Op | Assign Op Op | Var String | Const Int
+  deriving (Show)
+
+data Monkey = Monkey
+  { op :: Op
+  , tester :: Int
+  , successTarget :: Int
+  , failTarget :: Int
+  }
+  deriving (Show)
+
+divides :: (Integral a) => a -> a -> Bool
+divides a b = b `mod` a == 0
+
+parser :: Parser [(Int, [Int], Monkey)]
+parser = someLines (monkey <* newline)
+ where
+  symbol = L.symbol space
+  lexeme = L.lexeme space
+
+  monkey :: Parser (Int, [Int], Monkey)
+  monkey = do
+    mId <- symbol "Monkey" *> number <* char ':' <* newline
+    items <- space *> symbol "Starting items:" *> sepBy (fromIntegral <$> number) (symbol ",")
+    op_ <- space *> symbol "Operation:" *> operation
+    test <- space *> symbol "Test:" *> symbol "divisible by " *> number
+    ifTrue <- space *> symbol "If true:" *> symbol "throw to monkey" *> number
+    ifFalse <- space *> symbol "If false:" *> symbol "throw to monkey" *> number
+    return
+      ( mId
+      , items
+      , Monkey{op = op_, tester = fromIntegral test, successTarget = ifTrue, failTarget = ifFalse}
+      )
+
+  operatorTable :: [[Operator Parser Op]]
+  operatorTable =
+    [ [InfixL (Mult <$ symbol "*")]
+    , [InfixL (Add <$ symbol "+")]
+    , [InfixN (Assign <$ symbol "=")]
+    ]
+
+  operation :: Parser Op
+  operation =
+    makeExprParser
+      ( choice
+          [ Var <$> lexeme (some letterChar)
+          , Const . fromIntegral <$> lexeme number
+          , between (symbol "(") (symbol ")") operation
+          ]
+      )
+      operatorTable
+
+type Context = Map String Int
+
+eval :: Op -> Context -> (Int, Context)
+eval (Mult a b) context =
+  let (av, ca) = eval a context
+      (bv, cb) = eval b ca
+   in (av * bv, cb)
+eval (Add a b) context =
+  let (av, ca) = eval a context
+      (bv, cb) = eval b ca
+   in (av + bv, cb)
+eval (Assign a b) context =
+  case a of
+    (Var v) ->
+      let (bv, cb) = eval b context
+       in (bv, M.insert v bv cb)
+    _ -> error "unexpected expr!"
+eval (Var v) context = (fromJust $ M.lookup v context, context)
+eval (Const c) context = (c, context)
+
+buildMaps :: [(Int, [Int], Monkey)] -> (Map Int [Int], Map Int Monkey)
+buildMaps l =
+  let (indices, values, monkeys) = unzip3 l
+   in (M.fromList (zip indices values), M.fromList (zip indices monkeys))
+
+data ProcessState = ProcessState
+  { _queues :: Map Int [Int]
+  , _counters :: Map Int Int
+  }
+  deriving (Show)
+
+$(makeLenses ''ProcessState)
+
+simulate :: (Int -> Int) -> Int -> Map Int [Int] -> Map Int Monkey -> Integer
+simulate dampen rounds initial monkeyMap =
+  let initialState = ProcessState{_queues = initial, _counters = M.empty}
+   in business $ execState (foldM_ (const monkeyDo) () turns) initialState
+ where
+  numMonkeys :: Int
+  numMonkeys = M.size monkeyMap
+
+  turns :: [Int]
+  turns = concat $ replicate rounds [0 .. (numMonkeys - 1)]
+
+  business :: ProcessState -> Integer
+  business =
+    ( \case
+        [a, b] -> toInteger a * toInteger b
+        _ -> error "Require 2 or more monkeys"
+    )
+      . take 2
+      . sortBy (comparing Down)
+      . M.elems
+      . _counters
+
+  monkeyDo :: Int -> State ProcessState ()
+  monkeyDo i =
+    let monkey = fromJust $ M.lookup i monkeyMap
+     in unsafeGetQueue i
+          >>= \q ->
+            let
+              l = length q
+              p = processQueue monkey q
+             in
+              foldM_ (const $ uncurry updateQueues) () p
+                >> clearQueue i
+                >> updateStats i l
+
+  unsafeGetQueue :: Int -> State ProcessState [Int]
+  unsafeGetQueue i = gets (fromJust . M.lookup i . _queues)
+
+  updateStats :: Int -> Int -> State ProcessState ()
+  updateStats monkeyId numSeen =
+    modify
+      . over counters
+      $ M.alter (Just . maybe numSeen (+ numSeen)) monkeyId
+
+  updateQueues :: Int -> Int -> State ProcessState ()
+  updateQueues k v = modify (over queues (M.adjust (v :) k))
+
+  clearQueue :: Int -> State ProcessState ()
+  clearQueue i = modify (over queues (M.insert i []))
+
+  processQueue :: Monkey -> [Int] -> [(Int, Int)]
+  processQueue m = map (processSingle m)
+
+  processSingle :: Monkey -> Int -> (Int, Int)
+  processSingle monkey value =
+    let evalResult = fst $ eval (op monkey) (M.fromList [("old", value)])
+        newValue = dampen evalResult
+        testResult = tester monkey `divides` newValue
+        target = (if testResult then successTarget else failTarget) monkey
+     in (target, newValue)
+
+part1 :: Map Int [Int] -> Map Int Monkey -> Integer
+part1 = simulate (`div` 3) 20
+
+part2 :: Map Int [Int] -> Map Int Monkey -> Integer
+part2 initialQueues monkeyMap =
+  let modulus = product . map tester $ M.elems monkeyMap
+   in simulate (`mod` modulus) 10000 initialQueues monkeyMap
+
+day :: Day
+day = parsecDay (parser <&> buildMaps) (definitive . uncurry part1, definitive . uncurry part2)