AoC 2024, day 11, part 2

Day11.hs

type HashTable s k v = C.HashTable s k v

type Stone = Int

type Memo s = HashTable s (Stone, Int) Int

solvePart2 :: BS.ByteString -> Result
solvePart2 contents =
  Result $ runST $ do
    memo :: Memo s <- H.new
    let stoneCount _ 0 = pure 1
        stoneCount stone n = do
          H.lookup memo (stone, n) >>= \case
            Just count -> pure count
            Nothing -> do
              foldM
                ( \acc stone' -> do
                    count' <- stoneCount stone' (n - 1)
                    H.insert memo (stone', n - 1) count'
                    pure $ acc + count'
                )
                0
                $ updateStone stone

    foldM (\acc stone -> (acc +) <$> stoneCount stone (75 :: Int)) 0 $
      parseStones contents

parseStones :: BS.ByteString -> Seq Stone
parseStones = Seq.fromList . map read . words . BS8.unpack

updateStone :: Stone -> Seq Stone
updateStone stone
  | stone == 0 = Seq.singleton 1
  | s <- show stone
  , len <- length s
  , even len
  , (s1, s2) <- splitAt (len `div` 2) s =
      Seq.fromList [read s1, read s2]
  | otherwise = Seq.singleton $ stone * 2024