diff --git a/Math/NumberTheory/Primes/Sieve.hs b/Math/NumberTheory/Primes/Sieve.hs index edca9e4ad..088d4cd14 100644 --- a/Math/NumberTheory/Primes/Sieve.hs +++ b/Math/NumberTheory/Primes/Sieve.hs @@ -27,8 +27,12 @@ module Math.NumberTheory.Primes.Sieve {-# DEPRECATED "Use 'Enum' instance of 'Ma , psieveList , psieveFrom , primeList + + , atkinPrimeList + , atkinSieve ) where +import Math.NumberTheory.Primes.Sieve.Atkin import Math.NumberTheory.Primes.Sieve.Eratosthenes -- $limits diff --git a/Math/NumberTheory/Primes/Sieve/Atkin.hs b/Math/NumberTheory/Primes/Sieve/Atkin.hs new file mode 100644 index 000000000..96807851d --- /dev/null +++ b/Math/NumberTheory/Primes/Sieve/Atkin.hs @@ -0,0 +1,308 @@ +-- | +-- Module: Math.NumberTheory.Primes.Sieve.Atkin +-- Copyright: (c) 2019 Andrew Lelechenko +-- Licence: MIT +-- Maintainer: Andrew Lelechenko +-- +-- Atkin sieve. +-- + +{-# LANGUAGE BangPatterns #-} + +module Math.NumberTheory.Primes.Sieve.Atkin + ( atkinPrimeList + , atkinSieve + ) where + +import Control.Monad +import Control.Monad.ST +import Data.Bit +import Data.Maybe +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as MU + +import Math.NumberTheory.Powers.Squares +import qualified Math.NumberTheory.Primes.Sieve.Eratosthenes as E +import Math.NumberTheory.Primes.Types +import Math.NumberTheory.Utils + +atkinPrimeList :: PrimeSieve -> [Int] +atkinPrimeList (PrimeSieve low len segments) + | len60 == 0 = [] + | otherwise = takeWhile (< high) $ dropWhile (< low) $ merge l0 l1 + where + list00 = map (\k -> 60 * (low60 + k) + fromWheel30 0) (listBits $ segments V.! 0) + list01 = map (\k -> 60 * (low60 + k) + fromWheel30 1) (listBits $ segments V.! 1) + list02 = map (\k -> 60 * (low60 + k) + fromWheel30 2) (listBits $ segments V.! 2) + list03 = map (\k -> 60 * (low60 + k) + fromWheel30 3) (listBits $ segments V.! 3) + list04 = map (\k -> 60 * (low60 + k) + fromWheel30 4) (listBits $ segments V.! 4) + list05 = map (\k -> 60 * (low60 + k) + fromWheel30 5) (listBits $ segments V.! 5) + list06 = map (\k -> 60 * (low60 + k) + fromWheel30 6) (listBits $ segments V.! 6) + list07 = map (\k -> 60 * (low60 + k) + fromWheel30 7) (listBits $ segments V.! 7) + list08 = map (\k -> 60 * (low60 + k) + fromWheel30 8) (listBits $ segments V.! 8) + list09 = map (\k -> 60 * (low60 + k) + fromWheel30 9) (listBits $ segments V.! 9) + list10 = map (\k -> 60 * (low60 + k) + fromWheel30 10) (listBits $ segments V.! 10) + list11 = map (\k -> 60 * (low60 + k) + fromWheel30 11) (listBits $ segments V.! 11) + list12 = map (\k -> 60 * (low60 + k) + fromWheel30 12) (listBits $ segments V.! 12) + list13 = map (\k -> 60 * (low60 + k) + fromWheel30 13) (listBits $ segments V.! 13) + list14 = map (\k -> 60 * (low60 + k) + fromWheel30 14) (listBits $ segments V.! 14) + list15 = map (\k -> 60 * (low60 + k) + fromWheel30 15) (listBits $ segments V.! 15) + + lst0 = merge list00 list01 + lst1 = merge list02 list03 + lst2 = merge list04 list05 + lst3 = merge list06 list07 + lst4 = merge list08 list09 + lst5 = merge list10 list11 + lst6 = merge list12 list13 + lst7 = merge list14 list15 + + ls0 = merge lst0 lst1 + ls1 = merge lst2 lst3 + ls2 = merge lst4 lst5 + ls3 = merge lst6 lst7 + + l0 = merge ls0 ls1 + l1 = merge ls2 ls3 + + low60 = low `quot` 60 + len60 = (low + len + 59) `quot` 60 - low60 + high = low + len + +merge :: Ord a => [a] -> [a] -> [a] +merge [] ys = ys +merge xs [] = xs +merge xs@(x:xs') ys@(y:ys') + | x < y = x : merge xs' ys + | otherwise = y : merge xs ys' + +data PrimeSieve = PrimeSieve + { _psLowBound :: !Int + , _psLength :: !Int + , _psSegments :: V.Vector (U.Vector Bit) + } deriving (Show) + +atkinSieve + :: Int + -> Int + -> PrimeSieve +atkinSieve low len = PrimeSieve low len segments + where + low60 = low `quot` 60 + len60 = (low + len + 59) `quot` 60 - low60 + params = V.generate 16 (\i -> SieveParams (fromWheel30 i) low60 len60) + segments = V.map sieveSegment params + +data SieveParams = SieveParams + { spDelta :: !Int + , spLowBound :: !Int + , spLength :: !Int + } deriving (Show) + +spHighBound :: SieveParams -> Int +spHighBound sp = spLowBound sp + spLength sp + +sieveSegment + :: SieveParams + -> U.Vector Bit +sieveSegment sp = runST $ do + vec <- MU.new (spLength sp) + U.forM_ (fgs V.! toWheel30 (spDelta sp)) $ + traverseLatticePoints sp vec + algo3steps456 sp vec + U.unsafeFreeze vec + +-- | Solutions of k * f^2 + l * g^2 = delta (mod 60) +-- where (k, l) = (4, 1) for delta = 1 (mod 4) +-- = (3, 1) for delta = 1 (mod 6) +-- = (3,-1) for delta =11 (mod 12) +fgs :: V.Vector (U.Vector (Int, Int)) +fgs = V.generate 16 (dispatch . fromWheel30) + where + dispatch delta + | delta `mod` 4 == 1 + = U.fromList [ (f, g) | f <- [1..15], g <- [1..30], (4*f*f + g*g - delta) `rem` 60 == 0] + | delta `mod` 6 == 1 + = U.fromList [ (f, g) | f <- [1..10], g <- [1..30], (3*f*f + g*g - delta) `rem` 60 == 0] + | delta `mod` 12 == 11 + = U.fromList [ (f, g) | f <- [1..10], g <- [1..30], (3*f*f - g*g - delta) `rem` 60 == 0] + | otherwise + = error "fgs: unexpected delta" + +traverseLatticePoints + :: SieveParams + -> MU.MVector s Bit + -> (Int, Int) + -> ST s () +traverseLatticePoints sp vec (x0, y0) + | spDelta sp `mod` 4 == 1 + = traverseLatticePoints1 sp vec (x0, y0) + | spDelta sp `mod` 6 == 1 + = traverseLatticePoints2 sp vec (x0, y0) + | spDelta sp `mod` 12 == 11 + = traverseLatticePoints3 sp vec (x0, y0) + | otherwise + = error "traverseLatticePoints: unexpected delta" + +traverseLatticePoints1 + :: SieveParams + -> MU.MVector s Bit + -> (Int, Int) + -> ST s () +traverseLatticePoints1 !sp vec (!x0, !y0) = + go kMax xMax y0 + where + forwardY (k, y) = (k + y + 15, y + 30) + forwardX (k, x) = (k + 2 * x + 15, x + 15) + backwardX (k, x) = (k - 2 * x + 15, x - 15) + + -- Step 1 + k0 = (4 * x0 * x0 + y0 * y0 - spDelta sp) `quot` 60 - spLowBound sp + + -- Step 2 + (kMax, xMax) + = backwardX + $ head + $ dropWhile (\(k, _) -> k < spLength sp) + $ iterate forwardX + $ (k0, x0) + + -- Step 4 + adjustY (!k, !y) + | k >= 0 + = (k, y) + | otherwise + = adjustY $ forwardY (k, y) + + -- Step 6 + doActions (!k, !y) + | k < spLength sp + = unsafeFlipBit vec k >> doActions (forwardY (k, y)) + | otherwise + = pure () + + go !k !x !y + | x <= 0 = pure () + | otherwise = do + let (k', y') = adjustY (k, y) + doActions (k', y') + let (k'', x') = backwardX (k', x) + go k'' x' y' + +traverseLatticePoints2 + :: SieveParams + -> MU.MVector s Bit + -> (Int, Int) + -> ST s () +traverseLatticePoints2 sp vec (x0, y0) = + go kMax xMax y0 + where + forwardY (k, y) = (k + y + 15, y + 30) + forwardX (k, x) = (k + x + 5, x + 10) + backwardX (k, x) = (k - x + 5, x - 10) + + -- Step 1 + k0 = (3 * x0 * x0 + y0 * y0 - spDelta sp) `quot` 60 - spLowBound sp + + -- Step 2 + (kMax, xMax) + = backwardX + $ head + $ dropWhile (\(k, _) -> k < spLength sp) + $ iterate forwardX + $ (k0, x0) + + -- Step 4 + adjustY (!k, !y) + | k >= 0 + = (k, y) + | otherwise + = adjustY $ forwardY (k, y) + + -- Step 6 + doActions (!k, !y) + | k < spLength sp + = unsafeFlipBit vec k >> doActions (forwardY (k, y)) + | otherwise + = pure () + + go !k !x !y + | x <= 0 = pure () + | otherwise = do + let (k', y') = adjustY (k, y) + doActions (k', y') + let (k'', x') = backwardX (k', x) + go k'' x' y' + +traverseLatticePoints3 + :: SieveParams + -> MU.MVector s Bit + -> (Int, Int) + -> ST s () +traverseLatticePoints3 sp vec (x0, y0) = + go k0 x0 y0 + where + forwardY (k, y) = (k - y - 15, y + 30) + forwardX (k, x) = (k + x + 5, x + 10) + + -- Step 1 + k0 = (3 * x0 * x0 - y0 * y0 - spDelta sp) `quot` 60 - spLowBound sp + + -- Step 6 + doActions (!k, !x, !y) + | k >= 0 && y < x + = unsafeFlipBit vec k >> (let (k', y') = forwardY (k, y) in doActions (k', x, y')) + | otherwise + = pure () + + go !k !x !y + | k >= spLength sp + , x <= y + = pure () + | k >= spLength sp + = let (k', y') = forwardY (k, y) in + go k' x y' + | otherwise + = do + doActions (k, x, y) + let (k', x') = forwardX (k, x) + go k' x' y + +-- | Perform steps 4-6 of Algorithm 3.X. +algo3steps456 + :: SieveParams + -> MU.MVector s Bit + -> ST s () +algo3steps456 sp vec = + forM_ ps $ \p -> + crossMultiples sp vec (p * p) + where + low = 7 + high = integerSquareRoot (60 * spHighBound sp - 1) + ps = takeWhile (<= high) $ dropWhile (< low) $ map unPrime E.primes + +-- | Cross out multiples of the first argument +-- in a given sieve. +crossMultiples + :: SieveParams + -> MU.MVector s Bit + -> Int -- coprime with 60 + -> ST s () +crossMultiples sp vec m = + forM_ [k1, k1 + m .. spHighBound sp - 1] $ + \k -> MU.unsafeWrite vec (k - spLowBound sp) (Bit False) + where + -- k0 is the smallest k such that 60k+delta = 0 (mod m) + k0 = solveCongruence (spDelta sp) m + -- k1 = k0 (mod m), k1 >= lowBound + (q, r) = spLowBound sp `quotRem` m + k1 = if r < k0 then q * m + k0 else (q + 1) * m + k0 + +-- Find the smallest k such that 60k+delta = 0 (mod m) +-- Should be equal to +-- (`quot` 60) $ fromJust $ chineseCoprime (delta, 60) (0, m) +solveCongruence :: Int -> Int -> Int +solveCongruence delta m = (- recip60 * delta) `mod` m + where + recip60 = fromInteger $ fromJust $ recipMod 60 (toInteger m) diff --git a/app/Atkin.hs b/app/Atkin.hs new file mode 100644 index 000000000..c4669ef18 --- /dev/null +++ b/app/Atkin.hs @@ -0,0 +1,9 @@ +module Main where + +import Math.NumberTheory.Primes.Sieve + +atkin :: (Int, Int) -> Int +atkin (p, q) = sum $ atkinPrimeList $ atkinSieve p q + +main :: IO () +main = print $ atkin (10000000000,100000000) diff --git a/arithmoi.cabal b/arithmoi.cabal index 6d0d86dfa..611b5f953 100644 --- a/arithmoi.cabal +++ b/arithmoi.cabal @@ -37,6 +37,7 @@ library build-depends: base >=4.9 && <5, array >=0.5 && <0.6, + bitvec >=0.2, containers >=0.5 && <0.7, constraints, deepseq, @@ -102,6 +103,7 @@ library Math.NumberTheory.Primes.Counting.Impl Math.NumberTheory.Primes.Factorisation.Montgomery Math.NumberTheory.Primes.Factorisation.TrialDivision + Math.NumberTheory.Primes.Sieve.Atkin Math.NumberTheory.Primes.Sieve.Eratosthenes Math.NumberTheory.Primes.Sieve.Indexing Math.NumberTheory.Primes.Testing.Certificates.Internal @@ -226,3 +228,14 @@ executable sequence-model main-is: SequenceModel.hs hs-source-dirs: app default-language: Haskell2010 + +executable atkin + build-depends: + base, + arithmoi, + containers + buildable: True + main-is: Atkin.hs + hs-source-dirs: app + default-language: Haskell2010 + ghc-options: -O2 -Wall diff --git a/benchmark/Math/NumberTheory/SequenceBench.hs b/benchmark/Math/NumberTheory/SequenceBench.hs index 23f9840bc..a7f559b29 100644 --- a/benchmark/Math/NumberTheory/SequenceBench.hs +++ b/benchmark/Math/NumberTheory/SequenceBench.hs @@ -23,6 +23,9 @@ eratosthenes (p, q) = sum $ takeWhile (<= q) $ dropWhile (< p) $ map unPrime $ i then primeList $ primeSieve q else concatMap primeList $ psieveFrom p +atkin :: (Integer, Integer) -> Integer +atkin (p, q) = toInteger $ sum $ atkinPrimeList $ atkinSieve (fromInteger p) (fromInteger $ q - p) + filterIsPrimeBench :: Benchmark filterIsPrimeBench = bgroup "filterIsPrime" $ map (\(x, y) -> bench (show (x, y)) $ nf filterIsPrime (x, x + y)) @@ -40,10 +43,20 @@ eratosthenesBench = bgroup "eratosthenes" $ , x == 10 || y == 7 ] +atkinBench :: Benchmark +atkinBench = bgroup "atkin" $ + map (\(x, y) -> bench (show (x, y)) $ nf atkin (x, x + y)) + [ (10 ^ x, 10 ^ y) + | x <- [10..17] + , y <- [6..x-1] + , x == 10 || y == 7 + ] + benchSuite :: Benchmark benchSuite = bgroup "Sequence" [ filterIsPrimeBench , eratosthenesBench + , atkinBench ] ------------------------------------------------------------------------------- diff --git a/test-suite/Math/NumberTheory/Primes/SieveTests.hs b/test-suite/Math/NumberTheory/Primes/SieveTests.hs index bb339c0c6..01e21ede1 100644 --- a/test-suite/Math/NumberTheory/Primes/SieveTests.hs +++ b/test-suite/Math/NumberTheory/Primes/SieveTests.hs @@ -52,6 +52,11 @@ primesProperty2 _ = assertEqual "primes matches isPrime" (map unPrime primes :: [a]) (filter (isPrime . toInteger) [1..maxBound]) +atkinPrimesProperty1 :: Assertion +atkinPrimesProperty1 = assertEqual "atkinPrimes matches isPrime" + (atkinPrimeList $ atkinSieve 1 lim1) + (filter (isPrime . toInteger) [1..lim1]) + -- | Check that 'primeList' from 'primeSieve' matches truncated 'primes'. primeSieveProperty1 :: AnySign Integer -> Bool primeSieveProperty1 (AnySign highBound') @@ -89,7 +94,8 @@ sieveFromProperty2 (AnySign lowBound') testSuite :: TestTree testSuite = testGroup "Sieve" - [ testGroup "primes" + [ testCase "atkinPrimes" atkinPrimesProperty1 + , testGroup "primes" [ testCase "Int" (primesProperty1 (Proxy :: Proxy Int)) , testCase "Word" (primesProperty1 (Proxy :: Proxy Word)) , testCase "Integer" (primesProperty1 (Proxy :: Proxy Integer))