ac-library-hs 1.2.5.0 → 1.2.6.0

raw patch · 15 files changed

+1307/−15 lines, 15 filesdep +semiringsdep ~randomPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: semirings

Dependency ranges changed: random

API changes (from Hackage documentation)

Files

• CHANGELOG.md +13/−4
• ac-library-hs.cabal +9/−1
• benchmarks/Bench/RepeatWithIndex.hs +67/−0
• benchmarks/Bench/RepeatWithoutIndex.hs +75/−0
• benchmarks/Main.hs +4/−0
• src/AtCoder/Extra/Graph.hs +3/−3
• src/AtCoder/Extra/Math.hs +322/−1
• src/AtCoder/Extra/Math/Montgomery64.hs +176/−0
• src/AtCoder/Extra/ModInt64.hs +259/−0
• src/AtCoder/Extra/SqrtDecomposition.hs +19/−3
• src/AtCoder/ModInt.hs +3/−3
• test/Main.hs +4/−0
• test/Tests/Extra/Math.hs +90/−0
• test/Tests/Extra/Math/Montgomery64.hs +44/−0
• test/Tests/Extra/ModInt64.hs +219/−0

@@ -1,9 +1,18 @@ # Revision history for acl-hs +## 1.2.6.0 -- April 2025++- Added `AtCoder.Extra.Math` functions+  - `isPrime`+  - `primes`+  - `primeFactors`+- Added `AtCoderExtra.Math.Montgomery64`+- Added `AtCoderExtra.ModInt64`+ ## 1.2.5.0 -- April 2025 -- Added AtCoder.`Extra.Mo`-- Added AtCoder.`Extra.SqrtDecomposition`+- Added `AtCoder.Extra.Mo`+- Added `AtCoder.Extra.SqrtDecomposition`  ## 1.2.4.0 -- April 2025 @@ -29,9 +38,9 @@  ## 1.2.2.0 -- Feb 2025 -- Added `AtCoder.Extra.KdTree` and `Extra.LazyKdTree`.+- Added `AtCoder.Extra.KdTree` and `AtCoder.Extra.LazyKdTree`. - Added `clear` function to the dynamic segment tree family.-- Fixed AtCoder.`Extra.Hld.new` for a tree with a single vertex.+- Fixed `AtCoder.Extra.Hld.new` for a tree with a single vertex.  ## 1.2.1.0 -- Feb 2025 

@@ -4,7 +4,7 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.2.5.0+version:         1.2.6.0 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive@@ -39,6 +39,7 @@     , bitvec             <1.2     , bytestring         <0.14     , primitive          >=0.6.4.0 && <0.10+    , random             >=1.2.0     , vector             >=0.13.0  && <0.14     , vector-algorithms  <0.10     , wide-word          <0.2@@ -70,7 +71,9 @@     AtCoder.Extra.KdTree     AtCoder.Extra.LazyKdTree     AtCoder.Extra.Math+    AtCoder.Extra.Math.Montgomery64     AtCoder.Extra.Mo+    AtCoder.Extra.ModInt64     AtCoder.Extra.Monoid     AtCoder.Extra.Monoid.Affine1     AtCoder.Extra.Monoid.Mat2x2@@ -155,6 +158,8 @@     Tests.Extra.KdTree     Tests.Extra.LazyKdTree     Tests.Extra.Math+    Tests.Extra.Math.Montgomery64+    Tests.Extra.ModInt64     Tests.Extra.Monoid     Tests.Extra.MultiSet     Tests.Extra.SegTree2d@@ -199,6 +204,7 @@     , QuickCheck     , quickcheck-classes     , random+    , semirings     , tasty     , tasty-hspec     , tasty-hunit@@ -219,6 +225,8 @@     Bench.Matrix     Bench.ModInt     Bench.PowMod+    Bench.RepeatWithIndex+    Bench.RepeatWithoutIndex     Bench.SwapDupe     BenchLib.AddMod     BenchLib.Matrix

