Merge pull request #223 from ephemient/hs/day22

hs/src/Day22.hs

@@ -1,26 +1,23 @@
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE UnboxedTuples #-}
-{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns -Wno-partial-type-signatures #-}
 
 -- |
 -- Module:         Day22
 -- Description:    <https://adventofcode.com/2024/day/22 Day 22: Monkey Market>
 module Day22 (part1, part2) where
 
-import Common (readMany, readSome)
+import Common (readMany)
+import Control.Monad (foldM)
 import Control.Monad.Primitive (PrimMonad, PrimState, primitive)
 import Control.Monad.ST (ST, runST)
 import Control.Monad.ST.Unsafe (unsafeInterleaveST)
-import Control.Parallel.Strategies (parTraversable, rseq, withStrategy)
+import Control.Parallel.Strategies (parList, rseq, withStrategy)
 import Data.Array.ST (MArray (newArray), STUArray, readArray, writeArray)
 import Data.Bits (Bits, bit, shiftL, shiftR, testBit, xor, (.&.))
-import Data.Foldable (Foldable (foldMap'))
 import Data.Ix (index, rangeSize)
-import Data.List (tails)
-import Data.Maybe (fromJust)
 import Data.Primitive (MutablePrimArray (MutablePrimArray), newPrimArray, setPrimArray)
-import Data.Semigroup (Max (Max, getMax), sconcat)
 import Data.Text (Text)
 import Data.Text.Read qualified as T (decimal)
 import Data.Vector.Unboxed qualified as V (generate, (!))
@@ -42,23 +39,26 @@ part1 input = do
 
 part2 :: Text -> Either String Int
 part2 input = do
-  (nums, _) <- readSome T.decimal input
+  (nums, _) <- readMany T.decimal input
   pure $ runST $ do
     acc <- newPrimArray $ rangeSize bounds
     setPrimArray acc 0 (rangeSize bounds) 0
-    let go num =
-          fromJust <$> do
-            seen <- newArray bounds False :: ST s (STUArray s _ _)
-            let f (a : b : c : d : e : _) =
-                  let key = (a - b, b - c, c - d, d - e)
-                   in readArray seen key >>= \case
-                        True -> pure Nothing
-                        False -> do
-                          writeArray seen key True
-                          Just . Max . (+) e <$> fetchAddIntArray acc (index bounds key) e
-                f _ = pure Nothing
-            foldMap' f $ tails $ map (`mod` 10) $ take 2001 $ iterate step num
-    getMax . sconcat . withStrategy (parTraversable rseq) <$> mapM (unsafeInterleaveST . go) nums
+    let go num = do
+          seen <- newArray bounds False :: ST s (STUArray s _ _)
+          let go' (a, b, c, d, k) e =
+                readArray seen key >>= \case
+                  False -> do
+                    writeArray seen key True
+                    k' <- max k . (+) e <$> fetchAddIntArray acc (index bounds key) e
+                    pure (b, c, d, e, k')
+                  True -> pure (b, c, d, e, k)
+                where
+                  key = (a - b, b - c, c - d, d - e)
+          (_, _, _, _, k) <- foldM go' (p0, p1, p2, p3, 0) $ take 1997 prices
+          pure k
+          where
+            p0 : p1 : p2 : p3 : prices = (`mod` 10) <$> iterate step num
+    foldl' max 0 . withStrategy (parList rseq) <$> mapM (unsafeInterleaveST . go) nums
   where
     bounds = ((-9, -9, -9, -9), (9, 9, 9, 9))