-- |
-- 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.Bit
import Data.Bits
import Data.Int
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU

-- | 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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
primeCountMaxArg = [Char] -> Integer
forall a. HasCallStack => [Char] -> a
error ([Char] -> Integer) -> [Char] -> Integer
forall a b. (a -> b) -> a -> b
$ [Char]
"primeCount: can't handle bound " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2     = Integer
0
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1000  = Int -> Integer
intToInteger (Int -> Integer) -> (Integer -> Int) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Integer] -> Int) -> (Integer -> [Integer]) -> Integer -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
n) ([Integer] -> [Integer])
-> (Integer -> [Integer]) -> Integer -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Prime Integer -> Integer) -> [Prime Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Prime Integer -> Integer
forall a. Prime a -> a
unPrime ([Prime Integer] -> [Integer])
-> (Integer -> [Prime Integer]) -> Integer -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimeSieve -> [Prime Integer]
forall a. Integral a => PrimeSieve -> [Prime a]
primeList (PrimeSieve -> [Prime Integer])
-> (Integer -> PrimeSieve) -> Integer -> [Prime Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> PrimeSieve
primeSieve (Integer -> Integer) -> Integer -> Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
242 Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
30000 = (forall s. ST s Integer) -> Integer
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Integer) -> Integer)
-> (forall s. ST s Integer) -> Integer
forall a b. (a -> b) -> a -> b
$ do
        ba <- Integer -> ST s (MVector s Bit)
forall s. Integer -> ST s (MVector s Bit)
sieveTo Integer
n
        let (s, e) = (0, MU.length ba - 1)
        ct <- countFromTo s e ba
        return (intToInteger $ ct+3)
    | Bool
otherwise =
        let !ub :: Int64
ub = Int64 -> Int64
cop (Int64 -> Int64) -> Int64 -> Int64
forall a b. (a -> b) -> a -> b
$ Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
n
            !sr :: Int64
sr = Int64 -> Int64
forall a. Integral a => a -> a
integerSquareRoot Int64
ub
            !cr :: Int64
cr = Int64 -> Int64
forall a. Integral a => a -> a
nxtEnd (Int64 -> Int64) -> Int64 -> Int64
forall a b. (a -> b) -> a -> b
$ Int64 -> Int64
forall a. Integral a => a -> a
integerCubeRoot Int64
ub Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
15
            nxtEnd :: a -> a
nxtEnd a
k = a
k a -> a -> a
forall a. Num a => a -> a -> a
- (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
30) a -> a -> a
forall a. Num a => a -> a -> a
+ a
31
            !phn1 :: Integer
phn1 = Int64 -> Int64 -> Integer
calc Int64
ub Int64
cr
            !cs :: Int64
cs = Int64
crInt64 -> Int64 -> Int64
forall 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 Integer -> Integer -> Integer
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 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
2
nthPrime Int
2 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
3
nthPrime Int
3 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
5
nthPrime Int
4 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
7
nthPrime Int
5 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
11
nthPrime Int
6 = Integer -> Prime Integer
forall a. a -> Prime a
Prime Integer
13
nthPrime Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
    = [Char] -> Prime Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"Prime indexing starts at 1"
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
200000
    = Integer -> Prime Integer
forall a. a -> Prime a
Prime (Integer -> Prime Integer) -> Integer -> Prime Integer
forall a b. (a -> b) -> a -> b
$ Int -> [PrimeSieve] -> Integer
countToNth (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) [Integer -> PrimeSieve
primeSieve (Integer
p0 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
p0 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
32 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
37)]
    | Integer
p0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
    = [Char] -> Prime Integer
forall a. HasCallStack => [Char] -> a
error ([Char] -> Prime Integer) -> [Char] -> Prime Integer
forall a b. (a -> b) -> a -> b
$ [Char]
"nthPrime: index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is too large to handle"
    | Int
miss Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    = Integer -> Prime Integer
forall a. a -> Prime a
Prime (Integer -> Prime Integer) -> Integer -> Prime Integer
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Integer
tooLow  Int
n (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
p0) Int
miss
    | Bool
otherwise
    = Integer -> Prime Integer
forall a. a -> Prime a
Prime (Integer -> Prime Integer) -> Integer -> Prime Integer
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Integer
tooHigh Int
n (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
p0) (Int -> Int
forall a. Num a => a -> a
negate Int
miss)
      where
        p0 :: Integer
p0 = Integer -> Integer
nthPrimeApprox (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n)
        miss :: Int
