Day9.hs

{-|
Module:         Day9
Description:    <https://adventofcode.com/2018/day/9 Day 9: Marble Mania>
-}
{-# LANGUAGE FlexibleContexts, Strict, TypeApplications #-}
module Day9 (day9a, day9b, play) where

import Control.Arrow (second)
import Control.Monad (foldM, forM_)
import Control.Monad.ST (ST, runST)
import Data.Array.Base (newArray, unsafeRead, unsafeWrite)
import Data.Array.ST (STUArray)
import Data.IntMap (elems, empty, insertWith)
import Text.Megaparsec (MonadParsec, parseMaybe, skipManyTill, takeRest)
import Text.Megaparsec.Char (satisfy)
import Text.Megaparsec.Char.Lexer (decimal)

parser :: (MonadParsec e String m, Integral players, Integral target) => m (players, target)
parser = (,) <$> decimal <*> skipManyTill (satisfy $ const True) decimal <* takeRest

play :: Int -> Int -> Int
play players target = maximum . elems $ runST $ do
    a <- newArray (0, target - 1) 0 :: ST s (STUArray s Int Int)
    let ix pos i = (pos + i) `mod` target
        acc (pos, size, scores) n
          | n `mod` 23 == 0 = do
                forM_ [1..6] $ \i -> unsafeRead a (ix pos $ size - i) >>= unsafeWrite a (ix pos (-i))
                m <- unsafeRead a $ ix pos $ size - 7
                let scores' = insertWith (+) (n `mod` players) (n + m) scores
                return (ix pos (-6), size - 1, scores')
          | otherwise = do
                unsafeRead a pos >>= unsafeWrite a (ix pos size)
                unsafeRead a (ix pos 1) >>= unsafeWrite a (ix pos $ size + 1)
                unsafeWrite a (ix pos 1) n
                return (ix pos 1, size + 1, scores)
    (_, _, scores) <- foldM acc (0, 1, empty) [1..target]
    return scores

day9a :: String -> Maybe Int
day9a = fmap (uncurry play) . parseMaybe @() parser

day9b :: String -> Maybe Int
day9b = fmap (uncurry play . second (* 100)) . parseMaybe @() parser