content: {-# LANGUAGE DataKinds #-}
  
  -- | A bit tedious tests.
  module Tests.MulMod (tests) where
  
  import BenchLib.MulMod.Barrett64 qualified as Barrett64
  import BenchLib.MulMod.BarrettWideWord qualified as BarrettWideWord
  import BenchLib.MulMod.Montgomery qualified as Montgomery
  import Data.Foldable (for_)
  import Data.Proxy (Proxy (..))
  import Data.Word (Word64)
  import GHC.TypeNats (KnownNat, natVal)
  import Test.QuickCheck qualified as QC
  import Test.Tasty
  import Test.Tasty.HUnit
  import Test.Tasty.QuickCheck qualified as QC
  
  newtype InRange s a = InRange a
    deriving (Show, Eq)
  
  instance (KnownNat s) => QC.Arbitrary (InRange s Word64) where
    arbitrary = InRange . (`mod` fromIntegral (natVal (Proxy @s))) <$> QC.arbitrary
  
  prop_barrettWideWord :: forall m. (KnownNat m) => InRange m Word64 -> InRange m Word64 -> Bool
  prop_barrettWideWord (InRange a) (InRange b) = a * b `mod` m == BarrettWideWord.mulMod bt a b
    where
      m = fromIntegral $ natVal (Proxy @m) :: Word64
      bt = BarrettWideWord.new64 m
  
  unit_barrettWideWordBounds :: Word64 -> TestTree
  unit_barrettWideWordBounds m = testCase ("barrett64 bounds " ++ show m) $ do
    let bt = BarrettWideWord.new64 m
    for_ [0 .. 4] $ \a -> do
      for_ [0 .. 4] $ \b -> do
        a * b `mod` m @=? BarrettWideWord.mulMod bt a b
    for_ [fromIntegral m - 5 .. fromIntegral m - 1] $ \a -> do
      for_ [fromIntegral m - 5 .. fromIntegral m - 1] $ \b -> do
        fromIntegral a * fromIntegral b `mod` m @=? BarrettWideWord.mulMod bt a b
  
  prop_barrett64 :: forall m. (KnownNat m) => InRange m Word64 -> InRange m Word64 -> Bool
  prop_barrett64 (InRange a) (InRange b) = (a * b `mod` m) == Barrett64.mulMod bt a b
    where
      m = fromIntegral $ natVal (Proxy @m)
      bt = Barrett64.new m
  
  unit_barrett64Bounds :: Word64 -> TestTree
  unit_barrett64Bounds m = testCase ("barrett64 bounds " ++ show m) $ do
    let bt = Barrett64.new m
    for_ [0 .. 4] $ \a -> do
      for_ [0 .. 4] $ \b -> do
        a * b `mod` m @=? Barrett64.mulMod bt a b
    for_ [m - 5 .. m - 1] $ \a -> do
      for_ [m - 5 .. m - 1] $ \b -> do
        a * b `mod` m @=? Barrett64.mulMod bt a b
  
  prop_montgomery :: forall m. (KnownNat m) => InRange m Word64 -> InRange m Word64 -> Bool
  prop_montgomery (InRange a) (InRange b) = a * b `mod` m == Montgomery.reduce mont (Montgomery.mulMod mont a b)
    where
      m = fromIntegral $ natVal (Proxy @m)
      mont = Montgomery.new m
  
  unit_montgomeryBounds :: Word64 -> TestTree
  unit_montgomeryBounds m = testCase ("montgomery64 bounds " ++ show m) $ do
    let mont = Montgomery.new m
    for_ [0 .. 4] $ \a -> do
      for_ [0 .. 4] $ \b -> do
        a * b `mod` m @=? Montgomery.reduce mont (Montgomery.mulMod mont a b)
    for_ [m - 4 .. m - 1] $ \a -> do
      for_ [m - 4 .. m - 1] $ \b -> do
        a * b `mod` m @=? Montgomery.reduce mont (Montgomery.mulMod mont a b)
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "barrettWideWord random 998244353" $ prop_barrettWideWord @998244353,
      unit_barrettWideWordBounds 998244353,
      QC.testProperty "barrett64 random 998244353" $ prop_barrett64 @998244353,
      unit_barrett64Bounds 998244353,
      QC.testProperty "montgomery random 998244353" $ prop_montgomery @998244353,
      unit_montgomeryBounds 998244353,
      -- TODO: run in a separate tree
      QC.testProperty "barrettWideWord random 2147483647" $ prop_barrettWideWord @2147483647,
      unit_barrettWideWordBounds 2147483647,
      QC.testProperty "barrett64 random 2147483647" $ prop_barrett64 @2147483647,
      unit_barrett64Bounds 2147483647,
      QC.testProperty "montgomery random 2147483647" $ prop_montgomery @2147483647,
      unit_montgomeryBounds 2147483647
    ]
  