-- |
-- Module:      Math.NumberTheory.Primes.Counting.Impl
-- Copyright:   (c) 2011 Daniel Fischer
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
--
-- Number of primes not exceeding @n@, @&#960;(n)@, and @n@-th prime.
--
{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -fspec-constr-count=24 #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Math.NumberTheory.Primes.Counting.Impl
    ( primeCount
    , primeCountMaxArg
    , nthPrime
    ) where

import Math.NumberTheory.Primes.Sieve.Eratosthenes
    (PrimeSieve(..), primeList, primeSieve, psieveFrom, sieveTo, sieveBits, sieveRange)
import Math.NumberTheory.Primes.Sieve.Indexing (toPrim, idxPr)
import Math.NumberTheory.Primes.Counting.Approximate (nthPrimeApprox, approxPrimeCount)
import Math.NumberTheory.Primes.Types
import Math.NumberTheory.Roots
import Math.NumberTheory.Utils.FromIntegral

import Control.Monad.ST
import Data.Array.Base
import Data.Array.ST
import Data.Bits
import Data.Int
import Unsafe.Coerce

-- | Maximal allowed argument of 'primeCount'. Currently 8e18.
primeCountMaxArg :: Integer
primeCountMaxArg :: Integer
primeCountMaxArg = Integer
8000000000000000000

-- | @'primeCount' n == &#960;(n)@ is the number of (positive) primes not exceeding @n@.
--
--   For efficiency, the calculations are done on 64-bit signed integers, therefore @n@ must
--   not exceed 'primeCountMaxArg'.
--
--   Requires @/O/(n^0.5)@ space, the time complexity is roughly @/O/(n^0.7)@.
--   For small bounds, @'primeCount' n@ simply counts the primes not exceeding @n@,
--   for bounds from @30000@ on, Meissel's algorithm is used in the improved form due to
--   D.H. Lehmer, cf.
--   <http://en.wikipedia.org/wiki/Prime_counting_function#Algorithms_for_evaluating_.CF.80.28x.29>.
primeCount :: Integer -> Integer
primeCount :: Integer -> Integer
primeCount Integer
n
    | Integer
n forall a. Ord a => a -> a -> Bool
> Integer
primeCountMaxArg = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"primeCount: can't handle bound " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Integer
n
    | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
2     = Integer
0
    | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
1000  = Int -> Integer
intToInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a. Ord a => a -> a -> Bool
<= Integer
n) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Prime a -> a
unPrime forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => PrimeSieve -> [Prime a]
primeList forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> PrimeSieve
primeSieve forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> a -> a
max Integer
242 Integer
n
    | Integer
n forall a. Ord a => a -> a -> Bool
< Integer
30000 = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
        STUArray s Int Bool
ba <- forall s. Integer -> ST s (STUArray s Int Bool)
sieveTo Integer
n
        (Int
s,Int
e) <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds STUArray s Int Bool
ba
        Int
ct <- forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
s Int
e STUArray s Int Bool
ba
        forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Integer
intToInteger forall a b. (a -> b) -> a -> b
$ Int
ctforall a. Num a => a -> a -> a
+Int
3)
    | Bool
otherwise =
        let !ub :: Int64
ub = Int64 -> Int64
cop forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
n
            !sr :: Int64
sr = forall a. Integral a => a -> a
integerSquareRoot Int64
ub
            !cr :: Int64
cr = forall a. Integral a => a -> a
nxtEnd forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> a
integerCubeRoot Int64
ub forall a. Num a => a -> a -> a
+ Int64
15
            nxtEnd :: a -> a
nxtEnd a
k = a
k forall a. Num a => a -> a -> a
- (a
k forall a. Integral a => a -> a -> a
`rem` a
30) forall a. Num a => a -> a -> a
+ a
31
            !phn1 :: Integer
phn1 = Int64 -> Int64 -> Integer
calc Int64
ub Int64
cr
            !cs :: Int64
cs = Int64
crforall a. Num a => a -> a -> a
+Int64
6
            !pdf :: Integer
pdf = Int64 -> Int64 -> Int64 -> Integer
sieveCount Int64
ub Int64
cs Int64
sr
        in Integer
phn1 forall a. Num a => a -> a -> a
- Integer
pdf

-- | @'nthPrime' n@ calculates the @n@-th prime. Numbering of primes is
--   @1@-based, so @'nthPrime' 1 == 2@.
--
--   Requires @/O/((n*log n)^0.5)@ space, the time complexity is roughly @/O/((n*log n)^0.7@.
--   The argument must be strictly positive.
nthPrime :: Int -> Prime Integer
nthPrime :: Int -> Prime Integer
nthPrime Int
1 = forall a. a -> Prime a
Prime Integer
2
nthPrime Int
2 = forall a. a -> Prime a
Prime Integer
3
nthPrime Int
3 = forall a. a -> Prime a
Prime Integer
5
nthPrime Int
4 = forall a. a -> Prime a
Prime Integer
7
nthPrime Int
5 = forall a. a -> Prime a
Prime Integer
11
nthPrime Int
6 = forall a. a -> Prime a
Prime Integer
13
nthPrime Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
1
    = forall a. HasCallStack => [Char] -> a
error [Char]
"Prime indexing starts at 1"
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
200000
    = forall a. a -> Prime a
Prime forall a b. (a -> b) -> a -> b
$ Int -> [PrimeSieve] -> Integer
countToNth (Int
n forall a. Num a => a -> a -> a
- Int
3) [Integer -> PrimeSieve
primeSieve (Integer
p0 forall a. Num a => a -> a -> a
+ Integer
p0 forall a. Integral a => a -> a -> a
`quot` Integer
32 forall a. Num a => a -> a -> a
+ Integer
37)]
    | Integer
p0 forall a. Ord a => a -> a -> Bool
> forall a. Integral a => a -> Integer
toInteger (forall a. Bounded a => a
maxBound :: Int)
    = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"nthPrime: index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
n forall a. [a] -> [a] -> [a]
++ [Char]
" is too large to handle"
    | Int
miss forall a. Ord a => a -> a -> Bool
> Int
0
    = forall a. a -> Prime a
Prime forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Integer
tooLow  Int
n (forall a. Num a => Integer -> a
fromInteger Integer
p0) Int
miss
    | Bool
otherwise
    = forall a. a -> Prime a
Prime forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Integer
tooHigh Int
n (forall a. Num a => Integer -> a
fromInteger Integer
p0) (forall a. Num a => a -> a
negate Int
miss)
      where
        p0 :: Integer
p0 = Integer -> Integer
nthPrimeApprox (forall a. Integral a => a -> Integer
toInteger Int
n)
        miss :: Int
miss = Int
n forall a. Num a => a -> a -> a
- forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
primeCount Integer
p0)

