APR素数判定

多分当っている．
コード長いです．
Miller Rabin はやはり非常にシンプル．
それに比べ，APRは…
まぁ，実装がヘタなんだよと，言われればそれまでだが．
しかし，綺麗なコードがあれば見てみたいが，すくなくとも簡単にはwebでは見つからなかった．

100桁とかの素数判定を10分とか20分あれば，やってくれると思います（計算機のパワーに依存）．
ちなみに，10^123の次の素数を求めるのに約3分＠Desktop PC．

module APR where
import Data.List (sort, genericTake, genericSplitAt, genericReplicate)
import Data.Array.IArray (Array, array, (!))
import Modulo (powMod, primitiveRoots, congruence)
import IsPrime (isPrimeDivision)
import Debug.Trace (trace)

debug = True
using x | debug     =  trace ("using " ++ show x) x
        | otherwise = x

type Polynomial = [Integer]
-- [a_0, a1, ..., a_n-1, a_n]
-- <=> a_0 + a_1 x + ... + a_n-1 x^n-1 + a_n x^n

--mkPoly :: Integer -> [Integer] -> Polynomial
mkPoly p xs = mkpart 0 0.sort.map (`mod` p) $ xs
    where mkpart i c []    = c : genericReplicate (p - i - 1) 0
          mkpart i c (y:ys)
              | i == y    = mkpart i (c + 1) ys
              | otherwise = c : mkpart (i+1) 0 (y:ys)

pMulMod :: Integer -> Polynomial -> Polynomial -> Polynomial
pMulMod n xs ys = foldl1 pAdd.zipWith (\x zs -> map ((`mod` n).(x*)) zs) xs.iterate shift $ ys
    where shift zs = last zs : init zs
          pAdd xs = zipWith add $ foldr seq xs xs -- リストの中も正格評価
          add x y = if x + y >= n then x + y - n else x + y

-- xs^k (mod n)
pPowMod :: Integer -> Polynomial -> Integer -> Polynomial
pPowMod n xs k
    | k == 1    = xs
    | even k    = pPowMod n (pMulMod n xs xs) (div k 2)
    | otherwise = pMulMod n xs $ pPowMod n xs (k - 1)

normal :: Integer -> Polynomial -> Polynomial
normal n (x:xs) = map ((`mod` n).(subtract x)) xs

-- c * x^a in Z[x_p]
unit :: Integer -> Integer -> Integer -> Polynomial
unit p a c = genericReplicate a' 0 ++ [c] ++ genericReplicate (mod (p - a' - 1) p) 0
    where a' = mod a p

-- is Unit ? i.e. xs = x^a in Z_n[x_p]
isUnit :: Integer -> Polynomial -> Bool
isUnit n xs = sum xs' == 1 || all (n - 1 ==) xs'
    where xs' = normal n xs

indArray :: Integer -> Array Integer Integer
indArray q = array (1, q - 1).zip inds $ [1..]
    where g = head.primitiveRoots $ q
          inds = genericTake (q - 1).iterate ((`mod` q).(g*)) $ g

-- Jacobi Sum J(X^e, X^f) for the prime p, q
jacobiSum :: Integer -> Integer -> Array Integer Integer -> Integer -> Integer -> Polynomial
jacobiSum p q inds e f = mkPoly p [e * ind a + f * ind (1 - a)| a <- [2..q-1]]
    where ind a = inds ! (mod a $ q)

-- Gauss Sum to Jacobi Sum for the prime p, q
-- G(X)^p (mod n) = X(-1)*q*J(X,X)*J(X,X^2)*...*J(X,X^(p-2))
gaussSumP :: Integer -> Integer -> Integer -> Array Integer Integer -> Polynomial
gaussSumP n p q inds = foldl (pMulMod n) xq js
    where js = map (jacobiSum p q inds 1) [1..p-2]
          xq = unit p (inds ! (q-1)) q

-- Gauss Sum toJacobi Sum for arbitrarily number n (> p)
-- G(X)^n = poly * G(X^t) where t = mod n p
gaussSum :: Integer -> Integer -> Integer -> Array Integer Integer -> Polynomial
gaussSum n p q inds = foldl (pMulMod n) gpa js
    where gpa = pPowMod n (gaussSumP n p q inds) a
          (a,b) = divMod n p
          js = map (jacobiSum p q inds 1) [1..b-1]

-- find min w s.t. G(X)^(u*p^w) == unit
toUnit :: Integer -> Integer -> Integer -> Integer -> (Integer, Polynomial)
toUnit n p q u = fst $ until (isUnit n.snd) next ((1,gpu), gpu)
    where gpu = pPowMod n (gaussSumP n p q ind) u
          next ((w, g0), g1) = ((w+1, g1), pPowMod n g1 p)
          ind = indArray q

