module Math.NumberTheory.Moduli.SquareRoots.Test ( prop_sqrtsPP, prop_sqrts) where import Data.Numbers.Primes (primes) import Math.NumberTheory.Moduli.SquareRoots (sqrts) import Test.QuickCheck newtype PrimePower = PrimePower (Integer, Int) deriving (Show, Eq) evalPP :: PrimePower -> Integer evalPP (PrimePower (p, e)) = p ^ e instance Arbitrary PrimePower where arbitrary = sized $ \size -> flip suchThat ((<= (fromIntegral $ size + 2)) . evalPP) $ do indexP <- choose (0, size) e <- choose (1, 20) return $ PrimePower (primes !! indexP, e) shrink (PrimePower (p, e)) = map PrimePower $ [(p', e) | p' <- takeWhile (< p) primes] ++ [(p, e') | e' <- [1, e - 1]] newtype ProblemPP = ProblemPP (Integer, PrimePower) deriving (Show, Eq) instance Arbitrary ProblemPP where arbitrary = do pp <- arbitrary a <- choose (0, evalPP pp - 1) return $ ProblemPP (a, pp) shrink (ProblemPP (a, pp)) = map ProblemPP $ [(a', pp) | a' <- [0, a - 1]] ++ [(a, pp') | pp' <- shrink pp, evalPP pp' > a] newtype Problem = Problem (Integer, Integer) deriving (Show, Eq) instance Arbitrary Problem where arbitrary = do m <- suchThat arbitrary (> 0) a <- choose (0, m - 1) return $ Problem (a, m) shrink (Problem (a, m)) = [Problem (a', m) | a' <- shrink a, a' >= 0] ++ [Problem (a, m') | m' <- shrink m, m' > a] naive :: Problem -> [Integer] naive (Problem (a, m)) = [x | x <- [0 .. (m - 1)], (x * x - a) `mod` m == 0] prop_sqrts :: Problem -> Property prop_sqrts p@(Problem (a, m)) = sqrts a m === naive p prop_sqrtsPP :: ProblemPP -> Property prop_sqrtsPP (ProblemPP (a, pp)) = prop_sqrts (Problem (a, evalPP pp))