--------------------------------------------------------------------------------
--                                The Works                                   --
--------------------------------------------------------------------------------

-- TODO: do something better in case we guess too high.
-- Not too pressing, since I think a) nthPrimeApprox always underestimates
-- in the range we can handle, and b) it's always "goodEnough"

tooLow :: Int -> Int -> Int -> Integer
tooLow :: Int -> Int -> Int -> Integer
tooLow Int
n Int
p0 Int
shortage
  | Integer
p1 forall a. Ord a => a -> a -> Bool
> forall a. Integral a => a -> Integer
toInteger (forall a. Bounded a => a
maxBound :: Int)
  = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"nthPrime: index " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Int
n forall a. [a] -> [a] -> [a]
++ [Char]
" is too large to handle"
  | Bool
goodEnough
  = Int -> Int -> Integer
lowSieve Int
p0 Int
shortage
  | Int
c1 forall a. Ord a => a -> a -> Bool
< Int
n
  = Int -> Int -> Integer
lowSieve (forall a. Num a => Integer -> a
fromInteger Integer
p1) (Int
nforall a. Num a => a -> a -> a
-Int
c1)
  | Bool
otherwise
  = Int -> Int -> Integer
lowSieve Int
p0 Int
shortage   -- a third count wouldn't make it faster, I think
  where
    gap :: Integer
gap = forall a b. (RealFrac a, Integral b) => a -> b
truncate (forall a. Floating a => a -> a
log (Int -> Double
intToDouble Int
p0 :: Double))
    est :: Integer
est = forall a. Integral a => a -> Integer
toInteger Int
shortage forall a. Num a => a -> a -> a
* Integer
gap
    p1 :: Integer
p1  = forall a. Integral a => a -> Integer
toInteger Int
p0 forall a. Num a => a -> a -> a
+ Integer
est
    goodEnough :: Bool
goodEnough = Integer
3forall a. Num a => a -> a -> a
*Integer
estforall a. Num a => a -> a -> a
*Integer
estforall a. Num a => a -> a -> a
*Integer
est forall a. Ord a => a -> a -> Bool
< Integer
2forall a. Num a => a -> a -> a
*Integer
p1forall a. Num a => a -> a -> a
*Integer
p1    -- a second counting would be more work than sieving
    c1 :: Int
c1  = forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
primeCount Integer
p1)

tooHigh :: Int -> Int -> Int -> Integer
tooHigh :: Int -> Int -> Int -> Integer
tooHigh Int
n Int
p0 Int
surplus
  | Int
c forall a. Ord a => a -> a -> Bool
< Int
n
  = Int -> Int -> Integer
lowSieve Int
b (Int
nforall a. Num a => a -> a -> a
-Int
c)
  | Bool
otherwise
  = Int -> Int -> Int -> Integer
tooHigh Int
n Int
b (Int
cforall a. Num a => a -> a -> a
-Int
n)
  where
    gap :: Int
gap = forall a b. (RealFrac a, Integral b) => a -> b
truncate (forall a. Floating a => a -> a
log (Int -> Double
intToDouble Int
p0 :: Double))
    b :: Int
b = Int
p0 forall a. Num a => a -> a -> a
- (Int
surplus forall a. Num a => a -> a -> a
* Int
gap forall a. Num a => a -> a -> a
* Int
11) forall a. Integral a => a -> a -> a
`quot` Int
10
    c :: Int
c = forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
primeCount (forall a. Integral a => a -> Integer
toInteger Int
b))

lowSieve :: Int -> Int -> Integer
lowSieve :: Int -> Int -> Integer
lowSieve Int
a Int
miss = Int -> [PrimeSieve] -> Integer
countToNth (Int
missforall a. Num a => a -> a -> a
+Int
rep) [PrimeSieve]
psieves
      where
        strt :: Int
strt = Int
a forall a. Num a => a -> a -> a
+ Int
1 forall a. Num a => a -> a -> a
+ (Int
a forall a. Bits a => a -> a -> a
.&. Int
1)
        psieves :: [PrimeSieve]
psieves@(PS Integer
vO UArray Int Bool
ba:[PrimeSieve]
_) = Integer -> [PrimeSieve]
psieveFrom (forall a. Integral a => a -> Integer
toInteger Int
strt)
        rep :: Int
rep | Integer
o0 forall a. Ord a => a -> a -> Bool
< Integer
0    = Int
0
            | Bool