miss = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Integer -> Int
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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
  = [Char] -> Integer
forall a. HasCallStack => [Char] -> a
error ([Char] -> Integer) -> [Char] -> Integer
forall a b. (a -> b) -> a -> b
$ [Char]
"nthPrime: index " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is too large to handle"
  | Bool
goodEnough
  = Int -> Int -> Integer
lowSieve Int
p0 Int
shortage
  | Int
c1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
  = Int -> Int -> Integer
lowSieve (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
p1) (Int
nInt -> Int -> Int
forall 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 = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
intToDouble Int
p0 :: Double))
    est :: Integer
est = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
shortage Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
gap
    p1 :: Integer
p1  = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
p0 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
est
    goodEnough :: Bool
goodEnough = Integer
3Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
estInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
estInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
est Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
p1Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
p1    -- a second counting would be more work than sieving
    c1 :: Int
c1  = Integer -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
  = Int -> Int -> Integer
lowSieve Int
b (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
c)
  | Bool
otherwise
  = Int -> Int -> Int -> Integer
tooHigh Int
n Int
b (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n)
  where
    gap :: Int
gap = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
intToDouble Int
p0 :: Double))
    b :: Int
b = Int
p0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
surplus Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
gap Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
11) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
10
    c :: Int
c = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
primeCount (Int -> Integer
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
missInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rep) [PrimeSieve]
psieves
      where
        strt :: Int
strt = Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
a Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
1)
        psieves :: [PrimeSieve]
psieves@(PS Integer
vO Vector Bit
ba:[PrimeSieve]
_) = Integer -> [PrimeSieve]
psieveFrom (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
strt)
        rep :: Int
rep | Integer
o0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0    = Int
0
            | Bool
otherwise = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int
1 | Int
i <- [Int
0 .. Int
r2], Bit -> Bool
unBit (Vector Bit
ba Vector Bit -> Int -> Bit
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int
i)]
              where
                o0 :: Integer
o0 = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
strt Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
vO Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
9   -- (strt - 2) - v0 - 7
                r0 :: Int
r0 = Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
o0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
30
                r1 :: Int
r1 = Int
r0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
3
                r2 :: Int
