ac-library-hs-1.2.2.0: test/Tests/Convolution.hs
{-# LANGUAGE DataKinds #-}
module Tests.Convolution (tests) where
import AtCoder.Convolution qualified as ACC
import AtCoder.Internal.Convolution qualified as ACIC
import AtCoder.ModInt qualified as AM
import Data.Proxy (Proxy (..))
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 Test.QuickCheck qualified as QC
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck qualified as QC
-- | Orphan instance..
instance AM.Modulus 924844033 where
isPrimeModulus _ = True
primitiveRootModulus _ = 5
convModNaive :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int
convModNaive modulus a b = VU.create $ do
c <- VUM.replicate (VU.length a + VU.length b - 1) 0
VU.iforM_ a $ \i ai -> do
VU.iforM_ b $ \j bj -> do
VGM.modify c ((`mod` modulus) . (+ ai * bj)) (i + j)
pure c
convMintNaive :: (AM.Modulus p) => VU.Vector (AM.ModInt p) -> VU.Vector (AM.ModInt p) -> VU.Vector (AM.ModInt p)
convMintNaive a b = VU.create $ do
c <- VUM.replicate (VU.length a + VU.length b - 1) 0
VU.iforM_ a $ \i ai -> do
VU.iforM_ b $ \j bj -> do
VGM.modify c (+ ai * bj) (i + j)
pure c
conv64Naive :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
conv64Naive a b = VU.create $ do
c <- VUM.replicate (VU.length a + VU.length b - 1) 0
VU.iforM_ a $ \i ai -> do
VU.iforM_ b $ \j bj -> do
VGM.modify c (+ ai * bj) (i + j)
pure c
unit_empty :: TestTree
unit_empty = testCase "empty" $ do
-- any `Integral a` is allowed
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @Int) (VU.empty @Int)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @Int) (VU.fromList @Int [1, 2])
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @Int [1, 2]) (VU.empty @Int)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @Int [1]) (VU.empty @Int)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @Word64) (VU.empty @Word64)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @Word64) (VU.fromList @Word64 [1, 2])
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @Word64 [1, 2]) (VU.empty @Word64)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @Word64 [1]) (VU.empty @Word64)
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @(AM.ModInt 998244353)) (VU.empty @(AM.ModInt 998244353))
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.empty @(AM.ModInt 998244353)) (VU.fromList @(AM.ModInt 998244353) [1, 2])
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @(AM.ModInt 998244353) [1, 2]) (VU.empty @(AM.ModInt 998244353))
VU.empty @=? ACC.convolutionRaw (Proxy @998244353) (VU.fromList @(AM.ModInt 998244353) [1]) (VU.empty @(AM.ModInt 998244353))
-- FIXME: the naive calculation seems like too slow
prop_mid :: TestTree
prop_mid = QC.testProperty "mid" $ do
-- any `Integral a` is allowed
let modInt :: Int -> AM.ModInt998244353
modInt = AM.new
a <- VU.map modInt . VU.fromList <$> QC.vectorOf 1234 (QC.arbitrary @Int)
b <- VU.map modInt . VU.fromList <$> QC.vectorOf 2345 (QC.arbitrary @Int)
pure $ convMintNaive a b QC.=== ACC.convolution a b
unit_butterfly :: TestTree
unit_butterfly = testCase "butterfly" $ do
let expected = VU.fromList [10, 998244351, 173167434, 825076915]
let xs = VU.create $ do
let modInt :: Int -> AM.ModInt998244353
modInt = AM.new
vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4]
info <- ACIC.newInfo @_ @998244353
ACIC.butterfly info vec
pure vec
expected @?= xs
unit_invButterfly :: TestTree
unit_invButterfly = testCase "invButterfly" $ do
let expected = VU.fromList [10, 911660634, 998244349, 86583717]
let xs = VU.create $ do
let modInt :: Int -> AM.ModInt998244353
modInt = AM.new
vec <- VU.unsafeThaw $ VU.map modInt $ VU.fromList [1, 2, 3, 4]
info <- ACIC.newInfo @_ @998244353
ACIC.butterflyInv info vec
pure vec
expected @=? xs
testWithRangeMint ::
forall p.
(AM.Modulus p) =>
(VU.Vector (AM.ModInt p) -> VU.Vector (AM.ModInt p) -> QC.Property) ->
QC.Gen QC.Property
testWithRangeMint f = do
m <- QC.chooseInt (1, 20 - 1)
n <- QC.chooseInt (1, 20 - 1)
a <- VU.map (AM.new @p) . VU.fromList <$> QC.vectorOf n (QC.arbitrary @Int)
b <- VU.map (AM.new @p) . VU.fromList <$> QC.vectorOf m (QC.arbitrary @Int)
pure $ f a b
testWithRangeInt :: forall p. (AM.Modulus p) => Proxy p -> (VU.Vector Int -> VU.Vector Int -> QC.Property) -> QC.Gen QC.Property
testWithRangeInt _ f = do
m <- QC.chooseInt (1, 20 - 1)
n <- QC.chooseInt (1, 20 - 1)
let !modulus = AM.modVal (Proxy @p)
a <- VU.map (`mod` modulus) . VU.fromList <$> QC.vectorOf n (QC.arbitrary @Int)
b <- VU.map (`mod` modulus) . VU.fromList <$> QC.vectorOf m (QC.arbitrary @Int)
pure $ f a b
prop_simpleSMod1 :: TestTree
prop_simpleSMod1 = QC.testProperty "simpleSMod1"
. testWithRangeMint @998244353
$ \a b -> convMintNaive a b QC.=== ACC.convolution a b
prop_simpleSMod2 :: TestTree
prop_simpleSMod2 = QC.testProperty "simpleSMod1"
. testWithRangeMint @924844033
$ \a b -> convMintNaive a b QC.=== ACC.convolution a b
prop_simpleInt0 :: TestTree
prop_simpleInt0 = QC.testProperty "simpleInt0"
. testWithRangeInt (Proxy @998244353)
$ \a b -> convModNaive modulus a b QC.=== VU.map AM.val (ACC.convolution (VU.map modInt a) (VU.map modInt b))
where
modInt = AM.new @998244353
!modulus = AM.modVal (Proxy @998244353)
prop_simpleInt1 :: TestTree
prop_simpleInt1 = QC.testProperty "simpleInt1"
. testWithRangeInt (Proxy @998244353)
$ \a b -> convModNaive modulus a b QC.=== ACC.convolutionRaw (Proxy @998244353) a b
where
!modulus = AM.modVal (Proxy @998244353)
prop_simpleInt2 :: TestTree
prop_simpleInt2 = QC.testProperty "simpleInt2"
. testWithRangeInt (Proxy @924844033)
$ \a b -> convModNaive modulus a b QC.=== ACC.convolutionRaw (Proxy @924844033) a b
where
!modulus = AM.modVal (Proxy @924844033)
testWithRange64 :: (VU.Vector Int -> VU.Vector Int -> QC.Property) -> QC.Gen QC.Property
testWithRange64 f = do
m <- QC.chooseInt (1, 20 - 1)
n <- QC.chooseInt (1, 20 - 1)
a <- VU.map (fromIntegral . subtract 500_000 . (`mod` 1_000_000)) . VU.fromList <$> QC.vectorOf n (QC.arbitrary @Word64)
b <- VU.map (fromIntegral . subtract 500_000 . (`mod` 1_000_000)) . VU.fromList <$> QC.vectorOf m (QC.arbitrary @Word64)
pure $ f a b
prop_conv64 :: TestTree
prop_conv64 = QC.testProperty "conv64"
. testWithRange64
$ \a b -> conv64Naive a b QC.=== ACC.convolution64 a b
-- TODO: Shall we use Word64? (#41)
prop_conv64Bound :: [TestTree]
prop_conv64Bound =
[ QC.testProperty "conv64Bound_1" $ do
i <- fromIntegral <$> QC.chooseInt (-1000, 1000)
let a = VU.singleton (minBound @Int - m1m2 - m1m3 - m2m3 + i)
let b = VU.singleton 1
pure $ a QC.=== ACC.convolution64 a b,
QC.testProperty "conv64Bound_2" $ do
i <- fromIntegral <$> QC.chooseInt (0, 1000 - 1)
let a = VU.singleton $ (0 :: Int) + i
let b = VU.singleton 1
pure $ a QC.=== ACC.convolution64 a b,
QC.testProperty "conv64Bound_3" $ do
i <- fromIntegral <$> QC.chooseInt (0, 1000 - 1)
let a = VU.singleton $ maxBound @Int + i
let b = VU.singleton 1
pure $ a QC.=== ACC.convolution64 a b
]
where
mod1 :: Int = 469762049
mod2 :: Int = 167772161
mod3 :: Int = 754974721
m2m3 :: Int = mod2 * mod3
m1m3 :: Int = mod1 * mod3
m1m2 :: Int = mod1 * mod2
instance AM.Modulus 641 where
isPrimeModulus _ = True
primitiveRootModulus _ = 3
prop_conv641 :: TestTree
prop_conv641 = QC.testProperty "prop_conv641" $ do
let modulus = 641
a <- VU.fromList <$> QC.vectorOf 64 (QC.chooseInt (0, modulus - 1))
b <- VU.fromList <$> QC.vectorOf 65 (QC.chooseInt (0, modulus - 1))
pure $ convModNaive modulus a b QC.=== ACC.convolutionRaw (Proxy @641) a b
instance AM.Modulus 18433 where
isPrimeModulus _ = True
primitiveRootModulus _ = 5
prop_conv18433 :: TestTree
prop_conv18433 = QC.testProperty "prop_conv18433" $ do
let modulus = 18433
a <- VU.fromList <$> QC.vectorOf 1024 (QC.chooseInt (0, modulus - 1))
b <- VU.fromList <$> QC.vectorOf 1025 (QC.chooseInt (0, modulus - 1))
pure $ convModNaive modulus a b QC.=== ACC.convolutionRaw (Proxy @18433) a b
instance AM.Modulus 2 where
isPrimeModulus _ = True
primitiveRootModulus _ = 1
unit_conv2 :: TestTree
unit_conv2 = testCase "prop_conv2" $ do
VU.empty @Int @=? ACC.convolutionRaw (Proxy @257) VU.empty VU.empty
instance AM.Modulus 257 where
isPrimeModulus _ = True
primitiveRootModulus _ = 3
prop_conv257 :: TestTree
prop_conv257 = QC.testProperty "prop_conv257" $ do
let modulus = 257
a <- VU.fromList <$> QC.vectorOf 128 (QC.chooseInt (0, modulus - 1))
b <- VU.fromList <$> QC.vectorOf 129 (QC.chooseInt (0, modulus - 1))
pure $ convModNaive modulus a b QC.=== ACC.convolutionRaw (Proxy @257) a b
prop_conv2147483647 :: TestTree
prop_conv2147483647 = QC.testProperty "prop_conv2147483647" $ do
a <- VU.map (AM.new @2147483647) . VU.fromList <$> QC.vectorOf 1 (QC.arbitrary @Int)
b <- VU.map (AM.new @2147483647) . VU.fromList <$> QC.vectorOf 2 (QC.arbitrary @Int)
pure $ convMintNaive a b QC.=== ACC.convolution a b
instance AM.Modulus 2130706433 where
isPrimeModulus _ = True
primitiveRootModulus _ = 3
prop_conv2130706433 :: TestTree
prop_conv2130706433 = QC.testProperty "prop_conv2130706433" $ do
a <- VU.map (AM.new @2130706433) . VU.fromList <$> QC.vectorOf 1024 (QC.arbitrary @Int)
b <- VU.map (AM.new @2130706433) . VU.fromList <$> QC.vectorOf 1024 (QC.arbitrary @Int)
pure $ convMintNaive a b QC.=== ACC.convolution a b
unit_bigPrime :: TestTree
unit_bigPrime = testCase "bigPrime" $ do
let a = VU.map (AM.new @2130706433) $ VU.fromList [1, 2130706432, 1, 0, 2130706432]
let b = VU.map (AM.new @2130706433) $ VU.fromList [1, 1, 2130706432, 0, 0]
convMintNaive a b @=? ACIC.convolutionFft a b
tests :: [TestTree]
tests =
[ unit_empty,
-- prop_mid, -- slow
unit_butterfly,
unit_invButterfly,
unit_bigPrime,
prop_simpleSMod1,
prop_simpleSMod2,
prop_simpleInt0,
prop_simpleInt1,
prop_simpleInt2,
prop_conv64,
prop_conv641,
prop_conv18433,
unit_conv2,
prop_conv257,
prop_conv2147483647,
prop_conv2130706433
]
++ prop_conv64Bound