otherwise = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int
1 | Int
i <- [Int
0 .. Int
r2], UArray Int Bool
ba forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` Int
i]
              where
                o0 :: Integer
o0 = forall a. Integral a => a -> Integer
toInteger Int
strt forall a. Num a => a -> a -> a
- Integer
vO forall a. Num a => a -> a -> a
- Integer
9   -- (strt - 2) - v0 - 7
                r0 :: Int
r0 = forall a. Num a => Integer -> a
fromInteger Integer
o0 forall a. Integral a => a -> a -> a
`rem` Int
30
                r1 :: Int
r1 = Int
r0 forall a. Integral a => a -> a -> a
`quot` Int
3
                r2 :: Int
r2 = forall a. Ord a => a -> a -> a
min Int
7 (if Int
r1 forall a. Ord a => a -> a -> Bool
> Int
5 then Int
r1forall a. Num a => a -> a -> a
-Int
1 else Int
r1)

-- highSieve :: Integer -> Integer -> Integer -> Integer
-- highSieve a surp gap = error "Oh shit"

sieveCount :: Int64 -> Int64 -> Int64 -> Integer
sieveCount :: Int64 -> Int64 -> Int64 -> Integer
sieveCount Int64
ub Int64
cr Int64
sr = forall a. (forall s. ST s a) -> a
runST (forall s. Int64 -> Int64 -> Int64 -> ST s Integer
sieveCountST Int64
ub Int64
cr Int64
sr)

sieveCountST :: forall s. Int64 -> Int64 -> Int64 -> ST s Integer
sieveCountST :: forall s. Int64 -> Int64 -> Int64 -> ST s Integer
sieveCountST Int64
ub Int64
cr Int64
sr = do
    let psieves :: [PrimeSieve]
psieves = Integer -> [PrimeSieve]
psieveFrom (Int64 -> Integer
int64ToInteger Int64
cr)
        pisr :: Int64
pisr = forall a. Integral a => a -> a
approxPrimeCount Int64
sr
        picr :: Int64
picr = forall a. Integral a => a -> a
approxPrimeCount Int64
cr
        diff :: Int64
diff = Int64
pisr forall a. Num a => a -> a -> a
- Int64
picr
        size :: Int
size = Int64 -> Int
int64ToInt (Int64
diff forall a. Num a => a -> a -> a
+ Int64
diff forall a. Integral a => a -> a -> a
`quot` Int64
50) forall a. Num a => a -> a -> a
+ Int
30
    STUArray s Int Int64
store <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
unsafeNewArray_ (Int
0,Int
sizeforall a. Num a => a -> a -> a
-Int
1) :: ST s (STUArray s Int Int64)
    let feed :: Int64 -> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
        feed :: Int64
-> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
feed Int64
voff !Int
wi !Int
ri UArray Int Bool
uar [PrimeSieve]
sves
          | Int
ri forall a. Eq a => a -> a -> Bool
== Int
sieveBits = case [PrimeSieve]
sves of
                                (PS Integer
vO UArray Int Bool
ba : [PrimeSieve]
more) -> Int64
-> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
feed (forall a. Num a => Integer -> a
fromInteger Integer
vO) Int
wi Int
0 UArray Int Bool
ba [PrimeSieve]
more
                                [PrimeSieve]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"prime stream ended prematurely"
          | Int64
pval forall a. Ord a => a -> a -> Bool
> Int64
sr   = do
              STUArray s Int Bool
stu <- forall i (a :: * -> * -> *) e (b :: * -> * -> *) (m :: * -> *).
(Ix i, IArray a e, MArray b e m) =>
a i e -> m (b i e)
unsafeThaw UArray Int Bool
uar
              Integer
-> Integer
-> Int64
-> Int
-> Int
-> STUArray s Int Bool
-> [PrimeSieve]
-> ST s Integer
eat Integer
0 Integer
0 Int64
voff (Int
wiforall a. Num a => a -> a -> a
-Int
1) Int
ri STUArray s Int Bool
stu [PrimeSieve]
sves
          | UArray Int Bool
uar forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` Int
ri = do
              forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
store Int
wi (Int64
ub forall a. Integral a => a -> a -> a
`quot` Int64
pval)
              Int64
-> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
feed Int64
voff (Int
wiforall a. Num a => a -> a -> a
+Int
1) (Int
riforall a. Num a => a -> a -> a
+Int
1) UArray Int Bool
uar [PrimeSieve]
sves
          | Bool
otherwise = Int64
-> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
feed Int64
voff Int
wi (Int
riforall a. Num a => a -> a -> a
+Int
1) UArray Int Bool
uar [PrimeSieve]
sves
            where
              pval :: Int64
pval = Int64
voff forall a. Num a => a -> a -> a
+ forall a. Num a => Int -> a
toPrim Int
ri
        eat :: Integer -> Integer -> Int64 -> Int -> Int -> STUArray s Int Bool -> [PrimeSieve] -> ST s Integer
        eat :: Integer
-> Integer
-> Int64
-> Int
-> Int
-> STUArray s Int Bool
-> [PrimeSieve]
-> ST s Integer
eat !Integer
acc !Integer
btw Int64
voff !Int
wi !Int
si STUArray s Int Bool
stu [PrimeSieve]
sves
            | Int
si forall a. Eq a => a -> a -> Bool
== Int
sieveBits =
                case [PrimeSieve]
sves of
                  [] -> forall a. HasCallStack => [Char] -> a
error [Char]
"Premature end of prime stream"
                  (PS Integer
vO UArray Int Bool
ba : [PrimeSieve]
more) -> do
                      STUArray s Int Bool
nstu <- forall i (a :: * -> * -> *) e (b :: * -> * -> *) (m :: * -> *).
(Ix i, IArray a e, MArray b e m) =>
a i e -> m (b i e)
unsafeThaw UArray Int Bool
ba
                      Integer
-> Integer
-> Int64
-> Int
-> Int
-> STUArray s Int Bool
-> [PrimeSieve]
-> ST s Integer
eat Integer
acc Integer
btw (forall a. Num a => Integer -> a
fromInteger Integer
vO) Int
wi Int
0 STUArray s Int Bool
nstu [PrimeSieve]
more
            | Int
wi forall a. Ord a => a -> a -> Bool
< Int
0    = forall (m :: * -> *) a. Monad m => a -> m a
return Integer
acc
            | Bool
otherwise = do
                Int64
qb <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
store Int
wi
                let dist :: Int64
dist = Int64
qb forall a. Num a => a -> a -> a
- Int64
voff forall a. Num a => a -> a -> a
- Int64
7
                if Int64
dist forall a. Ord a => a -> a -> Bool
< Int -> Int64
intToInt64 Int
sieveRange
                  then do
                      let (Int
b,Int
j) = forall a. Integral a => a -> (Int, Int)
idxPr (Int64
distforall a. Num a => a -> a -> a
+Int64
7)
                          !li :: Int
