Day11.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module:         Day11
-- Description:    <https://adventofcode.com/2024/day/11 Day 11: Plutonian Pebbles>
module Day11 (part1, part2, solve) where

import Common (readMany)
import Control.Monad (forM_)
import Control.Monad.Primitive (PrimMonad, PrimState, primitive)
import Control.Monad.ST.Unsafe (unsafeInterleaveST)
import Control.Parallel.Strategies (parTraversable, rseq, withStrategy)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap (empty, toList)
import Data.IntMap.Strict qualified as IntMap (insertWith)
import Data.Primitive.Array (Array, MutableArray (marray#), createArray, sizeofArray, sizeofMutableArray)
import Data.Text (Text)
import Data.Text.Read qualified as T (decimal)
import GHC.Conc (getNumCapabilities)
import GHC.Exts (Int (I#), casArray#, readArray#)

part1, part2 :: Text -> IO Int
part1 = solve 25
part2 = solve 75

solve :: Int -> Text -> IO Int
solve n input = do
  (nums, _) <- either fail pure $ readMany T.decimal input
  numCapabilities <- max 1 <$> getNumCapabilities
  let start = createArray numCapabilities IntMap.empty $ \array ->
        forM_ nums $ \x -> insert array x 1
      end = iterate step start !! n
  pure $ foldl' (foldl' (+)) 0 end

step :: Array (IntMap Int) -> Array (IntMap Int)
step array = createArray (sizeofArray array) IntMap.empty $ \array' -> do
  let go (0, n) = insert array' 1 n
      go (x, n)
        | Just (y, z) <- splitDigits x = insert array' y n >> insert array' z n
        | otherwise = insert array' (2024 * x) n
  results <- mapM (unsafeInterleaveST . mapM_ go . IntMap.toList) array
  pure $! foldl' (flip seq) () $ withStrategy (parTraversable rseq) results

splitDigits :: (Integral a) => a -> Maybe (a, a)
splitDigits x = splitDigits' x 10
  where
    splitDigits' y n
      | y < 10 = Nothing
      | y < 100 = Just $ x `divMod` n
      | otherwise = splitDigits' (y `div` 100) $! 10 * n

insert :: (PrimMonad m, Num a) => MutableArray (PrimState m) (IntMap a) -> Int -> a -> m (IntMap a)
insert array x = modifyArray array (x `mod` sizeofMutableArray array) . IntMap.insertWith (+) x

modifyArray :: (PrimMonad m) => MutableArray (PrimState m) a -> Int -> (a -> a) -> m a
modifyArray array (I# index#) f = primitive $ \s1# ->
  let !(# s2#, a #) = readArray# array# index# s1#
   in modifyArray# s2# a
  where
    array# = marray# array
    modifyArray# s1# a =
      case casArray# array# index# a (f a) s1# of
        (# s2#, 0#, b #) -> (# s2#, b #)
        (# s2#, _, b #) -> modifyArray# s2# b