Merge pull request #222 from ephemient/hs/day11

ephemient · Jan 7, 2025 · 5ba3ba3 · 5ba3ba3
2 parents b15f043 + 685125b
commit 5ba3ba3
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 29 deletions.
diff --git a/.github/workflows/hs.yml b/.github/workflows/hs.yml
@@ -90,9 +90,9 @@ jobs:
     steps:
       - uses: actions/download-artifact@v4
       - run: chmod +x aoc2024-hs-9.10/aoc2024 aoc2024-hs-9.12/aoc2024
-      - run: aoc2024-hs-9.10/aoc2024
+      - run: aoc2024-hs-9.10/aoc2024 +RTS -s
         env:
           AOC2024_DATADIR: inputs
-      - run: aoc2024-hs-9.12/aoc2024
+      - run: aoc2024-hs-9.12/aoc2024 +RTS -s
         env:
           AOC2024_DATADIR: inputs
diff --git a/hs/src/Day11.hs b/hs/src/Day11.hs
@@ -7,19 +7,18 @@
 module Day11 (part1, part2, solve) where
 
 import Common (readMany)
-import Control.Monad (foldM, forM_)
+import Control.Monad (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 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, readArray, sizeofArray, sizeofMutableArray, writeArray)
+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#, toList)
+import GHC.Exts (Int (I#), casArray#, readArray#)
 
 part1, part2 :: Text -> IO Int
 part1 = solve 25
@@ -29,32 +28,38 @@ 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 $ \num -> do
-        let ix = num `mod` sizeofMutableArray array
-        readArray array ix >>= writeArray array ix . IntMap.insertWith (+) num 1
+  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 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 ()
+  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 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#
+  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