li = (Int
b forall a. Bits a => a -> Int -> a
`shiftL` Int
3) forall a. Bits a => a -> a -> a
.|. Int
j
                      Int
new <- if Int
li forall a. Ord a => a -> a -> Bool
< Int
si then forall (m :: * -> *) a. Monad m => a -> m a
return Int
0 else forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
si Int
li STUArray s Int Bool
stu
                      let nbtw :: Integer
nbtw = Integer
btw forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new forall a. Num a => a -> a -> a
+ Integer
1
                      Integer
-> Integer
-> Int64
-> Int
-> Int
-> STUArray s Int Bool
-> [PrimeSieve]
-> ST s Integer
eat (Integer
accforall a. Num a => a -> a -> a
+Integer
nbtw) Integer
nbtw Int64
voff (Int
wiforall a. Num a => a -> a -> a
-Int
1) (Int
liforall a. Num a => a -> a -> a
+Int
1) STUArray s Int Bool
stu [PrimeSieve]
sves
                  else do
                      let (Int64
cpl,Int64
fds) = Int64
dist forall a. Integral a => a -> a -> (a, a)
`quotRem` Int -> Int64
intToInt64 Int
sieveRange
                          (Int
b,Int
j) = forall a. Integral a => a -> (Int, Int)
idxPr (Int64
fdsforall a. Num a => a -> a -> a
+Int64
7)
                          !li :: Int
li = (Int
b forall a. Bits a => a -> Int -> a
`shiftL` Int
3) forall a. Bits a => a -> a -> a
.|. Int
j
                          ctLoop :: Integer -> t -> [PrimeSieve] -> ST s Integer
ctLoop !Integer
lac t
0 (PS Integer
vO UArray Int Bool
ba : [PrimeSieve]
more) = do
                              STUArray s Int Bool
nstu <- forall i (a :: * -> * -> *) e (b :: * -> * -> *) (m :: * -> *).
(Ix i, IArray a e, MArray b e m) =>
a i e -> m (b i e)
unsafeThaw UArray Int Bool
ba
                              Int
new <- forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
0 Int
li STUArray s Int Bool
nstu
                              let nbtw :: Integer
nbtw = Integer
btw forall a. Num a => a -> a -> a
+ Integer
lac forall a. Num a => a -> a -> a
+ Integer
1 forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new
                              Integer
-> Integer
-> Int64
-> Int
-> Int
-> STUArray s Int Bool
-> [PrimeSieve]
-> ST s Integer
eat (Integer
accforall a. Num a => a -> a -> a
+Integer
nbtw) Integer
nbtw (Integer -> Int64
integerToInt64 Integer
vO) (Int
wiforall a. Num a => a -> a -> a
-Int
1) (Int
liforall a. Num a => a -> a -> a
+Int
1) STUArray s Int Bool
nstu [PrimeSieve]
more
                          ctLoop Integer
lac t
s (PrimeSieve
ps : [PrimeSieve]
more) = do
                              let !new :: Int
new = PrimeSieve -> Int
countAll PrimeSieve
ps
                              Integer -> t -> [PrimeSieve] -> ST s Integer
ctLoop (Integer
lac forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new) (t
sforall a. Num a => a -> a -> a
-t
1) [PrimeSieve]
more
                          ctLoop Integer
_ t
_ [] = forall a. HasCallStack => [Char] -> a
error [Char]
"Primes ended"
                      Int
new <- forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
si (Int
sieveBitsforall a. Num a => a -> a -> a
-Int
1) STUArray s Int Bool
stu
                      forall {t}.
(Eq t, Num t) =>
Integer -> t -> [PrimeSieve] -> ST s Integer
ctLoop (Int -> Integer
intToInteger Int
new) (Int64
cplforall a. Num a => a -> a -> a
-Int64
1) [PrimeSieve]
sves
    case [PrimeSieve]
psieves of
      (PS Integer
vO UArray Int Bool
ba : [PrimeSieve]
more) -> Int64
-> Int -> Int -> UArray Int Bool -> [PrimeSieve] -> ST s Integer
feed (forall a. Num a => Integer -> a
fromInteger Integer
vO) Int
0 Int
0 UArray Int Bool
ba [PrimeSieve]
more
      [PrimeSieve]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"No primes sieved"

calc :: Int64 -> Int64 -> Integer
calc :: Int64 -> Int64 -> Integer
calc Int64
lim Int64
plim = forall a. (forall s. ST s a) -> a
runST (forall s. Int64 -> Int64 -> ST s Integer
calcST Int64
lim Int64
plim)

calcST :: forall s. Int64 -> Int64 -> ST s Integer
calcST :: forall s. Int64 -> Int64 -> ST s Integer
calcST Int64
lim Int64
plim = do
    !STUArray s Int Bool
parr <- forall s. Integer -> ST s (STUArray s Int Bool)
sieveTo (Int64 -> Integer
int64ToInteger Int64
plim)
    (Int
plo,Int
phi) <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds STUArray s Int Bool
parr
    !Int
pct <- forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
plo Int
phi STUArray s Int Bool
parr
    !STUArray s Int Int64
ar1 <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
unsafeNewArray_ (Int
0,Int
endforall a. Num a => a -> a -> a
-Int
1)
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
ar1 Int
0 Int64
lim
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
ar1 Int
1 Int64
1
    !STUArray s Int Int64
ar2 <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
unsafeNewArray_ (Int
0,Int
endforall a. Num a => a -> a -> a
-Int
1)
    let go :: Int -> Int -> STUArray s Int Int64 -> STUArray s Int Int64 -> ST s Integer
        go :: Int
-> Int
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Integer
go Int
cap Int
pix STUArray s Int Int64
old STUArray s Int Int64
new
            | Int
pix forall a. Eq a => a -> a -> Bool
== Int
2  =   Int -> STUArray s Int Int64 -> ST s Integer
coll Int
cap STUArray s Int Int64
old
            | Bool
otherwise = do
                Bool
isp <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Bool
parr Int
pix
                if Bool
isp
                    then do
                        let !n :: Int64