-- coefficient condition; False => composite
coefCond :: Integer -> Integer -> Polynomial -> Bool
coefCond n p xs = all cpart [0..p - 1]
    where cpart i = any ((1==).gcd n) $ descend i
          descend j = let (as, (b:bs)) = genericSplitAt j xs
                      in as ++ [b - 1] ++ bs

unitSearch :: Integer -> Polynomial -> Integer
unitSearch n xs | x > 1     = 0
                | otherwise = upart 1 $ x:xs'
    where (x:xs') = normal n xs
          upart a (1:ys) = a
          upart a (y:ys) = upart (a + 1) ys

aprPrimality :: Integer -> Integer -> Integer -> Integer -> [Integer] -> Bool -> Bool
aprPrimality _ _ _ _ [] _ = True
aprPrimality n t s p (q:qs) sub
    | normal n g /= normal n c = False
    | a /= 0 || r /= 1         = if all ((0/=).(mod n)) rs
                                 then aprPrimality n t s p qs sub
                                 else False
    | sub                      = aprPrimality n t s p qs sub
    | subtest                  = aprPrimality n t s p qs True
    | otherwise                = False
    where rs = genericTake t.takeWhile (1/=) $ iterate ((`mod` s).(n*)) (mod n s)
          g | p == 2    = let a = mod (div n 2 * ind ! (q - 1)) 2
                          in unit 2 a $ powMod q (div n 2) n
            | otherwise = gaussSum n p q $ ind
          c = unit p (mod (-n) p * ind ! (mod n q)) 1
          ind = indArray q
          a = mod (ind ! (mod n q)) p
          r = powMod n (p - 1) (p ^ 2)
          subtest | debug     = subPrimality n t s $ trace ("sub-test " ++ show p) p
                  | otherwise = subPrimality n t s p

subPrimality :: Integer -> Integer -> Integer -> Integer -> Bool
subPrimality n t s p
    | not.all (coefCond n p) $ xss = False
    | otherwise                    = all ((0/=).(mod n)) rs
    where u = until ((0/=).(`mod` p)) (`div` p) $ n ^ (p - 1) - 1
          qs = filter ((0==).(`mod` p).pred).firstPrimes.initialPrimes $ n
          (w, xss) = foldl search (1, []) qs
          search (w, xss) q
              | w > w'  = (w, xss)
              | w' > w  = (w', [xs])
              | w' == 1 = (w, [])
              | w' == w = if all (n - 1 ==).normal n.pPowMod n xs $ p
                          then (w, xs:xss) else (w, xss)
              where (w', xs) = toUnit n p q u
          unitRoot q =let z = pPowMod n (gaussSumP n p q (indArray q)) $ p ^ (w - 1) * u
                          a = unitSearch n z
                          g = head $ primitiveRoots q
                      in powMod g a q
          m = congruence qs.map unitRoot $ qs
          rs = genericTake t.takeWhile (1/=) $ iterate ((`mod` s).(m*)) (mod m s)

isPrimeAPR :: Integer -> Bool
isPrimeAPR n
    | gcd n t /= 1 = False
    | gcd n s /= 1 = False
    | otherwise    = and [aprPrimality n t s (using p) qs' False|
                          p <- ps, let qs' = filter ((0==).(`mod` p).pred) qs]
    where ps = initialPrimes $ if debug then trace ("test " ++ show n) n else n
          qs = firstPrimes ps
          t = product ps
          s = product qs

firstPrimes :: [Integer] -> [Integer]
firstPrimes ps = filter isPrimeDivision [succ.product $ ps' | ps' <- subsequences ps]

initialPrimes :: Integer -> [Integer]
initialPrimes n
    | n <= 138      = []
    | n <= 10 ^ 4   = [2,11]
    | n <= 10 ^ 7   = [2,5,7]
    | n <= 10 ^ 8   = [2,3,5]
    | n <= 10 ^ 12  = [2,3,5,17]
    | n <= 10 ^ 19  = [2,3,5,7]
    | n <= 10 ^ 35  = [2,3,5,7,17]
    | n <= 10 ^ 46  = [2,3,5,7,13]
    | n <= 10 ^ 89  = [2,3,5,7,11,13]
    | n <= 10 ^ 166 = [2,3,5,7,11,13,17]
    | n <= 10 ^ 329 = [2,3,5,7,11,13,17,23]
    | n <= 10 ^ 620 = [2,3,5,7,11,13,17,19,23]

-- subsequences from Data.List GHC 6.8.10
subsequences :: [a] -> [[a]]
subsequences xs =  [] : nonEmptySubsequences xs
nonEmptySubsequences :: [a] -> [[a]]
nonEmptySubsequences []      =  []
nonEmptySubsequences (x:xs)  =  [x] : foldr f [] (nonEmptySubsequences xs)
  where f ys r = ys : (x : ys) : r

すこーし，コードがカオス．
しかし，そもそも場合分けが多いアルゴリズムだから，しょうがないのか？
primitive root とか

module Modulo (gcd', productMod, invMod, powMod, jacobi, sqrtMod, sqSum, primitiveRoots, congruence) where
import Data.List (foldl', findIndex, nub)

-- input; m,n:正整数
-- output; a,b s.t. am+bn=gcd(m,n)
gcd' :: Integral a => a -> a -> (a, a)
gcd' m n | m < n = let (a,b) = gcd'  n m in (b,a)
         | r == 0 = (1,1-q)
         | otherwise = let (a,b) = gcd' n r
                       in (b,a-b*q)
         where (q,r) = divMod m n

-- product xs (mod p)

productMod :: Integral a => a -> [a] -> a
productMod p = foldl' (mulMod p) 1
    where mulMod p a b = mod (a*b) p

-- xn=1(mod p) なる x を返す
invMod :: Integral a => a -> a -> a
invMod n = fst.gcd' n

-- a^n (mod p)
powMod :: (Integral a, Integral b) => a -> b -> a -> a
powMod a n p | n == 0 = 1
             | even n = powMod (mod (a*a) p) (div n 2) p
             | otherwise = (a * powMod a (n-1) p)`mod` p

-- ============================================================
-- 4n+1型素数の平方和分解
-- ============================================================

jacobi :: Integral a => a -> a -> a
jacobi a n | a < 0 && mod n 4 == 1 = jacobi (-a) n
           | a < 0 && mod n 4 == 3 = - jacobi (-a) n
           | a < 2 = a
           | even a && (mod n 8 == 1 || mod n 8 == 7) = jacobi (div a 2) n
           | even a && (mod n 8 == 3 || mod n 8 == 5) = - jacobi (div a 2) n
           | mod a 4 == 3 && mod n 4 == 3 = - jacobi (mod n a) a
           | otherwise = jacobi (mod n a) a

-- sqrt n (mod p)
sqrtMod :: Integral a => a -> a -> a
sqrtMod n p =
    let m = fst.gcd' n$ p
        w = head.filter((==(-1)).flip jacobi p)$[2..]
        (s,q) = until(odd.snd) div2$ (0,p-1)
        v = powMod w q p
        step u = case findIndex (==1).take s.iterate sqMod $ mod (u^2*m) p of
                   Just 0 -> u
                   Just i -> step $ (u*(powMod v (2^(s-i-1)) p)) `mod` p
                   Nothing -> error$ "There is no such x that x^2 ="++show n++" (mod "++show p++")"
    in step $ powMod n (div (q+1) 2) p
    where div2 (s,q) = (s+1,div q 2)
          sqMod x = powMod x 2 p

-- 4n+1型素数の平方和分解
sqSum :: Integral a => a -> (a, a)
sqSum p = let r = sqrtMod (p-1) p
          in fst.until((==1).snd) next$ ((1,r),div (r*r+1) p)
    where next ((x,y),k) =
              let u = g$ x `mod` k
                  v = g$ y `mod` k
                  x' = abs$ (x*v-y*u) `div` k
                  y' = abs$ (x*u+y*v) `div` k
                  k' = (u*u+v*v) `div` k
              in ((x',y'),k')
              where g q = if q < div k 2 then q else q-k

-- ============================================================
-- 原始根
-- ============================================================

-- Given an odd prime p, find a primitive root modulo p
primitiveRoots :: Integer -> [Integer]
primitiveRoots p = filter isRoot [2..p - 1]
    where isRoot a = all (/=1) [powMod a (div (p - 1) q) p | q <- ps]
          ps = nub.primeFactors $ p - 1

-- ============================================================
-- 連立合同方程式
-- ============================================================

-- Input: 互いに素な整数 m1,m2,...,mn, 整数 a1,a2,...,an
-- Output: 連立合同方程式の解 x in Zm s.t. x = ai (mod mi)

congruence :: (Integral a) => [a] -> [a] -> a
congruence ms as = (`mod` m).sum.zipWith3 add as vs $ us
    where m = product ms
          us = [div m mi | mi <- ms]
          vs = zipWith invMod us ms
          add a v u = a * v * u

-- ============================================================
-- Primality check by trial division
-- ============================================================

primes :: Integral a => [a]
primes = 2:filter isPrimeDivision [3,5..]

isPrimeDivision :: Integral a => a -> Bool
isPrimeDivision 1 = False
isPrimeDivision n = all ((0/=).mod n).takeWhile (<=u) $ primes
    where u = floor.sqrt.fromIntegral $ n

primeFactors n = pf n primes
    where pf m (p:ps)
              | p * p > m    = [m]
              | mod m p == 0 = p : pf (div m p) (p:ps)
              | otherwise    = pf m ps