Advent of Code 2022 Day 11

advent2022-day11.hs

import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L

type Parser a = Parsec Void String a

unsafeParse :: Parser a -> String -> a
unsafeParse p s = fromJust <| parseMaybe p s

sc :: Parser ()
sc = L.space space1 empty empty

integer :: Parser Int
integer = L.signed sc L.decimal

data Operation = Add Int | Multiply Int deriving (Show, Eq)

data Monkey = Monkey
  { inspectItems :: Int
  , startingItems :: [Int]
  , operation :: Operation
  , testDivisibleBy :: Int
  , throwTo :: (Int, Int)
  }
  deriving (Show, Eq)

old :: Int
old = -1

solve :: Int
solve =
  allMonkeys
    |> repeats 10000 (monkeysInspectItems 0)
    |> fmap inspectItems
    |> sort
    |> reverse
    |> (\(a : b : _) -> a * b)
 where
  allMonkeys = unsafeParse (sepBy1 pMonkey eol) input

  repeats :: Int -> (a -> a) -> a -> a
  repeats 0 _ a = a
  repeats n f a =
    repeats (n - 1) f (f a)

  monkeysInspectItems :: Int -> [Monkey] -> [Monkey]
  monkeysInspectItems i monkeys =
    if i == length monkeys
      then monkeys
      else monkeysInspectItems (i + 1) (monkeyInspectItems i monkeys)

  monkeyInspectItems :: Int -> [Monkey] -> [Monkey]
  monkeyInspectItems from monkeys =
    let monkey = monkeys !! from
     in if length monkey.startingItems == 0
          then monkeys
          else
            let to = (if inspectItem monkey then fst else snd) monkey.throwTo
             in monkeyInspectItems from (throwItem from to monkeys)

  inspectItem :: Monkey -> Bool
  inspectItem monkey =
    monkey
      |> worryLevel
      |> (`isDivisible` monkey.testDivisibleBy)

  isDivisible :: Int -> Int -> Bool
  isDivisible a b =
    a `rem` b == 0

  allDivisibleBy :: Int
  allDivisibleBy = product <| fmap testDivisibleBy allMonkeys

  worryLevel :: Monkey -> Int
  worryLevel monkey =
    let (item : _) = monkey.startingItems
     in case monkey.operation of
          -- avoid overflow by dividing the product of all monkeys' testDivisibleBy
          Add n -> (item + n) `mod` allDivisibleBy
          Multiply (-1) -> (item * item) `mod` allDivisibleBy
          Multiply n -> (item * n) `mod` allDivisibleBy

  throwItem :: Int -> Int -> [Monkey] -> [Monkey]
  throwItem from to monkeys =
    let f = monkeys !! from
        t = monkeys !! to
        _ : rest = f.startingItems
        f' = f{startingItems = rest, inspectItems = f.inspectItems + 1}
        t' = t{startingItems = t.startingItems <> [worryLevel f]}
     in monkeys
          |> setAt from f'
          |> setAt to t'

  pId :: Parser ()
  pId = do
    _ <- string "Monkey "
    _ <- integer
    _ <- string ":"
    return ()

  pStartingItems :: Parser [Int]
  pStartingItems = do
    _ <- string "Starting items: "
    sepBy1 integer (string ", ")

  pOp :: Parser (Int -> Operation)
  pOp = do
    c <- printChar
    return
      <| if c == '+'
        then Add
        else Multiply

  pOperation :: Parser Operation
  pOperation = do
    _ <- string "Operation: new = old "
    op <- pOp
    _ <- space1
    a <- try (old <$ string "old") <|> integer
    return <| op a

  pTest :: Parser Int
  pTest = do
    _ <- string "Test: "
    _ <- string "divisible by "
    integer

  pThrowTo :: Parser (Int, Int)
  pThrowTo = do
    _ <- string "If true: throw to monkey "
    a <- integer
    _ <- eol
    _ <- string "If false: throw to monkey "
    b <- integer
    return <| (a, b)

  pMonkey :: Parser Monkey
  pMonkey = do
    _ <- pId
    _ <- eol
    items <- pStartingItems
    _ <- eol
    operation <- pOperation
    _ <- eol
    test <- pTest
    _ <- eol
    throwTo <- pThrowTo
    _ <- eol
    return <| Monkey 0 items operation test throwTo

input :: String
input =
  "Monkey 0:\n\
  \Starting items: 53, 89, 62, 57, 74, 51, 83, 97\n\
  \Operation: new = old * 3\n\
  \Test: divisible by 13\n\
  \If true: throw to monkey 1\n\
  \If false: throw to monkey 5\n\
  \\n\
  \Monkey 1:\n\
  \Starting items: 85, 94, 97, 92, 56\n\
  \Operation: new = old + 2\n\
  \Test: divisible by 19\n\
  \If true: throw to monkey 5\n\
  \If false: throw to monkey 2\n\
  \\n\
  \Monkey 2:\n\
  \Starting items: 86, 82, 82\n\
  \Operation: new = old + 1\n\
  \Test: divisible by 11\n\
  \If true: throw to monkey 3\n\
  \If false: throw to monkey 4\n\
  \\n\
  \Monkey 3:\n\
  \Starting items: 94, 68\n\
  \Operation: new = old + 5\n\
  \Test: divisible by 17\n\
  \If true: throw to monkey 7\n\
  \If false: throw to monkey 6\n\
  \\n\
  \Monkey 4:\n\
  \Starting items: 83, 62, 74, 58, 96, 68, 85\n\
  \Operation: new = old + 4\n\
  \Test: divisible by 3\n\
  \If true: throw to monkey 3\n\
  \If false: throw to monkey 6\n\
  \\n\
  \Monkey 5:\n\
  \Starting items: 50, 68, 95, 82\n\
  \Operation: new = old + 8\n\
  \Test: divisible by 7\n\
  \If true: throw to monkey 2\n\
  \If false: throw to monkey 4\n\
  \\n\
  \Monkey 6:\n\
  \Starting items: 75\n\
  \Operation: new = old * 7\n\
  \Test: divisible by 5\n\
  \If true: throw to monkey 7\n\
  \If false: throw to monkey 0\n\
  \\n\
  \Monkey 7:\n\
  \Starting items: 92, 52, 85, 89, 68, 82\n\
  \Operation: new = old * old\n\
  \Test: divisible by 2\n\
  \If true: throw to monkey 0\n\
  \If false: throw to monkey 1\n"