n = forall a. Num a => Integer -> a
fromInteger (forall a. Num a => Int -> a
toPrim Int
pix)
                        !Int
ncap <- forall s.
Int
-> Int64
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Int
treat Int
cap Int64
n STUArray s Int Int64
old STUArray s Int Int64
new
                        Int
-> Int
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Integer
go Int
ncap (Int
pixforall a. Num a => a -> a -> a
-Int
1) STUArray s Int Int64
new STUArray s Int Int64
old
                    else Int
-> Int
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Integer
go Int
cap (Int
pixforall a. Num a => a -> a -> a
-Int
1) STUArray s Int Int64
old STUArray s Int Int64
new
        coll :: Int -> STUArray s Int Int64 -> ST s Integer
        coll :: Int -> STUArray s Int Int64 -> ST s Integer
coll Int
stop STUArray s Int Int64
ar =
            let cgo :: Integer -> Int -> m Integer
cgo !Integer
acc Int
i
                    | Int
i forall a. Ord a => a -> a -> Bool
< Int
stop  = do
                        !Int64
k <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
ar Int
i
                        !Int64
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
ar (Int
iforall a. Num a => a -> a -> a
+Int
1)
                        Integer -> Int -> m Integer
cgo (Integer
acc forall a. Num a => a -> a -> a
+ Int64 -> Integer
int64ToInteger Int64
vforall a. Num a => a -> a -> a
*Int64 -> Integer
cp6 Int64
k) (Int
iforall a. Num a => a -> a -> a
+Int
2)
                    | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return (Integer
accforall a. Num a => a -> a -> a
+Int -> Integer
intToInteger Int
pctforall a. Num a => a -> a -> a
+Integer
2)
            in forall {m :: * -> *}.
MArray (STUArray s) Int64 m =>
Integer -> Int -> m Integer
cgo Integer
0 Int
0
    Int
-> Int
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Integer
go Int
2 Int
start STUArray s Int Int64
ar1 STUArray s Int Int64
ar2
  where
    (Int
bt,Int
ri) = forall a. Integral a => a -> (Int, Int)
idxPr Int64
plim
    !start :: Int
start = Int
8forall a. Num a => a -> a -> a
*Int
bt forall a. Num a => a -> a -> a
+ Int
ri
    !size :: Int
size = Int64 -> Int
int64ToInt forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> a
integerSquareRoot Int64
lim forall a. Integral a => a -> a -> a
`quot` Int64
4
    !end :: Int
end = Int
2forall a. Num a => a -> a -> a
*Int
size

treat :: Int -> Int64 -> STUArray s Int Int64 -> STUArray s Int Int64 -> ST s Int
treat :: forall s.
Int
-> Int64
-> STUArray s Int Int64
-> STUArray s Int Int64
-> ST s Int
treat Int
end Int64
n STUArray s Int Int64
old STUArray s Int Int64
new = do
    Int
qi0 <- forall s. Int64 -> Int -> Int -> STUArray s Int Int64 -> ST s Int
locate Int64
n Int
0 (Int
end forall a. Integral a => a -> a -> a
`quot` Int
2 forall a. Num a => a -> a -> a
- Int
1) STUArray s Int Int64
old
    let collect :: Int64 -> Int64 -> Int -> m (Int64, Int)
collect Int64
stop !Int64
acc Int
ix
            | Int
ix forall a. Ord a => a -> a -> Bool
< Int
end  = do
                !Int64
k <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old Int
ix
                if Int64
k forall a. Ord a => a -> a -> Bool
< Int64
stop
                    then do
                        Int64
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old (Int
ixforall a. Num a => a -> a -> a
+Int
1)
                        Int64 -> Int64 -> Int -> m (Int64, Int)
collect Int64
stop (Int64
accforall a. Num a => a -> a -> a
-Int64
v) (Int
ixforall a. Num a => a -> a -> a
+Int
2)
                    else forall (m :: * -> *) a. Monad m => a -> m a
return (Int64
acc,Int
ix)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return (Int64
acc,Int
ix)
        goTreat :: Int -> Int -> Int -> ST s Int
goTreat !Int
wi !Int
ci Int
qi
            | Int
qi forall a. Ord a => a -> a -> Bool
< Int
end  = do
                !Int64
key <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old Int
qi
                !Int64
val <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old (Int
qiforall a. Num a => a -> a -> a
+Int
1)
                let !q0 :: Int64
q0 = Int64
key forall a. Integral a => a -> a -> a
`quot` Int64
n
                    !r0 :: Int
r0 = Int64 -> Int
int64ToInt (Int64
q0 forall a. Integral a => a -> a -> a
`rem` Int64
30030)
                    !nkey :: Int64
nkey = Int64
q0 forall a. Num a => a -> a -> a
- Int8 -> Int64
int8ToInt64 (UArray Int Int8
cpDfAr forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` Int
r0)
                    nk0 :: Int64
nk0 = Int64
q0 forall a. Num a => a -> a -> a
+ Int8 -> Int64
int8ToInt64 (UArray Int Int8
cpGpAr forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` (Int
r0forall a. Num a => a -> a -> a
+Int
1) forall a. Num a => a -> a -> a
+ Int8
1)
                    !nlim :: Int64
nlim = Int64
nforall a. Num a => a -> a -> a
*Int64
nk0
                (Int
wi1,Int
ci1) <- forall s.
Int
-> Int64
-> STUArray s Int Int64
-> Int
-> STUArray s Int Int64
-> Int
-> ST s (Int, Int)
copyTo Int
end Int64
nkey STUArray s Int Int64
old Int
ci STUArray s Int Int64
new Int
wi
                Int64
ckey <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old Int
ci1
                (!Int64
acc, !Int
ci2) <- if Int64
ckey forall a. Eq a => a -> a -> Bool
== Int64
nkey
                                  then do
                                    !Int64
ov <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old (Int
ci1forall a. Num a => a -> a -> a
+Int
1)
                                    forall (m :: * -> *) a. Monad m => a -> m a