@@ -0,0 +1,67 @@+-- | Benchmark for monadic streams.+module Bench.RepeatWithIndex (benches) where++import Criterion+import Control.Monad (when, replicateM_)+import Control.Monad.ST (runST)+import Data.Foldable (for_)+import Data.Vector.Fusion.Stream.Monadic qualified as MS+import Data.Vector.Generic.Mutable as VGM+import Data.Vector.Unboxed as VU+import Data.Vector.Unboxed.Mutable as VUM++len :: Int+len = 10 ^ 7++list :: Int -> Int+list x = runST $ do+  res <- VUM.replicate 1 x+  for_ [0 .. len - 1] $ \dx -> do+    VGM.modify res (+ dx) 0+  VGM.read res 0++-- | @cojna/iota+(..<) :: (Monad m) => Int -> Int -> MS.Stream m Int+(..<) !l !r = MS.Stream step l+  where+    step x+      | x < r = return $ MS.Yield x (x + 1)+      | otherwise = return MS.Done+    {-# INLINE [0] step #-}+{-# INLINE [1] (..<) #-}++stream :: Int -> Int+stream x = runST $ do+  res <- VUM.replicate 1 x+  flip MS.mapM_ (0 ..< len) $ \dx -> do+    VGM.modify res (+ dx) 0+  VGM.read res 0++vector :: Int -> Int+vector x = runST $ do+  res <- VUM.replicate 1 x+  VU.forM_ (VU.generate len id) $ \dx -> do+    VGM.modify res (+ dx) 0+  VGM.read res 0++recursion :: Int -> Int+recursion x = runST $ do+  res <- VUM.replicate 1 x+  let run i = do+        when (i < len) $ do+          VGM.modify res (+ i) 0+          run $ i + 1+  run 0+  VGM.read res 0++-- The result is suspicious.. the stream version is slower than vector, recursive or replicateM_ is+-- faster by 10 times?+benches :: Benchmark+benches =+  bgroup+    "repeat-with-index"+    [ bench "list" $ nf list 0,+      bench "stream" $ nf stream 0,+      bench "vector" $ nf vector 0,+      bench "recursion" $ nf recursion 0+    ]

@@ -0,0 +1,75 @@+-- | Benchmark for monadic streams.+module Bench.RepeatWithoutIndex (benches) where++import Criterion+import Control.Monad (when, replicateM_)+import Control.Monad.ST (runST)+import Data.Foldable (for_)+import Data.Vector.Fusion.Stream.Monadic qualified as MS+import Data.Vector.Generic.Mutable as VGM+import Data.Vector.Unboxed as VU+import Data.Vector.Unboxed.Mutable as VUM++len :: Int+len = 10 ^ 7++list :: Int -> Int+list x = runST $ do+  res <- VUM.replicate 1 x+  for_ [0 .. len - 1] $ \_ -> do+    VGM.modify res (+ 1) 0+  VGM.read res 0++-- | @cojna/iota+(..<) :: (Monad m) => Int -> Int -> MS.Stream m Int+(..<) !l !r = MS.Stream step l+  where+    step x+      | x < r = return $ MS.Yield x (x + 1)+      | otherwise = return MS.Done+    {-# INLINE [0] step #-}+{-# INLINE [1] (..<) #-}++stream :: Int -> Int+stream x = runST $ do+  res <- VUM.replicate 1 x+  flip MS.mapM_ (0 ..< len) $ \_ -> do+    VGM.modify res (+ 1) 0+  VGM.read res 0++vector :: Int -> Int+vector x = runST $ do+  res <- VUM.replicate 1 x+  VU.forM_ (VU.generate len id) $ \_ -> do+    VGM.modify res (+ 1) 0+  VGM.read res 0++recursion :: Int -> Int+recursion x = runST $ do+  res <- VUM.replicate 1 x+  let run i = do+        when (i < len) $ do+          VGM.modify res (+ 1) 0+          run $ i + 1+  run 0+  VGM.read res 0++rep :: Int -> Int+rep x = runST $ do+  res <- VUM.replicate 1 x+  replicateM_ len $ do+    VGM.modify res (+ 1) 0+  VGM.read res 0++-- The result is suspicious.. the stream version is slower than vector, recursive or replicateM_ is+-- faster by 10 times?+benches :: Benchmark+benches =+  bgroup+    "repeat-without-index"+    [ bench "list" $ nf list 0,+      bench "stream" $ nf stream 0,+      bench "vector" $ nf vector 0,+      bench "recursion" $ nf recursion 0,+      bench "rep" $ nf rep 0+    ]

@@ -7,6 +7,8 @@ import Bench.ModInt qualified import Bench.MulMod qualified import Bench.PowMod qualified+import Bench.RepeatWithIndex qualified+import Bench.RepeatWithoutIndex qualified import Bench.SwapDupe qualified import Criterion.Main @@ -22,5 +24,7 @@       Bench.AddMod.benches,       Bench.PowMod.benches,       Bench.Matrix.benches,+      Bench.RepeatWithIndex.benches,+      Bench.RepeatWithoutIndex.benches,       Bench.SwapDupe.benches     ]

@@ -99,7 +99,7 @@ import AtCoder.Internal.MinHeap qualified as MH import AtCoder.Internal.Queue qualified as Q import AtCoder.Internal.Scc qualified as ACISCC-import Control.Monad (when)+import Control.Monad (replicateM_, when) import Control.Monad.Fix (fix) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Control.Monad.ST (ST, runST)@@ -443,7 +443,7 @@                           VGM.unsafeWrite next 0 (nxt + 1)                           B.pushBack edges (nxt, v)                           len <- B.length st-                          for_ [1 .. len - s] $ \_ -> do+                          replicateM_ (len - s) $ do                             back <- fromJust <$> B.popBack st                             B.pushBack edges (nxt, back)                         pure (child', k')@@ -477,7 +477,7 @@   n' <- VGM.unsafeRead next 0   Csr.build' n' <$> B.unsafeFreeze edges --- | \(O(n + m)\) Returns a [blocks (biconnected comopnents)](https://en.wikipedia.org/wiki/Biconnected_component)+-- | \(O(n + m)\) Returns [blocks (biconnected comopnents)](https://en.wikipedia.org/wiki/Biconnected_component) -- of the graph. -- -- ==== __Example__

@@ -7,6 +7,14 @@     ACIM.invGcd,     primitiveRoot32, +    -- * Prime numbers and divisors+    primes,+    isPrime,+    primeFactors,+    primeFactorsUnsorted,+    divisors,+    divisorsUnsorted,+     -- * Binary exponentiation      -- | ==== __Examples__@@ -26,10 +34,22 @@   ) where +import AtCoder.Extra.Math.Montgomery64 qualified as M64 import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Math qualified as ACIM-import Data.Bits ((.>>.))+import Control.Monad (unless, when)+import Data.Bit (Bit (..))+import Data.Bits (bit, countTrailingZeros, (.<<.), (.>>.))+import Data.Foldable (for_)+import Data.Maybe (fromJust)+import Data.Vector.Algorithms.Intro qualified as VAI+import Data.Vector.Algorithms.Radix qualified as VAR+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import Data.Word (Word64) import GHC.Stack (HasCallStack)+import System.Random  -- | \(O(k \log^3 n) (k = 3)\). Returns whether the given `Int` value is a prime number. --@@ -57,6 +77,307 @@ primitiveRoot32 x = ACIM.primitiveRoot x   where     !_ = ACIA.runtimeAssert (x < (1 .>>. 32)) $ "AtCoder.Extra.Math.primitiveRoot32: given too large number `" ++ show x ++ "`"++-- | \(O(n \log \log n)\) Creates an array of prime numbers up to the given limit, using Sieve of+-- Eratosthenes.+--+-- The minimum computational complexity is \(\Omega(B \log \log B)\), where \(B = 2^{15}\) is the+-- length of segment. This constraint comes from the use of segmented sieve.+--+-- ==== Constraints+-- - The upper limit must be less than or equal to \(2^{30} (\gt 10^9)\), otherwise the returned+-- prime table is incorrect.+--+-- @since 1.2.6.0+{-# INLINEABLE primes #-}+primes :: Int -> VU.Vector Int+primes upperLimit+  | upperLimit <= 1 = VU.empty+  | otherwise = VU.create $ do+      -- segment length (TODO: isn't it 32767?)+      let !s = 32768 :: Int -- 2 ^ 15++      -- sieve length (TODO: use \sqrt limit? do benchmark)+      let !sieveMax = s++      -- Is it like LT bound??+      let !limit = upperLimit + 1++      -- base primes with index+      (!ps, !is) <- do+        sieve <- VUM.replicate (sieveMax + 1) $ Bit False+        ps <- VUM.unsafeNew (sieveMax `div` 2)+        is <- VUM.unsafeNew (sieveMax `div` 2)+        -- FIXME: carry index?+        iNext <- VUM.replicate 1 (0 :: Int)+        for_ [3, 5 .. s] $ \p1 -> do+          Bit b <- VGM.read sieve p1+          unless b $ do+            at <- VGM.read iNext 0+            VGM.write iNext 0 $ at + 1+            -- (base prime, next index (odd numbers only, so `div` 2)+            VGM.write ps at p1+            VGM.write is at $ p1 * p1 `div` 2+            -- NOTE: if `j` is a composite number, it's already enumerated by a smaller prime+            -- number than `p0`, so skip to `p1 * p1` and iterate through odd numbers only:+            for_ [p1 * p1, p1 * p1 + 2 * p1 .. sieveMax] $ \np1 -> do+              VGM.write sieve np1 $ Bit True+        len <- VGM.read iNext 0+        (,VGM.take len is) <$> VU.unsafeFreeze (VGM.take len ps)++      -- https://en.wikipedia.org/wiki/Prime-counting_function+      let !maxPrimeCount :: Int+            -- NOTE: 1,700 is a point where the next function estimates better as far as I tested:+            | limit < 1700 = round (1.25506 * fromIntegral limit / log (fromIntegral limit) :: Double)+            -- Rosser and Schoenfeld Boundsh (1962): holds for x > e^{3/2}:+            | limit < 60184 = round (fromIntegral limit / (log (fromIntegral limit) - 1.5) :: Double)+            -- Pierre Dusart (2010): holds for x >= 60184:+            | otherwise = ceiling (fromIntegral limit / (log (fromIntegral limit) - 1.1) :: Double)++      -- let f x = round (1.25506 * fromIntegral x / log (fromIntegral x) :: Double)+      -- let g x = round (fromIntegral x / (log (fromIntegral x) - 1.5) :: Double)+      -- let h x = ceiling (fromIntegral x / (log (fromIntegral x) - 1.1) :: Double)+      -- let p x = (f x, g x, h x)++      result <- VUM.replicate maxPrimeCount (-1)+      VGM.write result 0 2+      -- FIXME: carry index?+      nPrimes <- VUM.replicate 1 (1 :: Int)++      -- Sieve of Eratosthenes by block of size `s`, ignoring even numers+      -- FIXME: block length of size `s/2` should make more sense?+      block <- VUM.unsafeNew s+      let !r = limit `div` 2+      for_ [1, 1 + s .. r] $ \l -> do+        VGM.set block $ Bit False++        VU.iforM_ ps $ \idx p -> do+          -- FIXME: cut out the ps beforehand+          when (p <= limit) $ do+            i0 <- VGM.read is idx+            let run i = do+                  if i < l + s+                    then do+                      -- within the block+                      VGM.write block (i - l) $ Bit True+                      run $ i + p+                    else do+                      -- went out of the block+                      VGM.write is idx i+            run i0++        block' <- VU.take (min s (r - l)) <$> VU.unsafeFreeze block+        VU.iforM_ block' $ \i (Bit b) -> do+          unless b $ do+            at <- VGM.read nPrimes 0+            when (at < maxPrimeCount) $ do+              VGM.write nPrimes 0 $ at + 1+              VGM.write result at $ (l + i) * 2 + 1++      len <- VGM.read nPrimes 0+      pure $ VGM.take len result++-- | \(O(w \log^3 n)\) Miller–Rabin primality test, where \(w = 3\) for \(x \lt 2^{32}\) and+-- \(w = 7\) for \(x \ge 3^{32}\).+--+-- @since 1.2.6.0+{-# INLINEABLE isPrime #-}+isPrime :: Int -> Bool+isPrime x+  | x <= 1 = False+  -- Up to 11^2:+  | x == 2 || x == 3 || x == 5 || x == 7 = True+  | even x || x `rem` 3 == 0 || x `rem` 5 == 0 || x `rem` 7 == 0 = False+  | x < 121 = True+isPrime x+  -- http://miller-rabin.appspot.com/+  -- \| x < bit 32 = all test [2, 7, 61]+  | x < bit 32 = test 2 && test 7 && test 61+  -- \| otherwise = all test [2, 325, 9375, 28178, 450775, 9780504, 1795265022]+  | otherwise = test 2 && test 325 && test 9375 && test 28178 && test 450775 && test 9780504 && test 1795265022+  where+    !x64 :: Word64 = fromIntegral x+    !d :: Word64 = (x64 - 1) .>>. countTrailingZeros (x64 - 1)+    !mont = M64.fromVal x64+    !one = M64.encode mont 1+    !minusOne = M64.encode mont (x64 - 1)+    test a = inner (M64.powMod mont (M64.encode mont a) (fromIntegral d)) d+      where+        inner :: Word64 -> Word64 -> Bool+        inner y t+          | not (M64.eq x64 y one) && not (M64.eq x64 y minusOne) && t /= x64 - 1 = inner (M64.mulMod mont y y) (t .<<. 1)+          | not (M64.eq x64 y minusOne) && even t = False+          | otherwise = True++-- | Pollard's Rho algorithm.+{-# INLINEABLE rho #-}+rho :: (HasCallStack) => Word64 -> Int -> Int -> Int+rho modVal n c+  | n < 1 = error $ "AtCoder.Extra.Math.rho: given value less than or equal to `1`: `" ++ show n ++ show "`"+  | otherwise = fromIntegral $! inner 1 (M64.encode mont 1) (M64.encode mont 2) (M64.encode mont 1) (M64.encode mont 1) 1+  where+    -- what a mess!!+    !mont = M64.fromVal modVal+    !n64 :: Word64 = fromIntegral n+    !cc = M64.encode mont $ fromIntegral c+    f !x = M64.addMod modVal (M64.mulMod mont x x) cc+    fn 0 !x = x+    fn n_ !x = fn (n_ - 1) $! f x+    !m2 :: Int = bit $ floor (logBase (2.0 :: Double) (fromIntegral n)) `div` 5+    inner r _lastY0 y0 z0 q0 g0+      | g0 == 1 =+          let !y = fn r y0+              (!y', !z', !q', !g') = inner2 0 y z0 q0 g0+           in inner (r .<<. 1) y0 y' z' q' g'+      -- FIXME: It can sometimes slow, depending on the seed value+      -- \| g0 == n64 = inner3 z0+      | otherwise = g0+      where+        inner2 !k !y !z !q !g+          | k >= r || g /= 1 = (y, z, q, g)+          | otherwise =+              let (!y', !q') = fn2 (min m2 (r - k)) y q+                  !g' = gcd (M64.decode mont q) n64+               in inner2 (k + m2) y' y q' g'+          where+            fn2 0 !y_ !q_ = (y_, q_)+            fn2 n_ !y_ !q_ =+              let !y' = f y_+                  !q' = M64.mulMod mont q_ (M64.subMod modVal y0 y')+               in fn2 (n_ - 1) y' q'++-- FIXME: it can sometimes slow, depending on the seed value+-- inner3 !z+--   | g == 1 = inner3 z'+--   | otherwise = g+--   where+--     !z' = f z+--     !g = gcd (M64.decode mont (M64.subMod modVal lastY0 z')) n64++-- | Tries to find a prime factor for the given value, running Pollard's Rho algorithm.+--+-- ==== Constrants+-- - \(x \gt 1\)+{-# INLINEABLE findPrimeFactor #-}+findPrimeFactor :: (HasCallStack) => StdGen -> Int -> (Maybe Int, StdGen)+findPrimeFactor gen0 n+  | n <= 1 = error $ "AtCoder.Extra.Math.findPrimeFactor: given value less than or equal to `1`: " ++ show n ++ show "`"+  | isPrime n = (Just n, gen0)+  | otherwise = tryN 200 gen0+  where+    tryN :: Int -> StdGen -> (Maybe Int, StdGen)+    tryN 0 gen = (Nothing, gen)+    tryN i gen+      | isPrime m = (Just m, gen')+      | otherwise = tryN (i - 1) gen'+      where+        (!rnd, !gen') = uniformR (0, n - 1) gen+        !m = rho (fromIntegral n) n rnd++-- | Returns prime factors in run-length encoding \((p_i, n_i)\), sorted by \(p_i\).+--+-- ==== Constraints+-- - \(x \ge 1\)+--+-- @since 1.2.6.0+{-# INLINE primeFactors #-}+primeFactors :: (HasCallStack) => Int -> VU.Vector (Int, Int)+primeFactors = VU.modify VAI.sort . primeFactorsUnsorted++-- | Returns prime factors in run-length encoding \((p_i, n_i)\) in arbitrary order.+--+-- It internally uses probabilistic method (Pollard's rho algorithm) and it can actually result in+-- runtime error, however, the probability is very low and the API does not return `Maybe`.+--+-- @since 1.2.6.0+{-# INLINEABLE primeFactorsUnsorted #-}+primeFactorsUnsorted :: (HasCallStack) => Int -> VU.Vector (Int, Int)+primeFactorsUnsorted n+  | n < 1 = error $ "AtCoder.Extra.Math.primeFactorsUnsorted: given non-positive value `" ++ show n ++ "`"+  | otherwise = VU.create $ do+      buf <- VUM.unsafeNew (ceiling (logBase (2 :: Double) (fromIntegral n)))++      -- for small prime factors, try them all:+      let runDiv cur iWrite [] = pure (cur, iWrite)+          runDiv cur iWrite (d : rest)+            | d * d > cur = pure (cur, iWrite)+            | otherwise = case tryDiv 0 cur d of+                Just (!cur', !nd) -> do+                  VGM.write buf iWrite (d, nd)+                  runDiv cur' (iWrite + 1) rest+                Nothing -> runDiv cur iWrite rest++      (!n', !iWrite0) <- runDiv n 0 (2 : [3, 5 .. 97])++      -- for bigger prime numbers, use Polland's rho algorithm:+      let runRho !gen !cur !iWrite+            | cur > 1 = case findPrimeFactor gen cur of+                (Just p, !gen') -> do+                  let (!cur', !np) = fromJust $ tryDiv 0 cur p+                  VGM.write buf iWrite (p, np)+                  runRho gen' cur' (iWrite + 1)+                (Nothing, !_gen') -> do+                  -- we could return `Nothing` instead+                  error $ "unable to find prime factor for " ++ show cur+            | otherwise = pure iWrite++      -- NOTE: The seed value ifs fixed here. We could decide it at runtime for possibly faster+      -- submissions on TLE redjuge, however, 're rather preferring deterministic result:+      len <- runRho (mkStdGen 123456789) n' iWrite0+      pure $ VUM.take len buf+  where+    tryDiv :: Int -> Int -> Int -> Maybe (Int, Int)+    tryDiv !nDiv x d+      | r == 0 = tryDiv (nDiv + 1) q d+      | nDiv > 0 = Just (x, nDiv)+      | otherwise = Nothing+      where+        (!q, !r) = x `quotRem` d++-- | Enumerates divisors of the input value.+--+-- ==== Constraints+-- - \(x \ge 1\)+--+-- @since 1.2.6.0+{-# INLINE divisors #-}+divisors :: Int -> VU.Vector Int+-- TODO: use intro sort?+divisors = VU.modify VAR.sort . divisorsUnsorted++-- | Enumerates divisors of the input value.+--+-- ==== Constraints+-- - \(x \ge 1\)+--+-- @since 1.2.6.0+{-# INLINEABLE divisorsUnsorted #-}+divisorsUnsorted :: Int -> VU.Vector Int+divisorsUnsorted x = VU.create $ do+  vec <- VUM.unsafeNew nDivisors+  VGM.write vec 0 1+  VU.foldM'_+    ( \lenSofar (!p, !np) -> do+        (fst <$>)+          $ VU.foldM'+            ( \(!offset, !pp) _ -> do+                let !pp' = pp * p+                -- multiply to all the values sofar:+                VGM.iforM_ (VGM.take lenSofar vec) $ \i vx -> do+                  VGM.write vec (offset + i) $! vx * pp'+                  pure pp'+                pure (offset + lenSofar, pp')+            )+            (lenSofar, 1 :: Int)+          $ VU.generate np (+ 1)+    )+    (1 :: Int)+    pns+  pure vec+  where+    pns = primeFactors x+    (!_, !ns) = VU.unzip pns+    nDivisors = VU.foldl' (\ !acc n -> acc * (n + 1)) (1 :: Int) ns  -- | Calculates \(x^n\) with custom multiplication operator using the binary exponentiation -- technique.

@@ -0,0 +1,176 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE NamedFieldPuns #-}++-- | Fast modular multiplication for `Word64` using Montgomery multiplication. If the modulus value+-- is known to fit in 32 bits, use the @AtCoder.Internal.Barrett@ module instead.+--+-- @since 1.2.6.0+module AtCoder.Extra.Math.Montgomery64+  ( -- * Montgomery64+    Montgomery64,++    -- * Constructor+    new,+    fromVal,++    -- * Accessor+    umod,++    -- * Montgomery form encoding+    encode,+    decode,+    reduce,++    -- * Calculations+    addMod,+    subMod,+    mulMod,+    powMod,+    eq,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Data.Bits (bit, (!>>.))+import Data.WideWord.Word128 (Word128 (..))+import Data.Word (Word64)+import GHC.Exts (Proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal')++-- TODO: provide with newtype for Montgomery form?++-- | Fast modular multiplication for `Word64` using Montgomery64 multiplication.+--+-- @since 1.2.6.0+data Montgomery64 = Montgomery64+  { mM64 :: {-# UNPACK #-} !Word64,+    rM64 :: {-# UNPACK #-} !Word64,+    n2M64 :: {-# UNPACK #-} !Word64+  }+  deriving+    ( -- | @since 1.2.6.0+      Eq,+      -- | @since 1.2.6.0+      Show+    )++-- TODO: add unasfePerformIO?+-- TODO: remove NOINLINE?++-- | \(O(1)\) Static, shared storage of `Montgomery64`.+--+-- ==== Constraints+-- - \(m \le 2^{62})+-- - \(m\) is odd+--+-- @since 1.2.6.0+{-# NOINLINE new #-}+new :: forall a. (KnownNat a) => Proxy# a -> Montgomery64+-- FIXME: test allocated once+new p = fromVal . fromIntegral $! natVal' p++-- | \(O(1)\) Creates a `Montgomery64` for a modulus value \(m\) of type `Word64` value.+--+-- ==== Constraints+-- - \(m \le 2^{62})+-- - \(m\) is odd+--+-- @since 1.2.6.0+{-# INLINE fromVal #-}+fromVal :: Word64 -> Montgomery64+fromVal m =+  let !m128 :: Word128 = fromIntegral m+      !n2 = word128Lo64 $ (-m128) `mod` m128+      !r = getR m 0+      !_ = ACIA.runtimeAssert (r * m == -1) "AtCoder.Extra.Montgomery64.fromVal: internal implementation error"+   in Montgomery64 m r n2+  where+    !_ = ACIA.runtimeAssert (odd m && m <= bit 62) $ "AtCoder.Extra.Montgomery64.fromVal: not given odd modulus value that is less than or equal to 2^62: " ++ show m+    getR :: Word64 -> Int -> Word64+    getR !acc i+      | i >= 5 = -acc+      | otherwise = getR (acc * (2 - m * acc)) (i + 1)++-- | \(O(1)\) Retrieves the modulus \(m\).+--+-- @since 1.2.6.0+{-# INLINE umod #-}+umod :: Montgomery64 -> Word64+umod Montgomery64 {mM64} = mM64++-- | \(O(1)\) Converts the given `Word64` to Montgomery form.+--+-- @since 1.2.6.0+{-# INLINE encode #-}+encode :: Montgomery64 -> Word64 -> Word64+encode mont@Montgomery64 {n2M64} x = reduce mont $! fromIntegral x * fromIntegral n2M64++-- | \(O(1)\) Retrieves the value from a Montgomery form of value.+--+-- @since 1.2.6.0+{-# INLINE decode #-}+decode :: Montgomery64 -> Word64 -> Word64+decode mont@Montgomery64 {mM64} x =+  let !res = reduce mont $! fromIntegral x+   in if res >= mM64 then res - mM64 else res++-- | \(O(1)\) Takes the mod in Montgomery form.+--+-- @since 1.2.6.0+{-# INLINE reduce #-}+reduce :: Montgomery64 -> Word128 -> Word64+reduce Montgomery64 {mM64, rM64} x =+  word128Hi64 $!+    (x + fromIntegral (word128Lo64 x * rM64) * fromIntegral mM64)++-- | \(O(1)\) Calculates \(a + b \bmod m\) in the Montgomery form.+{-# INLINE addMod #-}+addMod :: Word64 -> Word64 -> Word64 -> Word64+addMod m a b+    | x' >= m = x' - m+    | otherwise = x'+  where+    !x' = a + b++-- | \(O(1)\) Calculates \(a - b \bmod m\) in the Montgomery form.+{-# INLINE subMod #-}+subMod :: Word64 -> Word64 -> Word64 -> Word64+subMod m a b+    | a >= b = a - b+    | otherwise = a - b + m++-- | \(O(1)\) Calculates \(a^n \bmod m\) in the Montgomery form.+--+-- @since 1.2.6.0+{-# INLINE mulMod #-}+mulMod :: Montgomery64 -> Word64 -> Word64 -> Word64+mulMod mont a b = reduce mont $! fromIntegral a * fromIntegral b++-- | \(O(w)\) Calculates \(a^n \bmod m\) in the Montgomery form.+--+-- @since 1.2.6.0+{-# INLINE powMod #-}+powMod :: (HasCallStack) => Montgomery64 -> Word64 -> Int -> Word64+powMod mont x0 n0 = inner n0 (encode mont 1) x0+  where+    !_ = ACIA.runtimeAssert (0 <= n0) $ "AtCoder.Extra.Math.Montgomery64.powMod: given negative exponential `n`: " ++ show n0 ++ show "`"+    inner :: Int -> Word64 -> Word64 -> Word64+    inner !n !r !y+      | n == 0 = r+      | otherwise =+          let !r' = if odd n then mulMod mont r y else r+              !y' = mulMod mont y y+           in inner (n !>>. 1) r' y'++-- | \(O(1)\) Compares two values of Montgomery form and returns whether they represent the same+-- value.+--+-- @since 1.2.6.0+{-# INLINE eq #-}+eq :: Word64 -> Word64 -> Word64 -> Bool+eq mM64 a b = a' == b'+  where+    !a' = if a < mM64 then a else a - mM64+    !b' = if b < mM64 then b else b - mM64+

@@ -0,0 +1,259 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE TypeFamilies #-}++-- | @ModInt@ for 64 bit modulus values.+--+-- @since 1.2.6.0+module AtCoder.Extra.ModInt64+  ( -- * ModInt64+    ModInt64 (..),++    -- * Constructors++    -- ** Safe constructors+    new,+    new64,++    -- ** Unsafe constructor+    unsafeNew,++    -- * Accessors++    -- ** Modulus value+    modulus,++    -- ** Internal value+    val,+    val64,++    -- * Operators+    pow,+    inv,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Extra.Math.Montgomery64 qualified as M64+import Data.Ratio (denominator, numerator)+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Primitive qualified as P+import Data.Vector.Unboxed qualified as U+import Data.Vector.Unboxed qualified as VU+import Data.Word (Word64)+import GHC.Exts (proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal')++-- | `Word64` value that treats the modular arithmetic.+--+-- @since 1.2.6.0+newtype ModInt64 a = ModInt64+  { -- | Montgomery form of the value. Use `val` to retrieve the value.+    --+    -- @since 1.2.6.0+    unModInt64 :: Word64+  }+  deriving+    ( -- | @since 1.2.6.0+      P.Prim+    )++-- | @since 1.2.6.0+instance (KnownNat a) => Eq (ModInt64 a) where+  {-# INLINE (==) #-}+  ModInt64 x == ModInt64 y = M64.eq (fromIntegral (natVal' (proxy# @a))) x y++-- | @since 1.2.6.0+instance (KnownNat a) => Ord (ModInt64 a) where+  {-# INLINE compare #-}+  compare (ModInt64 a) (ModInt64 b) = compare a b++-- | @since 1.2.6.0+instance (KnownNat a) => Read (ModInt64 a) where+  {-# INLINE readsPrec #-}+  readsPrec p s = [(fromInteger x, r) | (!x, !r) <- readsPrec p s]++-- | @since 1.2.6.0+instance (KnownNat a) => Show (ModInt64 a) where+  {-# INLINE show #-}+  show = show . val++-- | \(O(1)\) Creates a `ModInt64` from an `Int` value taking the mod.+--+-- @since 1.2.6.0+{-# INLINE new #-}+new :: forall a. (KnownNat a) => Int -> ModInt64 a+new = ModInt64 . M64.encode (M64.new (proxy# @a)) . fromIntegral . (`mod` m)+  where+    !m = fromIntegral $ natVal' (proxy# @a)++-- | \(O(1)\) Creates a `ModInt64` from a `Word64` value taking the mod.+--+-- @since 1.2.6.0+{-# INLINE new64 #-}+new64 :: forall a. (KnownNat a) => Word64 -> ModInt64 a+new64 = ModInt64 . M64.encode (M64.new (proxy# @a))++-- | \(O(1)\) Creates `ModInt64` from a Montgomery form with no validation.+--+-- @since 1.2.6.0+{-# INLINE unsafeNew #-}+unsafeNew :: (KnownNat a) => Word64 -> ModInt64 a+unsafeNew = ModInt64++-- | \(O(1)\) Retrieve the mod from a `ModInt64` object.+--+-- ==== Complecity+-- - \(O(1)\)+--+-- @since 1.2.6.0+{-# INLINE modulus #-}+modulus :: forall a. (KnownNat a) => ModInt64 a -> Int+modulus _ = fromIntegral (natVal' (proxy# @a))++-- | \(O(1)\) Returns the internal value in `Int`.+--+-- ==== Complecity+-- - \(O(1)\)+--+-- @since 1.2.6.0+{-# INLINE val #-}+val :: forall a. (KnownNat a) => ModInt64 a -> Int+val = fromIntegral . val64++-- | \(O(1)\) Returns the internal value in `Word64`.+--+-- ==== Complecity+-- - \(O(1)\)+--+-- @since 1.2.6.0+{-# INLINE val64 #-}+val64 :: forall a. (KnownNat a) => ModInt64 a -> Word64+val64 (ModInt64 x) = M64.decode (M64.new (proxy# @a)) x++-- | \(O(\log n\) Returns \(x^n\). The implementation is a bit more efficient than `^`.+--+-- ==== Constraints+-- - \(0 \le n\)+--+-- @since 1.2.6.0+{-# INLINE pow #-}+pow :: forall a. (HasCallStack, KnownNat a) => ModInt64 a -> Int -> ModInt64 a+pow (ModInt64 x) n = ModInt64 $! M64.powMod (M64.new (proxy# @a)) x n++-- TODO: move invMod to Montgomery64+-- TODO: time complexity of `inv`?++-- | Returns \(y\) such that \(xy \equiv 1\) holds.+--+-- ==== Constraints+-- - The value must not be zero.+--+-- @since 1.2.6.0+{-# INLINE inv #-}+inv :: forall a. (HasCallStack, KnownNat a) => ModInt64 a -> ModInt64 a+-- TODO: assert zero division?+inv self = inner (val self) m 1 0+  where+    !_ = ACIA.runtimeAssert (val self /= 0) "AtCoder.Extra.ModInt64.inv: given zero"+    !m = fromIntegral (natVal' (proxy# @a))+    inner x y u v+      | y <= 0 = new u+      | otherwise = inner x' y' u' v'+      where+        x' = y+        y' = x - t * y+        u' = v+        v' = u - t * v+        t = x `div` y++-- https://github.com/NyaanNyaan/library/blob/master/modint/montgomery-modint.hpp+-- constexpr mint inverse() const {+--   int x = get(), y = mod, u = 1, v = 0, t = 0, tmp = 0;+--   while (y > 0) {+--     t = x / y;+--     x -= t * y, u -= t * v;+--     tmp = x, x = y, y = tmp;+--     tmp = u, u = v, v = tmp;+--   }+--   return mint{u};+-- }++-- | @since 1.2.6.0+deriving newtype instance (KnownNat p) => Real (ModInt64 p)++-- | @since 1.2.6.0+instance forall p. (KnownNat p) => Num (ModInt64 p) where+  {-# INLINE (+) #-}+  (ModInt64 !x1) + (ModInt64 !x2) = ModInt64 $! M64.addMod m x1 x2+    where+      !m = fromIntegral (natVal' (proxy# @p))+  {-# INLINE (-) #-}+  (ModInt64 !x1) - (ModInt64 !x2) = ModInt64 $! M64.subMod m x1 x2+    where+      !m = fromIntegral (natVal' (proxy# @p))+  {-# INLINE (*) #-}+  (ModInt64 !x1) * (ModInt64 !x2) = ModInt64 $! M64.mulMod (M64.new (proxy# @p)) x1 x2+  {-# INLINE negate #-}+  negate x = 0 - x+  {-# INLINE abs #-}+  abs = id+  {-# INLINE signum #-}+  signum _ = ModInt64 $ M64.encode (M64.new (proxy# @p)) 1+  -- because the input value can be negative, be sure to take the mod:+  {-# INLINE fromInteger #-}+  fromInteger = ModInt64 . M64.encode (M64.new (proxy# @p)) . fromInteger . (`mod` m)+    where+      !m = toInteger $ natVal' (proxy# @p)++-- | @since 1.2.6.0+instance (KnownNat p) => Bounded (ModInt64 p) where+  {-# INLINE minBound #-}+  minBound = ModInt64 0+  {-# INLINE maxBound #-}+  maxBound = ModInt64 . M64.encode (M64.new (proxy# @p)) $! fromIntegral (natVal' (proxy# @p)) - 1++-- | @since 1.2.6.0+instance (KnownNat p) => Enum (ModInt64 p) where+  {-# INLINE toEnum #-}+  toEnum = new+  {-# INLINE fromEnum #-}+  fromEnum = fromIntegral . val++-- | @since 1.2.6.0+instance (KnownNat p) => Integral (ModInt64 p) where+  -- FIXME: THIS IS COMPLETELY WRONG. Compare with `ModInt`.+  {-# INLINE quotRem #-}+  -- quotRem x y =+  --   let !x' = val x+  --       !y' = val y+  --       (!q, !r) = x' `quotRem` y'+  --    in (new q, new r)+  quotRem x y = (x / y, x - x / y * y)+  {-# INLINE toInteger #-}+  toInteger = toInteger . val++-- | @since 1.2.6.0+instance (KnownNat p) => Fractional (ModInt64 p) where+  {-# INLINE recip #-}+  recip = inv+  {-# INLINE fromRational #-}+  fromRational q = fromInteger (numerator q) / fromInteger (denominator q)++-- | @since 1.2.6.0+newtype instance VU.MVector s (ModInt64 a) = MV_ModInt64 (VU.MVector s Word64)++-- | @since 1.2.6.0+newtype instance VU.Vector (ModInt64 a) = V_ModInt64 (VU.Vector Word64)++-- | @since 1.2.6.0+deriving newtype instance VGM.MVector VU.MVector (ModInt64 a)++-- | @since 1.2.6.0+deriving newtype instance VG.Vector VU.Vector (ModInt64 a)++-- | @since 1.2.6.0+instance VU.Unbox (ModInt64 a)

@@ -30,7 +30,7 @@ import Data.Foldable (for_) import Data.Vector.Unboxed qualified as VU --- INLINE all the functions, even if the performance gain is just a little bit.+-- INLINE all the functions, even if the performance gain is just a little bit, in case it matters.  -- | \(O(\sqrt n)\) Runs user function for each block. {-# INLINE forM_ #-}@@ -67,10 +67,14 @@  -- | \(O(\sqrt n)\) Runs user function for each block and concatanate their monoid output. --+-- ==== Constraints+-- - \(l \le r\)+-- - If an empty interval is queried, the @readPart@ function must return a valid value.+-- -- @since 1.2.5.0 {-# INLINE foldMapM #-} foldMapM ::-  (Monad m, Monoid a) =>+  (Monad m, Semigroup a) =>   -- | Context: block length.   Int ->   -- | Function: @readFull@ function that takes target block index and returns monoid value of it.@@ -89,6 +93,10 @@ -- | \(O(\sqrt n)\) Runs user function for each block and concatanates their output with user -- function. --+-- ==== Constraints+-- - \(l \le r\)+-- - If an empty interval is queried, the @readPart@ function must return a valid value.+-- -- @since 1.2.5.0 {-# INLINE foldMapWithM #-} foldMapWithM ::@@ -133,6 +141,9 @@  -- | \(O(\sqrt n)\) Runs user function for each block, performing left folding. --+-- ==== Constraints+-- - \(l \le r\)+-- -- @since 1.2.5.0 {-# INLINE foldM #-} foldM ::@@ -158,7 +169,9 @@   let (!ir, !remR) = r `divMod` blockLen   if il == ir     then do-      foldPart s0 il l r+      if remL == remR+        then pure s0+        else foldPart s0 il l r     else do       !sx <-         if remL == 0@@ -174,6 +187,9 @@         else foldPart sm ir (r - remR) r  -- | \(O(\sqrt n)\) `foldM` with return value discarded.+--+-- ==== Constraints+-- - \(l \le r\) -- -- @since 1.2.5.0 {-# INLINE foldM_ #-}

@@ -79,7 +79,7 @@ import GHC.Stack (HasCallStack) import GHC.TypeNats (KnownNat, natVal, natVal') --- | `KnownNat` with meta information used for modulus.+-- | `KnownNat` with meta information as a modulus value for convolution. -- -- @since 1.0.0.0 class (KnownNat a) => Modulus a where@@ -244,7 +244,7 @@ modulus :: forall a. (KnownNat a) => ModInt a -> Int modulus _ = fromIntegral (natVal' (proxy# @a)) --- | Returns the internal value converted to `Int`.+-- | Returns the internal value in `Int`. -- -- ==== Complecity -- - \(O(1)\)@@ -265,7 +265,7 @@ val32 :: (KnownNat a) => ModInt a -> Word32 val32 = unModInt --- | Returns the internal value converted to `Word32`.+-- | Returns the internal value in `Word32`. -- -- ==== Complecity -- - \(O(1)\)

@@ -20,6 +20,8 @@ import Tests.Extra.KdTree qualified import Tests.Extra.LazyKdTree qualified import Tests.Extra.Math qualified+import Tests.Extra.Math.Montgomery64 qualified+import Tests.Extra.ModInt64 qualified import Tests.Extra.Monoid qualified import Tests.Extra.MultiSet qualified import Tests.Extra.SegTree2d qualified@@ -75,6 +77,8 @@             testGroup "KdTree" Tests.Extra.KdTree.tests,             testGroup "LazyKdTree" Tests.Extra.LazyKdTree.tests,             testGroup "Math" Tests.Extra.Math.tests,+            testGroup "Math.Montgomery64" Tests.Extra.Math.Montgomery64.tests,+            testGroup "ModInt64" Tests.Extra.ModInt64.tests,             testGroup "Monoid" Tests.Extra.Monoid.tests,             testGroup "MultiSet" Tests.Extra.MultiSet.tests,             testGroup "SegTree2d" Tests.Extra.SegTree2d.tests,

@@ -1,10 +1,14 @@ module Tests.Extra.Math (tests) where  import AtCoder.Extra.Math qualified as ACEM+import Data.Foldable (for_)+import Data.List qualified as L import Data.Proxy (Proxy (..)) import Data.Semigroup (Max (..), Min (..), Sum (..), mtimesDefault, stimes)+import Data.Vector.Unboxed qualified as VU import Test.QuickCheck.Property qualified as QC import Test.Tasty+import Test.Tasty.HUnit import Test.Tasty.QuickCheck qualified as QC import Tests.Util (myForAllShrink) @@ -36,6 +40,84 @@     rhsS = "mtimes n s"     rhs (QC.NonNegative !n, !m) = mtimesDefault n m +-- | This is a solid, fast implementation of prime number enumeration for lists.+truePrimes :: [Int]+truePrimes = 2 : 3 : minus [5, 7 ..] (unionAll [[p * p, p * p + 2 * p ..] | p <- tail truePrimes])+  where+    minus (x : xs) (y : ys) = case compare x y of+      LT -> x : minus xs (y : ys)+      EQ -> minus xs ys+      GT -> minus (x : xs) ys+    minus xs _ = xs++    union (x : xs) (y : ys) = case compare x y of+      LT -> x : union xs (y : ys)+      EQ -> x : union xs ys+      GT -> y : union (x : xs) ys+    union xs [] = xs+    union [] ys = ys++    unionAll :: (Ord a) => [[a]] -> [a]+    unionAll ((x : xs) : t) = x : union xs (unionAll $ pairs t)+      where+        pairs ((x : xs) : ys : t) = (x : union xs ys) : pairs t+        pairs _ = error "unionAll _ pairs: unreachable"+    unionAll _ = error "unionAll: unreachable"++-- | This is a solid, fast implementation of prime number enumeration for lists.+truePrimeFactors :: Int -> [(Int, Int)]+truePrimeFactors !n_ = map (\ !xs -> (head xs, length xs)) . L.group $ inner n_ input+  where+    -- TODO: reuse `primes`?+    input = 2 : 3 : [y | x <- [5, 11 ..], y <- [x, x + 2]]+    inner n pps@(p : ps)+      | n == 1 = []+      | n < p * p = [n]+      | r == 0 = p : inner q pps+      | otherwise = inner n ps+      where+        (q, r) = divMod n p+    inner _ _ = error "unreachable"++trueDivisors :: Int -> [Int]+trueDivisors n = L.sort $ inner 1+  where+    inner k+      -- no dependency to `Int` square root function+      | k * k > n = []+      -- no divisor duplication+      | k * k == n = [k]+      -- not sorted yet+      | r == 0 = k : d : inner (succ k)+      -- ignore non divisors+      | otherwise = inner (succ k)+      where+        -- This strict evaluation and unboxing takes some effect, even though they're not always+        -- used.+        (!d, !r) = n `divMod` k++-- unit_primes :: TestTree+-- unit_primes = testCase "primes" $ do+--   for_ [0 .. 10 ^ 9] $ \upper -> do+--     when (upper `mod` 10000 == 0) $ do+--       let !_ = traceShow upper ()+--       pure ()+--     let expected = VU.fromList $ takeWhile (<= upper) truePrimes+--     let result = ACEM.primes upper+--     result @?= expected++prop_primeFactors :: QC.Positive Int -> QC.Property+prop_primeFactors (QC.Positive x) =+  let expected = VU.fromList $ truePrimeFactors x+      result = ACEM.primeFactors x+   in result QC.=== expected++prop_divisors :: QC.Positive Int -> QC.Property+prop_divisors (QC.Positive x) =+  let expected = VU.fromList $ trueDivisors x+      result = ACEM.divisors x+   in result QC.=== expected+ tests :: [TestTree] tests =   [ testGroup@@ -51,5 +133,13 @@         QC.testProperty "Product" (prop_mtimes' (Proxy @(Sum Int))),         QC.testProperty "Max" (prop_mtimes' (Proxy @(Max Int))),         QC.testProperty "Min" (prop_mtimes' (Proxy @(Min Int)))+      ],+    testGroup+      "primes"+      [ -- unit_primes+        QC.testProperty "primeFactors" prop_primeFactors,+        QC.testProperty "divisors" prop_divisors       ]   ]++-- unit_primes

@@ -0,0 +1,44 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE TypeFamilies #-}++module Tests.Extra.Math.Montgomery64 (tests) where++import AtCoder.Extra.Math qualified as ACEM+import AtCoder.Extra.Math.Montgomery64 qualified as M+import Data.Bits+import Data.WideWord.Word128 (Word128 (..))+import Data.Word (Word64)+import Test.Tasty+import Test.Tasty.QuickCheck qualified as QC++to128 :: (Integral a) => a -> Word128+to128 = fromIntegral++mulMod :: Word64 -> Word64 -> Word64 -> Word64+mulMod m x y = word128Lo64 $! (to128 (x `mod` m) * to128 (y `mod` m)) `mod` to128 m++p :: Word64 -> Bool+p m = odd m && m <= bit 62++prop_mulMod :: QC.Positive Word64 -> Word64 -> Word64 -> QC.Property+prop_mulMod (QC.Positive m) x y =+  p m QC.==>+    let !mont = M.fromVal m+        !res = M.decode mont $ M.mulMod mont (M.encode mont x) (M.encode mont y)+        !expected = mulMod m x y+     in res QC.=== expected++prop_powMod :: QC.Positive Word64 -> Word64 -> QC.Positive Int -> QC.Property+prop_powMod (QC.Positive m) x (QC.Positive n) =+  p m QC.==>+    let !mont = M.fromVal m+        !res = M.decode mont $ M.powMod mont (M.encode mont x) n+        !expected = ACEM.power (mulMod m) n (x `mod` m)+     in res QC.=== expected++tests :: [TestTree]+tests =+  [ QC.testProperty "mulMod" prop_mulMod,+    QC.testProperty "powMod" prop_powMod+  ]

@@ -0,0 +1,219 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE TypeFamilies #-}++module Tests.Extra.ModInt64 (tests) where++import AtCoder.Extra.Math qualified as ACEM+import AtCoder.Extra.ModInt64 qualified as M+import AtCoder.ModInt qualified as M32+import Data.Semiring (Ring (..), Semiring (..), WrappedNum (..))+import Data.WideWord.Word128 (Word128 (..))+import Data.Word (Word64 (..))+import GHC.Exts (Proxy#, proxy#)+import GHC.TypeNats (KnownNat, natVal')+import Test.QuickCheck.Classes qualified as QCC+import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck qualified as QC+import Tests.Util (laws)++deriving via (WrappedNum Word64) instance Semiring (M.ModInt64 a)++deriving via (WrappedNum Word64) instance Ring (M.ModInt64 a)++type M1 = 3++type M2 = 5++type M3 = 998244353++type M4 = 1000000007++type M5 = 4611686018427387847++instance M32.Modulus M1 where+  {-# INLINE isPrimeModulus #-}+  isPrimeModulus _ = True++  -- FIXME: wrong+  primitiveRootModulus _ = (-1)++instance M32.Modulus M2 where+  {-# INLINE isPrimeModulus #-}+  isPrimeModulus _ = True++  -- FIXME: wrong+  primitiveRootModulus _ = (-1)++instance (KnownNat a) => QC.Arbitrary (M.ModInt64 a) where+  arbitrary = M.new <$> QC.arbitrary++to128 :: (Integral a) => a -> Word128+to128 = fromIntegral++mulMod :: Int -> Int -> Int -> Int+mulMod m x y = fromIntegral . word128Lo64 $! (to128 (x `mod` m) * to128 (y `mod` m)) `mod` to128 m++unit_literal :: forall a. (KnownNat a) => Proxy# a -> TestTree+unit_literal proxy = testCase "literal" $ do+  let !m :: Int = fromIntegral $ natVal' proxy+  (@?= (0 `mod` m)) $ fromIntegral (0 :: M.ModInt64 a)+  (@?= ((-1) `mod` m)) $ fromIntegral (-1 :: M.ModInt64 a)+  (@?= (1 `mod` m)) $ fromIntegral (1 :: M.ModInt64 a)+  (@?= (m `mod` m)) $ fromIntegral (M.new @a m)+  (@?= ((m - 1) `mod` m)) $ fromIntegral (M.new @a (m - 1))+  (@?= ((m + 1) `mod` m)) $ fromIntegral (M.new @a (m + 1))++prop_addMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property+prop_addMod proxy x y =+  let !m = fromIntegral $ natVal' proxy+      !res = M.val $ M.new @a x + M.new @a y+      !expected = (x + y) `mod` m+   in res QC.=== expected++prop_subMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property+prop_subMod proxy x y =+  let !m = fromIntegral $ natVal' proxy+      !res = M.val $ M.new @a x - M.new @a y+      !expected = (x - y) `mod` m+   in res QC.=== expected++prop_mulMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property+prop_mulMod proxy x y =+  let !m = fromIntegral $ natVal' proxy+      !res = M.val $ M.new @a x * M.new @a y+      !expected = mulMod m x y+   in res QC.=== expected++prop_powMod :: forall a. (KnownNat a) => Proxy# a -> Int -> QC.Positive Int -> QC.Property+prop_powMod proxy x (QC.Positive n) =+  let !m = fromIntegral $ natVal' proxy+      !res = M.val $ M.pow (M.new @a x) n+      !expected = ACEM.power (mulMod m) n (x `mod` m)+   in res QC.=== expected++prop_inv :: forall a. (KnownNat a) => M.ModInt64 a -> QC.Property+prop_inv x =+  M.val x+    /= 0+    QC.==> QC.counterexample (show x)+    $ QC.conjoin+      [ M.inv x * x QC.=== M.new 1,+        M.new 1 QC.=== M.inv x * x+      ]++prop_quotRem :: forall a. (M32.Modulus a) => Proxy# a -> Int -> QC.NonZero Int -> QC.Property+prop_quotRem _ x (QC.NonZero y) =+  (y `mod` m /= 0) QC.==>+    let (!resQ, !resR) = M.new @a x `quotRem` M.new @a y+        (!expQ, !expR) = M32.new @a x `quotRem` M32.new @a y+     in QC.conjoin+          [ M.val resQ QC.=== M32.val expQ,+            M.val resR QC.=== M32.val expR+          ]+  where+    !m = fromIntegral $ natVal' (proxy# @a)++prop_eq :: forall a. (KnownNat a) => Proxy# a -> Word64 -> Word64 -> QC.Property+prop_eq _ x y = lhs QC.=== rhs+  where+    !lhs = M.new64 @a x == M.new64 @a y+    !m = fromIntegral $ natVal' (proxy# @a)+    !rhs = x `mod` m == y `mod` m++-- Cannot create list for unlifted types+--+-- {-# LANGUAGE ImpredicativeTypes #-}+-- modProps :: String -> [forall a. (KnownNat a) => Proxy# a -> QC.Property] -> TestTree+-- modProps title prop =+--   testGroup title $+--     map+--       (\proxy -> QC.testProperty (show (natVal' proxy)) (prop proxy))+--       [proxy# @M1, proxy# @M2, proxy# @M3, proxy# @M4, proxy# @M5]++tests :: [TestTree]+tests =+  [ testGroup+      "literal"+      [ unit_literal (proxy# @M1),+        unit_literal (proxy# @M2),+        unit_literal (proxy# @M3),+        unit_literal (proxy# @M4),+        unit_literal (proxy# @M5)+      ],+    testGroup+      "inv"+      [ QC.testProperty "1" (prop_inv @M1),+        QC.testProperty "2" (prop_inv @M2),+        QC.testProperty "3" (prop_inv @M3),+        QC.testProperty "4" (prop_inv @M4),+        QC.testProperty "5" (prop_inv @M5)+      ],+    testGroup+      "quotRem"+      [ QC.testProperty "1" (prop_quotRem (proxy# @M1)),+        QC.testProperty "2" (prop_quotRem (proxy# @M2)),+        QC.testProperty "3" (prop_quotRem (proxy# @M3)),+        QC.testProperty "4" (prop_quotRem (proxy# @M4))+        -- 64 bit+        -- QC.testProperty "5" (prop_quotRem (proxy# @M5))+      ],+    testGroup+      "eq"+      [ QC.testProperty "1" (prop_eq (proxy# @M1)),+        QC.testProperty "2" (prop_eq (proxy# @M2)),+        QC.testProperty "3" (prop_eq (proxy# @M3)),+        QC.testProperty "4" (prop_eq (proxy# @M4)),+        QC.testProperty "5" (prop_eq (proxy# @M5))+      ],+    testGroup+      "addMod"+      [ QC.testProperty "1" (prop_addMod (proxy# @M1)),+        QC.testProperty "2" (prop_addMod (proxy# @M2)),+        QC.testProperty "3" (prop_addMod (proxy# @M3)),+        QC.testProperty "4" (prop_addMod (proxy# @M4)),+        QC.testProperty "5" (prop_addMod (proxy# @M5))+      ],+    testGroup+      "subMod"+      [ QC.testProperty "1" (prop_subMod (proxy# @M1)),+        QC.testProperty "2" (prop_subMod (proxy# @M2)),+        QC.testProperty "3" (prop_subMod (proxy# @M3)),+        QC.testProperty "4" (prop_subMod (proxy# @M4)),+        QC.testProperty "5" (prop_subMod (proxy# @M5))+      ],+    testGroup+      "mulMod"+      [ QC.testProperty "1" (prop_mulMod (proxy# @M1)),+        QC.testProperty "2" (prop_mulMod (proxy# @M2)),+        QC.testProperty "3" (prop_mulMod (proxy# @M3)),+        QC.testProperty "4" (prop_mulMod (proxy# @M4)),+        QC.testProperty "5" (prop_mulMod (proxy# @M5))+      ],+    testGroup+      "powMod"+      [ QC.testProperty "1" (prop_powMod (proxy# @M1)),+        QC.testProperty "2" (prop_powMod (proxy# @M2)),+        QC.testProperty "3" (prop_powMod (proxy# @M3)),+        QC.testProperty "4" (prop_powMod (proxy# @M4)),+        QC.testProperty "5" (prop_powMod (proxy# @M5))+      ],+    testGroup+      "laws"+      [ laws @(M.ModInt64 M5)+          [ QCC.eqLaws,+            QCC.numLaws,+            QCC.integralLaws,+            QCC.ordLaws,+            QCC.enumLaws,+            QCC.boundedEnumLaws,+            QCC.primLaws,+            QCC.semiringLaws,+            QCC.ringLaws,+            QCC.showReadLaws+          ]+      ]+  ]