content: {-# 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
    ]
  