return (Int64
ovforall a. Num a => a -> a -> a
-Int64
val,Int
ci1forall a. Num a => a -> a -> a
+Int
2)
                                  else forall (m :: * -> *) a. Monad m => a -> m a
return (-Int64
val,Int
ci1)
                (!Int64
tot, !Int
nqi) <- forall {m :: * -> *}.
MArray (STUArray s) Int64 m =>
Int64 -> Int64 -> Int -> m (Int64, Int)
collect Int64
nlim Int64
acc (Int
qiforall a. Num a => a -> a -> a
+Int
2)
                forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
new Int
wi1 Int64
nkey
                forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
new (Int
wi1forall a. Num a => a -> a -> a
+Int
1) Int64
tot
                Int -> Int -> Int -> ST s Int
goTreat (Int
wi1forall a. Num a => a -> a -> a
+Int
2) Int
ci2 Int
nqi
            | Bool
otherwise = forall s.
Int
-> STUArray s Int Int64
-> Int
-> STUArray s Int Int64
-> Int
-> ST s Int
copyRem Int
end STUArray s Int Int64
old Int
ci STUArray s Int Int64
new Int
wi
    Int -> Int -> Int -> ST s Int
goTreat Int
0 Int
0 Int
qi0

--------------------------------------------------------------------------------
--                               Auxiliaries                                  --
--------------------------------------------------------------------------------

locate :: Int64 -> Int -> Int -> STUArray s Int Int64 -> ST s Int
locate :: forall s. Int64 -> Int -> Int -> STUArray s Int Int64 -> ST s Int
locate Int64
p Int
low Int
high STUArray s Int Int64
arr = do
    let go :: Int -> Int -> m Int
go Int
lo Int
hi
          | Int
lo forall a. Ord a => a -> a -> Bool
< Int
hi     = do
            let !md :: Int
md = (Int
loforall a. Num a => a -> a -> a
+Int
hi) forall a. Integral a => a -> a -> a
`quot` Int
2
            Int64
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
arr (Int
2forall a. Num a => a -> a -> a
*Int
md)
            case forall a. Ord a => a -> a -> Ordering
compare Int64
p Int64
v of
                Ordering
LT -> Int -> Int -> m Int
go Int
lo Int
md
                Ordering
EQ -> forall (m :: * -> *) a. Monad m => a -> m a
return (Int
2forall a. Num a => a -> a -> a
*Int
md)
                Ordering
GT -> Int -> Int -> m Int
go (Int
mdforall a. Num a => a -> a -> a
+Int
1) Int
hi
          | Bool
otherwise   = forall (m :: * -> *) a. Monad m => a -> m a
return (Int
2forall a. Num a => a -> a -> a
*Int
lo)
    forall {m :: * -> *}.
MArray (STUArray s) Int64 m =>
Int -> Int -> m Int
go Int
low Int
high