r2 = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
7 (if Int
r1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
5 then Int
r1Int -> Int -> Int
forall 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 s. ST s Integer) -> Integer
forall a. (forall s. ST s a) -> a
runST (Int64 -> Int64 -> Int64 -> ST s Integer
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 = Int64 -> Int64
forall a. Integral a => a -> a
approxPrimeCount Int64
sr
        picr :: Int64
picr = Int64 -> Int64
forall a. Integral a => a -> a
approxPrimeCount Int64
cr
        diff :: Int64
diff = Int64
pisr Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
picr
        size :: Int
size = Int64 -> Int
int64ToInt (Int64
diff Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
diff Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`quot` Int64
50) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
30
    store <- Int -> ST s (MVector (PrimState (ST s)) Int64)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.unsafeNew Int
size :: ST s (MU.MVector s Int64)
    let feed :: Int64 -> Int -> Int -> U.Vector Bit -> [PrimeSieve] -> ST s Integer
        feed Int64
voff !Int
wi !Int
ri Vector Bit
uar [PrimeSieve]
sves
          | Int
ri Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sieveBits = case [PrimeSieve]
sves of
                                (PS Integer
vO Vector Bit
ba : [PrimeSieve]
more) -> Int64 -> Int -> Int -> Vector Bit -> [PrimeSieve] -> ST s Integer
feed (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
vO) Int
wi Int
0 Vector Bit
ba [PrimeSieve]
more
                                [PrimeSieve]
_ -> [Char] -> ST s Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"prime stream ended prematurely"
          | Int64
pval Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
sr   = do
              stu <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
uar
              eat 0 0 voff (wi-1) ri stu sves
          | Bit -> Bool
unBit (Vector Bit
uar Vector Bit -> Int -> Bit
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int
ri) = do
              MVector (PrimState (ST s)) Int64 -> Int -> Int64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Int64
MVector (PrimState (ST s)) Int64
store Int
wi (Int64
ub Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`quot` Int64
pval)
              Int64 -> Int -> Int -> Vector Bit -> [PrimeSieve] -> ST s Integer
feed Int64
voff (Int
wiInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
riInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Vector Bit
uar [PrimeSieve]
sves
          | Bool
otherwise = Int64 -> Int -> Int -> Vector Bit -> [PrimeSieve] -> ST s Integer
feed Int64
voff Int
wi (Int
riInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Vector Bit
uar [PrimeSieve]
sves
            where
              pval :: Int64
pval = Int64
voff Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int -> Int64
forall a. Num a => Int -> a
toPrim Int
ri
        eat :: Integer -> Integer -> Int64 -> Int -> Int -> MU.MVector s Bit -> [PrimeSieve] -> ST s Integer
        eat !Integer
acc !Integer
btw Int64
voff !Int
wi !Int
si MVector s Bit
stu [PrimeSieve]
sves
            | Int
si Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sieveBits =
                case [PrimeSieve]
sves of
                  [] -> [Char] -> ST s Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"Premature end of prime stream"
                  (PS Integer
vO Vector Bit
ba : [PrimeSieve]
more) -> do
                      nstu <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
ba
                      eat acc btw (fromInteger vO) wi 0 nstu more
            | Int
wi Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0    = Integer -> ST s Integer
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
acc
            | Bool
otherwise = do
                qb <- MVector (PrimState (ST s)) Int64 -> Int -> ST s Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState (ST s)) Int64
store Int
wi
                let dist = Int64
qb Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
voff Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
7
                if dist < intToInt64 sieveRange
                  then do
                      let (b,j) = idxPr (dist+7)
                          !li = (Int
b Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
j
                      new <- if li < si then return 0 else countFromTo si li stu
                      let nbtw = Integer
btw Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
                      eat (acc+nbtw) nbtw voff (wi-1) (li+1) stu sves
                  else do
                      let (cpl,fds) = dist `quotRem` intToInt64 sieveRange
                          (b,j) = idxPr (fds+7)
                          !li = (Int
b Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
j
                          ctLoop !Integer
lac t
0 (PS Integer
vO Vector Bit
ba : [PrimeSieve]
more) = do
                              nstu <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
ba
                              new <- countFromTo 0 li nstu
                              let nbtw = Integer
btw Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
lac Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new
                              eat (acc+nbtw) nbtw (integerToInt64 vO) (wi-1) (li+1) nstu 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 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
intToInteger Int
new) (t
st -> t -> t
forall a. Num a => a -> a -> a
-t
1) [PrimeSieve]
more
                          ctLoop Integer
_ t
_ [] = [Char] -> ST s Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"Primes ended"
                      new <- countFromTo si (sieveBits-1) stu
                      ctLoop (intToInteger new) (cpl-1) sves
    case psieves of
      (PS Integer
vO Vector Bit
ba : [PrimeSieve]
more) -> Int64 -> Int -> Int -> Vector Bit -> [PrimeSieve] -> ST s Integer
feed (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
vO) Int
0 Int
0 Vector Bit
ba [PrimeSieve]
more
      [PrimeSieve]
_ -> [Char] -> ST s Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"No primes sieved"

calc :: Int64 -> Int64 -> Integer
calc :: Int64 -> Int64 -> Integer
calc Int64
lim Int64
plim = (forall s. ST s Integer) -> Integer
forall a. (forall s. ST s a) -> a
runST (Int64 -> Int64 -> ST s Integer
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
    !parr <- Integer -> ST s (MVector s Bit)
forall s. Integer -> ST s (MVector s Bit)
sieveTo (Int64 -> Integer
int64ToInteger Int64
plim)
    let (plo, phi) = (0, MU.length parr - 1)
    !pct <- countFromTo plo phi parr
    !ar1 <- MU.unsafeNew end
    MU.unsafeWrite ar1 0 lim
    MU.unsafeWrite ar1 1 1
    !ar2 <- MU.unsafeNew end
    let go :: Int -> Int -> MU.MVector s Int64 -> MU.MVector s Int64 -> ST s Integer
        go Int
cap Int
pix MVector s Int64
old MVector s Int64
new
            | Int
pix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2  =   Int -> MVector s Int64 -> ST s Integer
coll Int
cap MVector s Int64
old
            | Bool
otherwise = do
                Bit isp <- MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Bit
MVector (PrimState (ST s)) Bit
parr Int
pix
                if isp
                    then do
                        let !n = Integer -> Int64
forall a. Num a => Integer -> a
fromInteger (Int -> Integer
forall a. Num a => Int -> a
toPrim Int
pix)
                        !ncap <- treat cap n old new
                        go ncap (pix-1) new old
                    else go cap (pix-1) old new
        coll :: Int -> MU.MVector s Int64 -> ST s Integer
        coll Int
stop MVector s Int64
ar =
            let cgo :: Integer -> Int -> m Integer
cgo !Integer
acc Int
i
                    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
stop  = do
                        !k <- MVector (PrimState m) Int64 -> Int -> m Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState m) Int64
ar Int
i
                        !v <- MU.unsafeRead ar (i+1)
                        cgo (acc + int64ToInteger v*cp6 k) (i+2)
                    | Bool
otherwise = Integer -> m Integer
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer
accInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Int -> Integer
intToInteger Int
pctInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
2)
            in Integer -> Int -> ST s Integer
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Integer -> Int -> m Integer
cgo Integer
0 Int
0
    go 2 start ar1 ar2
  where
    (Int
bt,Int
ri) = Int64 -> (Int, Int)
forall a. Integral a => a -> (Int, Int)
idxPr Int64
plim
    !start :: Int
start = Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
bt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ri
    !size :: Int
size = Int64 -> Int
int64ToInt (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ Int64 -> Int64
forall a. Integral a => a -> a
integerSquareRoot Int64
lim Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`quot` Int64
4
    !end :: Int
end = Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
size

treat :: Int -> Int64 -> MU.MVector s Int64 -> MU.MVector s Int64 -> ST s Int
treat :: forall s.
Int -> Int64 -> MVector s Int64 -> MVector s Int64 -> ST s Int
treat Int
end Int64
n MVector s Int64
old MVector s Int64
new = do
    qi0 <- Int64 -> Int -> Int -> MVector s Int64 -> ST s Int
forall s. Int64 -> Int -> Int -> MVector s Int64 -> ST s Int
locate Int64
n Int
0 (Int
end Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) MVector s Int64
old
    let collect Int64
stop !Int64
acc Int
ix
            | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end  = do
                !k <- MVector (PrimState m) Int64 -> Int -> m Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState m) Int64
old Int
ix
                if k < stop
                    then do
                        v <- MU.unsafeRead old (ix+1)
                        collect stop (acc-v) (ix+2)
                    else return (acc,ix)
            | Bool
otherwise = (Int64, Int) -> m (Int64, Int)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int64
acc,Int
ix)
        goTreat !Int
wi !Int
ci Int
qi
            | Int
qi Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end  = do
                !key <- MVector (PrimState (ST s)) Int64 -> Int -> ST s Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState (ST s)) Int64
