-- | Prime numbers and related number theoretical stuff. {-# LANGUAGE BangPatterns #-} module Math.FiniteField.Primes ( -- * Integer logarithm integerLog2 , ceilingLog2 -- * Integer square root , isSquare , integerSquareRoot , ceilingSquareRoot , integerSquareRoot' , integerSquareRootNewton' -- * Elementary number theory , divides , divisors, squareFreeDivisors, squareFreeDivisors_ , divisorSum , divisorSum' , moebiusMu , eulerTotient , liouvilleLambda -- * List of prime numbers , primes , primesSimple , primesTMWE -- * Prime factorization , factorize, factorizeNaive , productOfFactors , integerFactorsTrialDivision , groupIntegerFactors -- * Modulo @m@ arithmetic , powerMod -- * Prime testing , isPrimeTrialDivision , millerRabinPrimalityTest , isProbablyPrime , isVeryProbablyPrime ) where -------------------------------------------------------------------------------- import Data.List ( group , sort , foldl' ) import Data.Bits import System.Random import Math.FiniteField.Sign import Math.FiniteField.Misc -------------------------------------------------------------------------------- -- Integer logarithm -- | Largest integer @k@ such that @2^k@ is smaller or equal to @n@ integerLog2 :: Integer -> Integer integerLog2 n = go n where go 0 = -1 go k = 1 + go (shiftR k 1) -- | Smallest integer @k@ such that @2^k@ is larger or equal to @n@ ceilingLog2 :: Integer -> Integer ceilingLog2 0 = 0 ceilingLog2 n = 1 + go (n-1) where go 0 = -1 go k = 1 + go (shiftR k 1) -------------------------------------------------------------------------------- -- Integer square root isSquare :: Integer -> Bool isSquare n = if (fromIntegral \$ mod n 32) `elem` rs then snd (integerSquareRoot' n) == 0 else False where rs = [0,1,4,9,16,17,25] :: [Int] -- | Integer square root (largest integer whose square is smaller or equal to the input) -- using Newton's method, with a faster (for large numbers) inital guess based on bit shifts. integerSquareRoot :: Integer -> Integer integerSquareRoot = fst . integerSquareRoot' -- | Smallest integer whose square is larger or equal to the input ceilingSquareRoot :: Integer -> Integer ceilingSquareRoot n = (if r>0 then u+1 else u) where (u,r) = integerSquareRoot' n -- | We also return the excess residue; that is -- -- > (a,r) = integerSquareRoot' n -- -- means that -- -- > a*a + r = n -- > a*a <= n < (a+1)*(a+1) integerSquareRoot' :: Integer -> (Integer,Integer) integerSquareRoot' n | n<0 = error "integerSquareRoot: negative input" | n<2 = (n,0) | otherwise = go firstGuess where k = integerLog2 n firstGuess = 2^(div (k+2) 2) -- !! note that (div (k+1) 2) is NOT enough !! go a = if m < a then go a' else (a, r + a*(m-a)) where (m,r) = divMod n a a' = div (m + a) 2 -- | Newton's method without an initial guess. For very small numbers (<10^10) it -- is somewhat faster than the above version. integerSquareRootNewton' :: Integer -> (Integer,Integer) integerSquareRootNewton' n | n<0 = error "integerSquareRootNewton: negative input" | n<2 = (n,0) | otherwise = go (div n 2) where go a = if m < a then go a' else (a, r + a*(m-a)) where (m,r) = divMod n a a' = div (m + a) 2 {- -- brute force test of integer square root isqrt_test n1 n2 = [ k | k<-[n1..n2] , let (a,r) = integerSquareRoot' k , (a*a+r/=k) || (a*a>k) || (a+1)*(a+1)<=k ] -} -------------------------------------------------------------------------------- -- | @d `divides` n@ divides :: Integer -> Integer -> Bool divides !d !n = (mod n d == 0) {-# SPECIALIZE moebiusMu :: Int -> Int #-} {-# SPECIALIZE moebiusMu :: Integer -> Integer #-} -- | The Moebius mu function (naive implementation) moebiusMu :: (Integral a, Num b) => a -> b moebiusMu n | any (>1) expos = 0 | even (length primes) = 1 | otherwise = -1 where factors = groupIntegerFactors \$ integerFactorsTrialDivision \$ fromIntegral n (primes,expos) = unzip factors {-# SPECIALIZE liouvilleLambda :: Int -> Int #-} {-# SPECIALIZE liouvilleLambda :: Integer -> Integer #-} -- | The Liouville lambda function (naive implementation) liouvilleLambda :: (Integral a, Num b) => a -> b liouvilleLambda n = if odd (foldl' (+) 0 \$ map snd grps) then -1 else 1 where grps = groupIntegerFactors \$ integerFactorsTrialDivision \$ fromIntegral n -- | Sum ofthe of the divisors divisorSum :: Integer -> Integer divisorSum n = foldl' (+) 0 [ d | d <- divisors n ] -- | Sum of @k@-th powers of the divisors divisorSum' :: Int -> Integer -> Integer divisorSum' k n = foldl' (+) 0 [ d^k | d <- divisors n ] -- | Euler's totient function eulerTotient :: Integer -> Integer eulerTotient n = div n prodp * prodp1 where grps = groupIntegerFactors \$ integerFactorsTrialDivision n ps = map fst grps prodp = foldl' (*) 1 [ p | p <- ps ] prodp1 = foldl' (*) 1 [ p-1 | p <- ps ] -- | Divisors of @n@ (note: the result is /not/ ordered!) divisors :: Integer -> [Integer] divisors n = [ f tup | tup <- tuples' expos ] where grps = groupIntegerFactors \$ integerFactorsTrialDivision n (ps,expos) = unzip grps f es = foldl' (*) 1 \$ zipWith (^) ps es -- | List of square-free divisors together with their Mobius mu value -- (note: the result is /not/ ordered!) squareFreeDivisors :: Integer -> [(Integer,Sign)] squareFreeDivisors n = map f (sublists primes) where grps = groupIntegerFactors \$ integerFactorsTrialDivision n primes = map fst grps f ps = ( foldl' (*) 1 ps , if even (length ps) then Plus else Minus) -- | List of square-free divisors -- (note: the result is /not/ ordered!) squareFreeDivisors_ :: Integer -> [Integer] squareFreeDivisors_ n = map f (sublists primes) where grps = groupIntegerFactors \$ integerFactorsTrialDivision n primes = map fst grps f ps = foldl' (*) 1 ps -- | To avoid cyclic dependencies, I made a local copy of this... sublists :: [a] -> [[a]] sublists [] = [[]] sublists (x:xs) = sublists xs ++ map (x:) (sublists xs) -------------------------------------------------------------------------------- -- List of prime numbers -- | Infinite list of primes, using the TMWE algorithm. primes :: [Integer] primes = primesTMWE -- | A relatively simple but still quite fast implementation of list of primes. -- By Will Ness primesSimple :: [Integer] primesSimple = 2 : 3 : sieve 0 primes' 5 where primes' = tail primesSimple sieve k (p:ps) x = noDivs k h ++ sieve (k+1) ps (t+2) where t = p*p h = [x,x+2..t-2] noDivs k = filter (\x -> all (\y -> rem x y /= 0) (take k primes')) -- | List of primes, using tree merge with wheel. Code by Will Ness. primesTMWE :: [Integer] primesTMWE = 2:3:5:7: gaps 11 wheel (fold3t \$ roll 11 wheel primes') where primes' = 11: gaps 13 (tail wheel) (fold3t \$ roll 11 wheel primes') fold3t ((x:xs): ~(ys:zs:t)) = x : union xs (union ys zs) `union` fold3t (pairs t) pairs ((x:xs):ys:t) = (x : union xs ys) : pairs t wheel = 2:4:2:4:6:2:6:4:2:4:6:6:2:6:4:2:6:4:6:8:4:2:4:2: 4:8:6:4:6:2:4:6:2:6:6:4:2:4:6:2:6:4:2:4:2:10:2:10:wheel gaps k ws@(w:t) cs@ ~(c:u) | k==c = gaps (k+w) t u | True = k : gaps (k+w) t cs roll k ws@(w:t) ps@ ~(p:u) | k==p = scanl (\c d->c+p*d) (p*p) ws : roll (k+w) t u | True = roll (k+w) t ps minus xxs@(x:xs) yys@(y:ys) = case compare x y of LT -> x : minus xs yys EQ -> minus xs ys GT -> minus xxs ys minus xs [] = xs minus [] _ = [] union xxs@(x:xs) yys@(y:ys) = case compare x y of LT -> x : union xs yys EQ -> x : union xs ys GT -> y : union xxs ys union xs [] = xs union [] ys =ys -------------------------------------------------------------------------------- -- Prime factorization factorize :: Integer -> [(Integer,Int)] factorize = factorizeNaive factorizeNaive :: Integer -> [(Integer,Int)] factorizeNaive = groupIntegerFactors . integerFactorsTrialDivision productOfFactors :: [(Integer,Int)] -> Integer productOfFactors = productInterleaved . map (uncurry pow) where pow _ 0 = 1 pow p 1 = p pow 2 n = shiftL 1 n pow p 2 = p*p pow p n = if even n then (pow p (shiftR n 1))^2 else p * (pow p (shiftR n 1))^2 -- | Groups integer factors. Example: from [2,2,2,3,3,5] we produce [(2,3),(3,2),(5,1)] groupIntegerFactors :: [Integer] -> [(Integer,Int)] groupIntegerFactors = map f . group . sort where f xs = (head xs, length xs) -- | The naive trial division algorithm. integerFactorsTrialDivision :: Integer -> [Integer] integerFactorsTrialDivision n | n<1 = error "integerFactorsTrialDivision: n should be at least 1" | otherwise = go primes n where go _ 1 = [] go rs k = sub ps k where sub [] k = [k] sub qqs@(q:qs) k = case mod k q of 0 -> q : go qqs (div k q) _ -> sub qs k ps = takeWhile (\p -> p*p <= k) rs {- go 1 = [] go k = sub ps k where sub [] k = [k] sub (q:qs) k = case mod k q of 0 -> q : go (div k q) _ -> sub qs k ps = takeWhile (\p -> p*p <= k) primes -} {- -- brute force testing of factors ifactorsTest :: (Integer -> [Integer]) -> Integer -> Bool ifactorsTest alg n = and [ product (alg k) == k | k<-[1..n] ] -} -------------------------------------------------------------------------------- -- Modulo @m@ arithmetic -- | Efficient powers modulo m. -- -- > powerMod a k m == (a^k) `mod` m powerMod :: Integer -> Integer -> Integer -> Integer powerMod a' k m = {- debug bs \$ -} go a bs where bs = bin k bin 0 = [] bin x = (x .&. 1 /= 0) : bin (shiftR x 1) a = mod a' m go _ [] = 1 go x (b:bs) = -- debug (x,b) \$ if b then mod (x*rest) m else rest where rest = go (mod (x*x) m) bs -------------------------------------------------------------------------------- -- * Prime testing -- | Prime testing using trial division isPrimeTrialDivision :: Integer -> Bool isPrimeTrialDivision n = and [ not (divides p n) | p <- ps ] where ps = takeWhile (<= nsqrt) primes nsqrt = integerSquareRoot n -- | Miller-Rabin Primality Test (taken from Haskell wiki). -- We test the primality of the first argument @n@ by using the second argument @a@ as a candidate witness. -- If it returs @False@, then @n@ is composite. If it returns @True@, then @n@ is either prime or composite. -- -- A random choice between @2@ and @(n-2)@ is a good choice for @a@. millerRabinPrimalityTest :: Integer -> Integer -> Bool millerRabinPrimalityTest n a | a <= 1 || a >= n-1 = error \$ "millerRabinPrimalityTest: a out of range (" ++ show a ++ " for "++ show n ++ ")" | n < 2 = False | even n = False | b0 == 1 || b0 == n' = True | otherwise = iter (tail b) where n' = n-1 (k,m) = find2km n' b0 = powMod n a m b = take (fromIntegral k) \$ iterate (squareMod n) b0 iter [] = False iter (x:xs) | x == 1 = False | x == n' = True | otherwise = iter xs {-# SPECIALIZE find2km :: Integer -> (Integer,Integer) #-} find2km :: Integral a => a -> (a,a) find2km n = f 0 n where f k m | r == 1 = (k,m) | otherwise = f (k+1) q where (q,r) = quotRem m 2 {-# SPECIALIZE pow' :: (Integer -> Integer -> Integer) -> (Integer -> Integer) -> Integer -> Integer -> Integer #-} pow' :: (Num a, Integral b) => (a -> a -> a) -> (a -> a) -> a -> b -> a pow' _ _ _ 0 = 1 pow' mul sq x' n' = f x' n' 1 where f !x !n !y | n == 1 = x `mul` y | r == 0 = f x2 q y | otherwise = f x2 q (x `mul` y) where (q,r) = quotRem n 2 x2 = sq x {-# SPECIALIZE mulMod :: Integer -> Integer -> Integer -> Integer #-} mulMod :: Integral a => a -> a -> a -> a mulMod !a !b !c = (b * c) `mod` a {-# SPECIALIZE squareMod :: Integer -> Integer -> Integer #-} squareMod :: Integral a => a -> a -> a squareMod !a !b = (b * b) `rem` a {-# SPECIALIZE powMod :: Integer -> Integer -> Integer -> Integer #-} powMod :: Integral a => a -> a -> a -> a powMod m = pow' (mulMod m) (squareMod m) -------------------------------------------------------------------------------- -- | For very small numbers, we use trial division, for larger numbers, we apply the -- Miller-Rabin primality test @log4(n)@ times, with candidate witnesses derived -- deterministically from @n@ using a pseudo-random sequence -- (which /should be/ based on a cryptographic hash function, but isn\'t, yet). -- -- Thus the candidate witnesses should behave essentially like random, but the -- resulting function is still a deterministic, pure function. -- -- TODO: implement the hash sequence, at the moment we use 'System.Random' instead... -- isProbablyPrime :: Integer -> Bool isProbablyPrime n | n < 2 = False | even n = (n==2) | n < 1000 = length (integerFactorsTrialDivision n) == 1 | otherwise = and [ millerRabinPrimalityTest n a | a <- witnessList ] where log2n = integerLog2 n nchecks = 1 + fromInteger (div log2n 2) :: Int witnessList = take nchecks pseudoRnds pseudoRnds = 2 : [ a | a <- integerRndSequence n , a > 1 && a < (n-1) ] -- | A more exhaustive version of 'isProbablyPrime', this one tests candidate -- witnesses both the first log4(n) prime numbers and then log4(n) pseudo-random -- numbers isVeryProbablyPrime :: Integer -> Bool isVeryProbablyPrime n | n < 2 = False | even n = (n==2) | n < 1000 = length (integerFactorsTrialDivision n) == 1 | otherwise = and [ millerRabinPrimalityTest n a | a <- witnessList ] where log2n = integerLog2 n nchecks = 1 + fromInteger (div log2n 2) :: Int witnessList = take nchecks primes ++ take nchecks pseudoRnds pseudoRnds = [ a | a <- integerRndSequence (n+3) , a > 1 && a < (n-1) ] -------------------------------------------------------------------------------- -- | Given an integer @n@, we initialize a system random generator with using a -- seed derived from @n@ (note that this uses at most 32 or 64 bits), and generate -- an infinite sequence of pseudo-random integers between @0..n-1@, generated by -- that random generator. -- -- Note that this is not really a preferred way of generating such sequences! -- integerRndSequence :: Integer -> [Integer] integerRndSequence n = randomRs (0,n-1) gen where gen = mkStdGen \$ fromInteger (n + 17 * integerLog2 n) --------------------------------------------------------------------------------