content: module Bench.MulMod (benches) where
  
  import BenchLib.MulMod.Barrett64 qualified as Barrett64
  import BenchLib.MulMod.BarrettWideWord qualified as BarrettWideWord
  import BenchLib.MulMod.Montgomery qualified as Montgomery
  import Criterion
  import Data.Bifunctor (first)
  import Data.Vector.Unboxed qualified as VU
  import Data.Word (Word32, Word64)
  import System.Random
  
  n :: Int
  n = 10000
  
  btWW :: BarrettWideWord.Barrett
  btWW = BarrettWideWord.new64 998244353
  
  bt64 :: Barrett64.Barrett
  bt64 = Barrett64.new 998244353
  
  mont :: Montgomery.Montgomery
  mont = Montgomery.new 998244353
  
  -- | This benchmark is almost nonsense, as the pre calculation' overhead is not considered while the
  -- real use case is @powMod@. However, it's useful when optimizing the successive calculation.
  benches32 :: Benchmark
  benches32 =
    bgroup
      "mulMod Word32 vector"
      [ bench "barrettWideWord" $ whnf (VU.foldl' (\acc -> BarrettWideWord.mulMod btWW acc . fromIntegral) w64) nonZeroRandomVec32,
        bench "barrett" $ whnf (VU.foldl' (\acc -> Barrett64.mulMod bt64 acc . fromIntegral) w64) nonZeroRandomVec32,
        -- NOTE: It skips the last reduce.
        bench "montgomery" $ whnf (VU.foldl' (\acc -> Montgomery.mulModGenerated mont acc . Montgomery.generate mont . fromIntegral) w64) nonZeroRandomVec32,
        bench "mod" $ whnf (VU.foldl' (\acc n -> acc * n `mod` 998244353) w32) nonZeroRandomVec32,
        bench "rem" $ whnf (VU.foldl' (\acc n -> acc * n `rem` 998244353) w32) nonZeroRandomVec32
      ]
    where
      w32 :: Word32
      w32 = 1
      w64 :: Word64
      w64 = 1
      -- [1, 998244383)
      nonZeroRandomVec32 :: VU.Vector Word32
      nonZeroRandomVec32 = VU.map fromIntegral $ VU.unfoldrExactN n (first (+ 1) . genWord64R (998244383 - 2)) (mkStdGen 123456789)
  
  -- | This benchmark is almost nonsense, as the pre calculation' overhead is not considered while the
  -- real use case is @powMod@. However, it's useful when optimizing the successive calculation.
  benches :: Benchmark
  benches =
    bgroup
      "mulMod Word64 vector"
      [ bench "barrettWideWord" $ whnf (VU.foldl' (BarrettWideWord.mulMod btWW) w64) nonZeroRandomVec64,
        bench "barrett" $ whnf (VU.foldl' (\acc -> Barrett64.mulMod bt64 acc) w64) nonZeroRandomVec64,
        -- NOTE: It skips the last reduce.
        bench "montgomery" $ whnf (VU.foldl' (\acc -> Montgomery.mulModGenerated mont acc . Montgomery.generate mont) w64) nonZeroRandomVec64,
        bench "mod" $ whnf (VU.foldl' (\acc n -> acc * n `mod` 998244353) w64) nonZeroRandomVec64,
        bench "rem" $ whnf (VU.foldl' (\acc n -> acc * n `rem` 998244353) w64) nonZeroRandomVec64
      ]
    where
      w64 :: Word64
      w64 = 1
      -- [1, 998244383)
      nonZeroRandomVec64 :: VU.Vector Word64
      nonZeroRandomVec64 = VU.map fromIntegral $ VU.unfoldrExactN n (first (+ 1) . genWord64R (998244383 - 2)) (mkStdGen 123456789)
  