Day 22: Use primops instead of C11 atomics

hs/aoc2024.cabal

@@ -20,7 +20,6 @@ source-repository head
 
 library
   hs-source-dirs: src
-  c-sources: src/cbits/atomics.c
   exposed-modules:
     Day1
     Day10
@@ -59,6 +58,7 @@ library
     heap ^>=1.0.4,
     megaparsec ^>=9.7.0,
     parallel ^>=3.2.2.0,
+    primitive ^>=0.9.0.0,
     split ^>=0.2.5,
     text ^>=2.1.2,
     vector ^>=0.13.2.0,

hs/app/Main.hs

@@ -79,7 +79,7 @@ main = do
   run 19 (uncurry (>>) . bimap print print) [Day19.solve]
   run 20 print [Day20.solve 2 100, Day20.solve 20 100]
   run 21 print [Day21.solve 2, Day21.solve 25]
-  run 22 (>>= print) [either fail pure . Day22.part1, Day22.part2]
+  run 22 (either fail print) [Day22.part1, Day22.part2]
   run 23 putStrLn [show . Day23.part1, T.unpack . Day23.part2]
   run 24 (either (fail . errorBundlePretty) (maybe (fail "error") putStrLn)) [fmap2 show . Day24.part1, fmap2 T.unpack . Day24.part2]
   run 25 print [Day25.part1]

hs/bench/Main.hs

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

hs/src/Day22.hs

@@ -1,21 +1,29 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# LANGUAGE UnboxedTuples #-}
+
 -- |
 -- Module:         Day22
 -- Description:    <https://adventofcode.com/2024/day/22 Day 22: Monkey Market>
 module Day22 (part1, part2) where
 
 import Common (readMany, readSome)
-import Control.Concurrent.Async (mapConcurrently)
-import Data.Array.IO (IOUArray)
-import Data.Array.MArray (newArray, readArray, writeArray)
+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 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, (!))
-import Foreign (Ptr, advancePtr, callocArray)
+import GHC.Exts (Int (I#), fetchAddIntArray#)
 
 step :: (Bits a, Num a) => a -> a
 step num = num3
@@ -31,24 +39,29 @@ part1 input = do
   where
     constants = V.generate 24 $ (!! 2000) . iterate step . bit
 
-part2 :: Text -> IO Int
-part2 input = case readSome T.decimal input of
-  Right (nums, _) -> do
-    acc <- callocArray $ rangeSize bounds
-    let go num = do
-          seen <- newArray @IOUArray bounds False
-          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) <$> atomic_fetch_add_int (acc `advancePtr` index bounds key) e
-              f _ = pure Nothing
-          foldMap' f (tails $ take 2001 $ map (`mod` 10) $ iterate step num)
-    mapConcurrently go nums >>= maybe (fail "error") (pure . getMax) . sconcat
-  Left err -> fail err
+part2 :: Text -> Either String Int
+part2 input = do
+  (nums, _) <- readSome 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
   where
     bounds = ((-9, -9, -9, -9), (9, 9, 9, 9))
 
-foreign import ccall unsafe atomic_fetch_add_int :: Ptr Int -> Int -> IO Int
+fetchAddIntArray :: (PrimMonad m) => MutablePrimArray (PrimState m) Int -> Int -> Int -> m Int
+fetchAddIntArray (MutablePrimArray mba#) (I# offset#) (I# incr#) = primitive $ \s1# ->
+  let !(# s2#, res# #) = fetchAddIntArray# mba# offset# incr# s1#
+   in (# s2#, I# res# #)

hs/test/Day22Spec.hs

@@ -5,7 +5,7 @@ module Day22Spec (spec) where
 import Data.Text (Text)
 import Data.Text qualified as T (unlines)
 import Day22 (part1, part2)
-import Test.Hspec (Spec, describe, it, shouldBe, shouldReturn)
+import Test.Hspec (Spec, describe, it, shouldBe)
 
 example1, example2 :: Text
 example1 =
@@ -30,4 +30,4 @@ spec = do
       part1 example1 `shouldBe` Right 37327623
   describe "part 2" $ do
     it "examples" $ do
-      part2 example2 `shouldReturn` 23
+      part2 example2 `shouldBe` Right 23