Merge pull request #219 from ephemient/hs/day11

hs/app/Main.hs

@@ -68,7 +68,7 @@ main = do
   run 8 print [Day8.part1, Day8.part2]
   run 9 print [Day9.part1, Day9.part2]
   run 10 print [Day10.part1, Day10.part2]
-  run 11 (either fail print) [Day11.part1, Day11.part2]
+  run 11 (>>= print) [Day11.part1, Day11.part2]
   run 12 print [Day12.part1, Day12.part2]
   run 13 (either (fail . errorBundlePretty) print) [Day13.part1, Day13.part2]
   run 14 (either (fail . errorBundlePretty) print) [Day14.part1, Day14.part2]

hs/bench/Main.hs

@@ -1,7 +1,7 @@
 module Main (main) where
 
 import Control.Arrow ((>>>))
-import Criterion.Main (bench, bgroup, defaultMain, env, envWithCleanup, nf)
+import Criterion.Main (bench, bgroup, defaultMain, env, envWithCleanup, nf, nfAppIO)
 import Data.Foldable (find)
 import Data.Maybe (fromMaybe)
 import Data.Text (Text)
@@ -111,8 +111,8 @@ main =
       env (getDayInput 11) $ \input ->
         bgroup
           "Day 11"
-          [ bench "part 1" $ nf Day11.part1 input,
-            bench "part 2" $ nf Day11.part2 input
+          [ bench "part 1" $ nfAppIO Day11.part1 input,
+            bench "part 2" $ nfAppIO Day11.part2 input
           ],
       env (getDayInput 12) $ \input ->
         bgroup

hs/src/Day11.hs

@@ -1,30 +1,60 @@
-{-# LANGUAGE MultiWayIf #-}
+{-# 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 (foldM, forM_)
+import Control.Monad.Primitive (PrimMonad, PrimState, primitive)
+import Control.Monad.ST.Unsafe (unsafeInterleaveST)
+import Control.Parallel.Strategies (parList, rseq, withStrategy)
+import Data.Bifunctor (Bifunctor (bimap))
 import Data.IntMap (IntMap)
-import Data.IntMap qualified as IntMap (toList)
-import Data.IntMap.Strict qualified as IntMap (fromListWith)
+import Data.IntMap qualified as IntMap (empty, toList)
+import Data.IntMap.Strict qualified as IntMap (insertWith)
+import Data.Primitive.Array (Array, MutableArray (marray#), createArray, readArray, sizeofArray, sizeofMutableArray, writeArray)
 import Data.Text (Text)
 import Data.Text.Read qualified as T (decimal)
+import GHC.Conc (getNumCapabilities)
+import GHC.Exts (Int (I#), casArray#, readArray#, toList)
 
-part1, part2 :: Text -> Either String Int
+part1, part2 :: Text -> IO Int
 part1 = solve 25
 part2 = solve 75
 
-solve :: Int -> Text -> Either String Int
+solve :: Int -> Text -> IO Int
 solve n input = do
-  (nums, _) <- readMany T.decimal input
-  pure $ foldl' (+) 0 $ iterate step (IntMap.fromListWith (+) $ (,1) <$> nums) !! n
+  (nums, _) <- either fail pure $ readMany T.decimal input
+  numCapabilities <- max 1 <$> getNumCapabilities
+  let start = createArray numCapabilities IntMap.empty $ \array -> forM_ nums $ \num -> do
+        let ix = num `mod` sizeofMutableArray array
+        readArray array ix >>= writeArray array ix . IntMap.insertWith (+) num 1
+      end = iterate step start !! n
+  pure $ foldl' (flip $ (+) . foldl' (+) 0) 0 end
 
-step :: IntMap Int -> IntMap Int
-step counts = IntMap.fromListWith (+) $ do
-  (x, n) <- IntMap.toList counts
-  if
-    | x == 0 -> [(1, n)]
-    | s <- show x, (l, 0) <- length s `divMod` 2, (a, b) <- splitAt l s -> [(read a, n), (read b, n)]
-    | otherwise -> [(2024 * x, n)]
+step :: Array (IntMap Int) -> Array (IntMap Int)
+step array = createArray (sizeofArray array) IntMap.empty $ \array' -> do
+  let insert x = modifyArray array' (x `mod` sizeofMutableArray array') . IntMap.insertWith (+) x
+      go _ (0, n) = insert 1 n
+      go _ (x, n)
+        | s <- show x,
+          (l, 0) <- length s `divMod` 2 = do
+            let (y, z) = bimap read read $ splitAt l s
+            insert y n
+            insert z n
+        | otherwise = insert (2024 * x) n
+  results <- mapM (unsafeInterleaveST . foldM go () . IntMap.toList) $ toList array
+  pure $! foldl' (flip seq) () $ withStrategy (parList rseq) results
+
+modifyArray :: (PrimMonad m) => MutableArray (PrimState m) a -> Int -> (a -> a) -> m ()
+modifyArray array (I# index#) f = primitive $ \s1# ->
+  let array# = marray# array
+      modifyArray# s3# b =
+        case casArray# array# index# b (f b) s3# of
+          (# s4#, 0#, _ #) -> (# s4#, () #)
+          (# s4#, _, c #) -> modifyArray# s4# c
+      !(# s2#, a #) = readArray# array# index# s1#
+   in modifyArray# s2# a

hs/test/Day11Spec.hs

@@ -4,7 +4,7 @@ module Day11Spec (spec) where
 
 import Data.Text (Text)
 import Day11 (solve)
-import Test.Hspec (Spec, describe, it, shouldBe)
+import Test.Hspec (Spec, describe, it, shouldReturn)
 
 example1, example2 :: Text
 example1 = "0 1 10 99 999\n"
@@ -14,6 +14,6 @@ spec :: Spec
 spec = do
   describe "part 1" $ do
     it "examples" $ do
-      solve 1 example1 `shouldBe` Right 7
-      solve 6 example2 `shouldBe` Right 22
-      solve 25 example2 `shouldBe` Right 55312
+      solve 1 example1 `shouldReturn` 7
+      solve 6 example2 `shouldReturn` 22
+      solve 25 example2 `shouldReturn` 55312