old Int
qi
                !val <- MU.unsafeRead old (qi+1)
                let !q0 = Int64
key Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`quot` Int64
n
                    !r0 = Int64 -> Int
int64ToInt (Int64
q0 Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`rem` Int64
30030)
                    !nkey = Int64
q0 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int8 -> Int64
int8ToInt64 (Vector Int8
cpDfAr Vector Int8 -> Int -> Int8
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int
r0)
                    nk0 = Int64
q0 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int8 -> Int64
int8ToInt64 (Vector Int8
cpGpAr Vector Int8 -> Int -> Int8
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` (Int
r0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int8 -> Int8 -> Int8
forall a. Num a => a -> a -> a
+ Int8
1)
                    !nlim = Int64
nInt64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
*Int64
nk0
                (wi1,ci1) <- copyTo end nkey old ci new wi
                ckey <- MU.unsafeRead old ci1
                (!acc, !ci2) <- if ckey == nkey
                                  then do
                                    !ov <- MU.unsafeRead old (ci1+1)
                                    return (ov-val,ci1+2)
                                  else return (-val,ci1)
                (!tot, !nqi) <- collect nlim acc (qi+2)
                MU.unsafeWrite new wi1 nkey
                MU.unsafeWrite new (wi1+1) tot
                goTreat (wi1+2) ci2 nqi
            | Bool
otherwise = Int -> MVector s Int64 -> Int -> MVector s Int64 -> Int -> ST s Int
forall s.
Int -> MVector s Int64 -> Int -> MVector s Int64 -> Int -> ST s Int
copyRem Int
end MVector s Int64
old Int
ci MVector s Int64
new Int
wi
    goTreat 0 0 qi0

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

locate :: Int64 -> Int -> Int -> MU.MVector s Int64 -> ST s Int
locate :: forall s. Int64 -> Int -> Int -> MVector s Int64 -> ST s Int
locate Int64
p Int
low Int
high MVector s Int64
arr = do
    let go :: Int -> Int -> m Int
go Int
lo Int
hi
          | Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
hi     = do
            let !md :: Int
md = (Int
loInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
hi) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2
            v <- MVector (PrimState m) Int64 -> Int -> m Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState m) Int64
arr (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
md)
            case compare p v of
                Ordering
LT -> Int -> Int -> m Int
go Int
lo Int
md
                Ordering
EQ -> Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
md)
                Ordering
GT -> Int -> Int -> m Int
go (Int
mdInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
hi
          | Bool
otherwise   = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
lo)
    Int -> Int -> ST s Int
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> Int -> m Int
go Int
low Int
high

{-# INLINE copyTo #-}
copyTo :: Int -> Int64 -> MU.MVector s Int64 -> Int
       -> MU.MVector s Int64 -> Int -> ST s (Int,Int)
copyTo :: forall s.
Int
-> Int64
-> MVector s Int64
-> Int
-> MVector s Int64
-> Int
-> ST s (Int, Int)
copyTo Int
end Int64
lim MVector s Int64
old Int
oi MVector s Int64
new Int
ni = do
    let go :: Int -> Int -> m (Int, Int)
go Int
ri Int
wi
            | Int
ri Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
end  = do
                ok <- MVector (PrimState m) Int64 -> Int -> m Int64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int64
MVector (PrimState m) Int64
old Int
ri
                if ok < lim
                    then do
                        !ov <- MU.unsafeRead old (ri+1)
                        MU.unsafeWrite new wi ok
                        MU.unsafeWrite new (wi+1) ov
                        go (ri+2) (wi+2)
                    else return (wi,ri)
            | Bool
otherwise = (Int, Int) -> m (Int, Int)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
wi,Int
ri)
    Int -> Int -> ST s (Int, Int)
forall {m :: * -> *}.
(PrimState m ~ s, PrimMonad m) =>
Int -> Int -> m (Int, Int)
go Int
oi Int
ni

{-# INLINE copyRem #-}
copyRem :: Int -> MU.MVector s Int64 -> Int -> MU.MVector s Int64 -> Int -> ST s Int
copyRem :: forall s.
Int -> MVector s Int64 -> Int -> MVector s Int64 -> Int -> ST s Int
copyRem Int
end MVector s Int64
old Int
oi MVector s Int64
new Int
ni = do
  let len :: Int
len = Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
oi
  MVector (PrimState (ST s)) Int64
-> MVector (PrimState (ST s)) Int64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
MU.copy (Int -> Int -> MVector s Int64 -> MVector s Int64
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
ni Int
len MVector s Int64
new) (Int -> Int -> MVector s Int64 -> MVector s Int64
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
oi Int
len MVector s Int64
old)
  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
ni Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len

{-# INLINE cp6 #-}
cp6 :: Int64 -> Integer
cp6 :: Int64 -> Integer
cp6 Int64
k =
  case Int64
k Int64 -> Int64 -> (Int64, Int64)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int64
30030 of
    (Int64
q,Int64
r) -> Integer
5760Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Int64 -> Integer
int64ToInteger Int64
q Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+
                Int16 -> Integer
int16ToInteger (Vector Int16
cpCtAr Vector Int16 -> Int -> Int16
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int64 -> Int
int64ToInt Int64
r)

cop :: Int64 -> Int64
cop :: Int64 -> Int64
cop Int64
m = Int64
m Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int8 -> Int64
int8ToInt64 (Vector Int8
cpDfAr Vector Int8 -> Int -> Int8
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int64 -> Int
int64ToInt (Int64
m Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`rem` Int64
30030))


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

cpCtAr :: U.Vector Int16
cpCtAr :: Vector Int16
cpCtAr = (forall s. ST s (Vector Int16)) -> Vector Int16
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int16)) -> Vector Int16)
-> (forall s. ST s (Vector Int16)) -> Vector Int16
forall a b. (a -> b) -> a -> b
$ do
    ar <- Int -> Int16 -> ST s (MVector (PrimState (ST s)) Int16)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
30030 Int16
1
    let zilch Int
s Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
30030 = MVector (PrimState m) Int16 -> Int -> Int16 -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Int16
MVector (PrimState m) Int16
ar Int
i Int16
0 m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
zilch Int
s (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate Int16
ct Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
30030 = do
                v <- MVector (PrimState m) Int16 -> Int -> m Int16
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int16
MVector (PrimState m) Int16
ar Int
i
                let !ct' = Int16
ctInt16 -> Int16 -> Int16
forall a. Num a => a -> a -> a
+Int16
v
                MU.unsafeWrite ar i ct'
                accumulate ct' (i+1)
            | Bool
otherwise = MVector s Int16 -> m (MVector s Int16)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Int16
ar
    zilch 2 0
    zilch 6 3
    zilch 10 5
    zilch 14 7
    zilch 22 11
    zilch 26 13
    vec <- accumulate 1 2
    U.unsafeFreeze vec

cpDfAr :: U.Vector Int8
cpDfAr :: Vector Int8
cpDfAr = (forall s. ST s (Vector Int8)) -> Vector Int8
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int8)) -> Vector Int8)
-> (forall s. ST s (Vector Int8)) -> Vector Int8
forall a b. (a -> b) -> a -> b
$ do
    ar <- Int -> Int8 -> ST s (MVector (PrimState (ST s)) Int8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
30030 Int8
0
    let note Int
s Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
30029 = MVector (PrimState m) Int8 -> Int -> Int8 -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Int8
MVector (PrimState m) Int8
ar Int
i Int8
1 m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
note Int
s (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate Int8
d Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
30029 = do
                v <- MVector (PrimState m) Int8 -> Int -> m Int8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int8
MVector (PrimState m) Int8
ar Int
i
                if v == 0
                    then accumulate 2 (i+2)
                    else do MU.unsafeWrite ar i d
                            accumulate (d+1) (i+1)
            | Bool
otherwise = MVector s Int8 -> m (MVector s Int8)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Int8
ar
    note 2 0
    note 6 3
    note 10 5
    note 14 7
    note 22 11
    note 26 13
    vec <- accumulate 2 3
    U.unsafeFreeze vec

cpGpAr :: U.Vector Int8
cpGpAr :: Vector Int8
cpGpAr = (forall s. ST s (Vector Int8)) -> Vector Int8
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int8)) -> Vector Int8)
-> (forall s. ST s (Vector Int8)) -> Vector Int8
forall a b. (a -> b) -> a -> b
$ do
    ar <- Int -> Int8 -> ST s (MVector (PrimState (ST s)) Int8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
30031 Int8
0
    MU.unsafeWrite ar 30030 1
    let note Int
s Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
30029 = MVector (PrimState m) Int8 -> Int -> Int8 -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Int8
MVector (PrimState m) Int8
ar Int
i Int8
1 m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> m ()
note Int
s (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
s)
            | Bool
otherwise = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        accumulate Int8
d Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1     = MVector s Int8 -> m (MVector s Int8)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Int8
ar
            | Bool
otherwise = do
                v <- MVector (PrimState m) Int8 -> Int -> m Int8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Int8
MVector (PrimState m) Int8
ar Int
i
                if v == 0
                    then accumulate 2 (i-2)
                    else do MU.unsafeWrite ar i d
                            accumulate (d+1) (i-1)
    note 2 0
    note 6 3
    note 10 5
    note 14 7
    note 22 11
    note 26 13
    vec <- accumulate 2 30027
    U.unsafeFreeze vec

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

-- 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
_ [] = [Char] -> Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"countToNth: Prime stream ended prematurely"
countToNth !Int
n (PS Integer
v0 Vector Bit
bs : [PrimeSieve]
more) = case Bit -> Int -> Vector Bit -> Maybe Int
nthBitIndex (Bool -> Bit
Bit Bool
True) Int
n Vector Bit
bs of
  Just Int
i -> Integer
v0 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Num a => Int -> a
toPrim Int
i
  Maybe Int
Nothing -> Int -> [PrimeSieve] -> Integer
countToNth (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Bit -> Int
countBits Vector Bit
bs) [PrimeSieve]
more

-- count all set bits in a chunk, do it wordwise for speed.
countAll :: PrimeSieve -> Int
countAll :: PrimeSieve -> Int
countAll (PS Integer
_ Vector Bit
bs) = Vector Bit -> Int
countBits Vector Bit
bs

-- count set bits between two indices (inclusive)
-- start and end must both be valid indices and start <= end
countFromTo :: Int -> Int -> MU.MVector s Bit -> ST s Int
countFromTo :: forall s. Int -> Int -> MVector s Bit -> ST s Int
countFromTo Int
start Int
end =
  (Vector Bit -> Int) -> ST s (Vector Bit) -> ST s Int
forall a b. (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Vector Bit -> Int
countBits (ST s (Vector Bit) -> ST s Int)
-> (MVector s Bit -> ST s (Vector Bit))
-> MVector s Bit
-> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVector s Bit -> ST s (Vector Bit)
MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (MVector s Bit -> ST s (Vector Bit))
-> (MVector s Bit -> MVector s Bit)
-> MVector s Bit
-> ST s (Vector Bit)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> MVector s Bit -> MVector s Bit
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
start (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)