{-# INLINE copyTo #-}
copyTo :: Int -> Int64 -> STUArray s Int Int64 -> Int
       -> STUArray s Int Int64 -> Int -> ST s (Int,Int)
copyTo :: forall s.
Int
-> Int64
-> STUArray s Int Int64
-> Int
-> STUArray s Int Int64
-> Int
-> ST s (Int, Int)
copyTo Int
end Int64
lim STUArray s Int Int64
old Int
oi STUArray s Int Int64
new Int
ni = do
    let go :: Int -> Int -> m (Int, Int)
go Int
ri Int
wi
            | Int
ri forall a. Ord a => a -> a -> Bool
< Int
end  = do
                Int64
ok <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old Int
ri
                if Int64
ok forall a. Ord a => a -> a -> Bool
< Int64
lim
                    then do
                        !Int64
ov <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old (Int
riforall a. Num a => a -> a -> a
+Int
1)
                        forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
new Int
wi Int64
ok
                        forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
new (Int
wiforall a. Num a => a -> a -> a
+Int
1) Int64
ov
                        Int -> Int -> m (Int, Int)
go (Int
riforall a. Num a => a -> a -> a
+Int
2) (Int
wiforall a. Num a => a -> a -> a
+Int
2)
                    else forall (m :: * -> *) a. Monad m => a -> m a
return (Int
wi,Int
ri)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return (Int
wi,Int
ri)
    forall {m :: * -> *}.
MArray (STUArray s) Int64 m =>
Int -> Int -> m (Int, Int)
go Int
oi Int
ni

{-# INLINE copyRem #-}
copyRem :: Int -> STUArray s Int Int64 -> Int -> STUArray s Int Int64 -> Int -> ST s Int
copyRem :: forall s.
Int
-> STUArray s Int Int64
-> Int
-> STUArray s Int Int64
-> Int
-> ST s Int
copyRem Int
end STUArray s Int Int64
old Int
oi STUArray s Int Int64
new Int
ni = do
    let go :: Int -> Int -> m Int
go Int
ri Int
wi
          | Int
ri forall a. Ord a => a -> a -> Bool
< Int
end    = do
            forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int64
old Int
ri forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int64
new Int
wi
            Int -> Int -> m Int
go (Int
riforall a. Num a => a -> a -> a
+Int
1) (Int
wiforall a. Num a => a -> a -> a
+Int
1)
          | Bool
otherwise   = forall (m :: * -> *) a. Monad m => a -> m a
return Int
wi
    forall {m :: * -> *}.
MArray (STUArray s) Int64 m =>
Int -> Int -> m Int
go Int
oi Int
ni

{-# INLINE cp6 #-}
cp6 :: Int64 -> Integer
cp6 :: Int64 -> Integer
cp6 Int64
k =
  case Int64
k forall a. Integral a => a -> a -> (a, a)
`quotRem` Int64
30030 of
    (Int64
q,Int64
r) -> Integer
5760forall a. Num a => a -> a -> a
*Int64 -> Integer
int64ToInteger Int64
q forall a. Num a => a -> a -> a
+
                Int16 -> Integer
int16ToInteger (UArray Int Int16
cpCtAr forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` Int64 -> Int
int64ToInt Int64
r)

cop :: Int64 -> Int64
cop :: Int64 -> Int64
cop Int64
m = Int64
m forall a. Num a => a -> a -> a
- Int8 -> Int64
int8ToInt64 (UArray Int Int8
cpDfAr forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
`unsafeAt` Int64 -> Int
int64ToInt (Int64
m forall a. Integral a => a -> a -> a
`rem` Int64
30030))


--------------------------------------------------------------------------------
--                           Ugly helper arrays                               --
--------------------------------------------------------------------------------

cpCtAr :: UArray Int Int16
cpCtAr :: UArray Int Int16
cpCtAr = forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int16
ar <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0,Int
30029) Int16
1
    let zilch :: Int -> Int -> m ()
zilch Int
s Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
30030 = forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int16
ar Int
i Int16
0 forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
zilch Int
s (Int
iforall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate :: Int16 -> Int -> m (STUArray s Int Int16)
accumulate Int16
ct Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
30030 = do
                Int16
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int16
ar Int
i
                let !ct' :: Int16
ct' = Int16
ctforall a. Num a => a -> a -> a
+Int16
v
                forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int16
ar Int
i Int16
ct'
                Int16 -> Int -> m (STUArray s Int Int16)
accumulate Int16
ct' (Int
iforall a. Num a => a -> a -> a
+Int
1)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return STUArray s Int Int16
ar
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
2 Int
0
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
6 Int
3
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
10 Int
5
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
14 Int
7
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
22 Int
11
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int -> Int -> m ()
zilch Int
26 Int
13
    forall {m :: * -> *}.
MArray (STUArray s) Int16 m =>
Int16 -> Int -> m (STUArray s Int Int16)
accumulate Int16
1 Int
2

cpDfAr :: UArray Int Int8
cpDfAr :: UArray Int Int8
cpDfAr = forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int8
ar <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0,Int
30029) Int8
0
    let note :: Int -> Int -> m ()
note Int
s Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
30029 = forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int8
ar Int
i Int8
1 forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
note Int
s (Int
iforall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate :: Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
d Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
30029 = do
                Int8
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int8
ar Int
i
                if Int8
v forall a. Eq a => a -> a -> Bool
== Int8
0
                    then Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
2 (Int
iforall a. Num a => a -> a -> a
+Int
2)
                    else do forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int8
ar Int
i Int8
d
                            Int8 -> Int -> m (STUArray s Int Int8)
accumulate (Int8
dforall a. Num a => a -> a -> a
+Int8
1) (Int
iforall a. Num a => a -> a -> a
+Int
1)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return STUArray s Int Int8
ar
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
2 Int
0
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
6 Int
3
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
10 Int
5
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
14 Int
7
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
22 Int
11
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
26 Int
13
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
2 Int
3

cpGpAr :: UArray Int Int8
cpGpAr :: UArray Int Int8
cpGpAr = forall i e. (forall s. ST s (STUArray s i e)) -> UArray i e
runSTUArray forall a b. (a -> b) -> a -> b
$ do
    STUArray s Int Int8
ar <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0,Int
30030) Int8
0
    forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int8
ar Int
30030 Int8
1
    let note :: Int -> Int -> m ()
note Int
s Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
30029 = forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int8
ar Int
i Int8
1 forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
note Int
s (Int
iforall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate :: Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
d Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
1     = forall (m :: * -> *) a. Monad m => a -> m a
return STUArray s Int Int8
ar
            | Bool
otherwise = do
                Int8
v <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Int8
ar Int
i
                if Int8
v forall a. Eq a => a -> a -> Bool
== Int8
0
                    then Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
2 (Int
iforall a. Num a => a -> a -> a
-Int
2)
                    else do forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite STUArray s Int Int8
ar Int
i Int8
d
                            Int8 -> Int -> m (STUArray s Int Int8)
accumulate (Int8
dforall a. Num a => a -> a -> a
+Int8
1) (Int
iforall a. Num a => a -> a -> a
-Int
1)
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
2 Int
0
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
6 Int
3
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
10 Int
5
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
14 Int
7
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
22 Int
11
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int -> Int -> m ()
note Int
26 Int
13
    forall {m :: * -> *}.
MArray (STUArray s) Int8 m =>
Int8 -> Int -> m (STUArray s Int Int8)
accumulate Int8
2 Int
30027

-------------------------------------------------------------------------------
-- Prime counting

rMASK :: Int
rMASK :: Int
rMASK = forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Num a => a -> a -> a
- Int
1

wSHFT :: (Bits a, Num a) => a
wSHFT :: forall a. (Bits a, Num a) => a
wSHFT = if forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Eq a => a -> a -> Bool
== Int
64 then a
6 else a
5

tOPB :: Int
tOPB :: Int
tOPB = forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) forall a. Bits a => a -> Int -> a
`shiftR` Int
1

tOPM :: (Bits a, Num a) => a
tOPM :: forall a. (Bits a, Num a) => a
tOPM = (a
1 forall a. Bits a => a -> Int -> a
`shiftL` Int
tOPB) forall a. Num a => a -> a -> a
- a
1

-- find the n-th set bit in a list of PrimeSieves,
-- aka find the (n+3)-rd prime
countToNth :: Int -> [PrimeSieve] -> Integer
countToNth :: Int -> [PrimeSieve] -> Integer
countToNth !Int
_ [] = forall a. HasCallStack => [Char] -> a
error [Char]
"countToNth: Prime stream ended prematurely"
countToNth !Int
n (PS Integer
v0 UArray Int Bool
bs : [PrimeSieve]
more) = Int -> Int -> Integer
go Int
n Int
0
  where
    wa :: UArray Int Word
    wa :: UArray Int Word
wa = forall a b. a -> b
unsafeCoerce UArray Int Bool
bs

    go :: Int -> Int -> Integer
go !Int
k Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== forall a b. (a, b) -> b
snd (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Int Word
wa)
      = Int -> [PrimeSieve] -> Integer
countToNth Int
k [PrimeSieve]
more
      | Bool
otherwise
      = let w :: Word
w = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
unsafeAt UArray Int Word
wa Int
i
            bc :: Int
bc = forall a. Bits a => a -> Int
popCount Word
w
        in if Int
bc forall a. Ord a => a -> a -> Bool
< Int
k
          then Int -> Int -> Integer
go (Int
kforall a. Num a => a -> a -> a
-Int
bc) (Int
iforall a. Num a => a -> a -> a
+Int
1)
          else let j :: Int
j = Int
bc forall a. Num a => a -> a -> a
- Int
k
                   px :: Int
px = Word -> Int -> Int -> Int
top Word
w Int
j Int
bc
               in Integer
v0 forall a. Num a => a -> a -> a
+ forall a. Num a => Int -> a
toPrim (Int
px forall a. Num a => a -> a -> a
+ (Int
i forall a. Bits a => a -> Int -> a
`shiftL` forall a. (Bits a, Num a) => a
wSHFT))

-- count all set bits in a chunk, do it wordwise for speed.
countAll :: PrimeSieve -> Int
countAll :: PrimeSieve -> Int
countAll (PS Integer
_ UArray Int Bool
bs) = Int -> Int -> Int
go Int
0 Int
0
  where
    wa :: UArray Int Word
    wa :: UArray Int Word
wa = forall a b. a -> b
unsafeCoerce UArray Int Bool
bs

    go :: Int -> Int -> Int
go !Int
ct Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== forall a b. (a, b) -> b
snd (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds UArray Int Word
wa)
      = Int
ct
      | Bool
otherwise
      = Int -> Int -> Int
go (Int
ct forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount (forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
unsafeAt UArray Int Word
wa Int
i)) (Int
iforall a. Num a => a -> a -> a
+Int
1)

-- Find the j-th highest of bc set bits in the Word w.
top :: Word -> Int -> Int -> Int
top :: Word -> Int -> Int -> Int
top Word
w Int
j Int
bc = forall {t}. (Num t, Bits t) => Int -> Int -> t -> Int -> t -> Int
go Int
0 Int
tOPB forall a. (Bits a, Num a) => a
tOPM Int
bn Word
w
    where
      !bn :: Int
bn = Int
bcforall a. Num a => a -> a -> a
-Int
j
      go :: Int -> Int -> t -> Int -> t -> Int
go !Int
_ Int
_ !t
_ !Int
_ t
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Too few bits set"
      go Int
bs Int
0 t
_ Int
_ t
wd = if t
wd forall a. Bits a => a -> a -> a
.&. t
1 forall a. Eq a => a -> a -> Bool
== t
0 then forall a. HasCallStack => [Char] -> a
error [Char]
"Too few bits, shift 0" else Int
bs
      go Int
bs Int
a t
msk Int
ix t
wd =
        case forall a. Bits a => a -> Int
popCount (t
wd forall a. Bits a => a -> a -> a
.&. t
msk) of
          Int
lc | Int
lc forall a. Ord a => a -> a -> Bool
< Int
ix  -> Int -> Int -> t -> Int -> t -> Int
go (Int
bsforall a. Num a => a -> a -> a
+Int
a) Int
a t
msk (Int
ixforall a. Num a => a -> a -> a
-Int
lc) (t
wd forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
a)
             | Bool
otherwise ->
               let !na :: Int
na = Int
a forall a. Bits a => a -> Int -> a
`shiftR` Int
1
               in Int -> Int -> t -> Int -> t -> Int
go Int
bs Int
na (t
msk forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
na) Int
ix t
wd

-- count set bits between two indices (inclusive)
-- start and end must both be valid indices and start <= end
countFromTo :: Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo :: forall s. Int -> Int -> STUArray s Int Bool -> ST s Int
countFromTo Int
start Int
end STUArray s Int Bool
ba = do
    STUArray s Int Word
wa <- (forall s ix a b. STUArray s ix a -> ST s (STUArray s ix b)
castSTUArray :: STUArray s Int Bool -> ST s (STUArray s Int Word)) STUArray s Int Bool
ba
    let !sb :: Int
sb = Int
start forall a. Bits a => a -> Int -> a
`shiftR` forall a. (Bits a, Num a) => a
wSHFT
        !si :: Int
si = Int
start forall a. Bits a => a -> a -> a
.&. Int
rMASK
        !eb :: Int
eb = Int
end forall a. Bits a => a -> Int -> a
`shiftR` forall a. (Bits a, Num a) => a
wSHFT
        !ei :: Int
ei = Int
end forall a. Bits a => a -> a -> a
.&. Int
rMASK
        count :: Int -> Int -> m Int
count !Int
acc Int
i
            | Int
i forall a. Eq a => a -> a -> Bool
== Int
eb = do
                Word
w <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Word
wa Int
i
                forall (m :: * -> *) a. Monad m => a -> m a
return (Int
acc forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount (Word
w forall a. Bits a => a -> Int -> a
`shiftL` (Int
rMASK forall a. Num a => a -> a -> a
- Int
ei)))
            | Bool
otherwise = do
                Word
w <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Word
wa Int
i
                Int -> Int -> m Int
count (Int
acc forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount Word
w) (Int
iforall a. Num a => a -> a -> a
+Int
1)
    if Int
sb forall a. Ord a => a -> a -> Bool
< Int
eb
      then do
          Word
w <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Word
wa Int
sb
          forall {m :: * -> *}.
MArray (STUArray s) Word m =>
Int -> Int -> m Int
count (forall a. Bits a => a -> Int
popCount (Word
w forall a. Bits a => a -> Int -> a
`shiftR` Int
si)) (Int
sbforall a. Num a => a -> a -> a
+Int
1)
      else do
          Word
w <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> m e
unsafeRead STUArray s Int Word
wa Int
sb
          let !w1 :: Word
w1 = Word
w forall a. Bits a => a -> Int -> a
`shiftR` Int
si
          forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Bits a => a -> Int
popCount (Word
w1 forall a. Bits a => a -> Int -> a
`shiftL` (Int
rMASK forall a. Num a => a -> a -> a
- Int
ei forall a. Num a => a -> a -> a
+ Int
si)))