{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Math.NumberTheory.ArithmeticFunctions.SieveBlock
( runFunctionOverBlock
, SieveBlockConfig(..)
, multiplicativeSieveBlockConfig
, additiveSieveBlockConfig
, sieveBlock
, sieveBlockUnboxed
, sieveBlockMoebius
) where
import Control.Monad (forM_, when)
import Control.Monad.ST (runST)
import Data.Bits
import Data.Coerce
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import GHC.Exts
import Math.NumberTheory.ArithmeticFunctions.Class
import Math.NumberTheory.ArithmeticFunctions.Moebius (Moebius, sieveBlockMoebius)
import Math.NumberTheory.Logarithms (wordLog2, integerLogBase')
import Math.NumberTheory.Primes
import Math.NumberTheory.Primes.Types
import Math.NumberTheory.Roots (integerSquareRoot)
import Math.NumberTheory.Utils.FromIntegral (wordToInt, intToWord)
data SieveBlockConfig a = SieveBlockConfig
{ SieveBlockConfig a -> a
sbcEmpty :: a
, SieveBlockConfig a -> Prime Word -> Word -> a
sbcFunctionOnPrimePower :: Prime Word -> Word -> a
, SieveBlockConfig a -> a -> a -> a
sbcAppend :: a -> a -> a
}
multiplicativeSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig :: (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig Prime Word -> Word -> a
f = SieveBlockConfig :: forall a.
a
-> (Prime Word -> Word -> a) -> (a -> a -> a) -> SieveBlockConfig a
SieveBlockConfig
{ sbcEmpty :: a
sbcEmpty = a
1
, sbcFunctionOnPrimePower :: Prime Word -> Word -> a
sbcFunctionOnPrimePower = Prime Word -> Word -> a
f
, sbcAppend :: a -> a -> a
sbcAppend = a -> a -> a
forall a. Num a => a -> a -> a
(*)
}
additiveSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig :: (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig Prime Word -> Word -> a
f = SieveBlockConfig :: forall a.
a
-> (Prime Word -> Word -> a) -> (a -> a -> a) -> SieveBlockConfig a
SieveBlockConfig
{ sbcEmpty :: a
sbcEmpty = a
0
, sbcFunctionOnPrimePower :: Prime Word -> Word -> a
sbcFunctionOnPrimePower = Prime Word -> Word -> a
f
, sbcAppend :: a -> a -> a
sbcAppend = a -> a -> a
forall a. Num a => a -> a -> a
(+)
}
runFunctionOverBlock
:: ArithmeticFunction Word a
-> Word
-> Word
-> V.Vector a
runFunctionOverBlock :: ArithmeticFunction Word a -> Word -> Word -> Vector a
runFunctionOverBlock (ArithmeticFunction Prime Word -> Word -> m
f m -> a
g) = ((m -> a) -> Vector m -> Vector a
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map m -> a
g (Vector m -> Vector a) -> (Word -> Vector m) -> Word -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Word -> Vector m) -> Word -> Vector a)
-> (Word -> Word -> Vector m) -> Word -> Word -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SieveBlockConfig m -> Word -> Word -> Vector m
forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock SieveBlockConfig :: forall a.
a
-> (Prime Word -> Word -> a) -> (a -> a -> a) -> SieveBlockConfig a
SieveBlockConfig
{ sbcEmpty :: m
sbcEmpty = m
forall a. Monoid a => a
mempty
, sbcAppend :: m -> m -> m
sbcAppend = m -> m -> m
forall a. Monoid a => a -> a -> a
mappend
, sbcFunctionOnPrimePower :: Prime Word -> Word -> m
sbcFunctionOnPrimePower = (Prime Word -> Word -> m) -> Prime Word -> Word -> m
coerce Prime Word -> Word -> m
f
}
sieveBlock
:: forall v a.
G.Vector v a
=> SieveBlockConfig a
-> Word
-> Word
-> v a
sieveBlock :: SieveBlockConfig a -> Word -> Word -> v a
sieveBlock SieveBlockConfig a
_ Word
_ Word
0 = v a
forall (v :: * -> *) a. Vector v a => v a
G.empty
sieveBlock (SieveBlockConfig a
empty Prime Word -> Word -> a
f a -> a -> a
append) !Word
lowIndex' Word
len' = (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
let lowIndex :: Int
lowIndex :: Int
lowIndex = Word -> Int
wordToInt Word
lowIndex'
len :: Int
len :: Int
len = Word -> Int
wordToInt Word
len'
highIndex :: Int
highIndex :: Int
highIndex = Int
lowIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
highIndex' :: Word
highIndex' :: Word
highIndex' = Int -> Word
intToWord Int
highIndex
ps :: [Int]
ps :: [Int]
ps = if Int
highIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
4 then [] else (Prime Int -> Int) -> [Prime Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Prime Int -> Int
forall a. Prime a -> a
unPrime [Int -> Prime Int
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime Int
2 .. Int -> Prime Int
forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
precPrime (Int -> Int
forall a. Integral a => a -> a
integerSquareRoot Int
highIndex)]
MVector s Word
as <- Int -> Word -> ST s (MVector (PrimState (ST s)) Word)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
len Word
1
Mutable v s a
bs <- Int -> a -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MG.replicate Int
len a
empty
let doPrime :: Int -> ST s ()
doPrime Int
2 = do
let fs :: Vector a
fs = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate (Word -> Int
wordLog2 Word
highIndex')
(\Int
k -> Prime Word -> Word -> a
f (Word -> Prime Word
forall a. a -> Prime a
Prime Word
2) (Int -> Word
intToWord Int
k Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1))
npLow :: Word
npLow = (Word
lowIndex' Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
npHigh :: Word
npHigh = Word
highIndex' Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
[Word] -> (Word -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word
npLow .. Word
npHigh] ((Word -> ST s ()) -> ST s ()) -> (Word -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \np :: Word
np@(W# Word#
np#) -> do
let ix :: Int
ix = Word -> Int
wordToInt (Word
np Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lowIndex :: Int
tz :: Int
tz = Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
ctz# Word#
np#))
MVector (PrimState (ST s)) Word -> (Word -> Word) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector s Word
MVector (PrimState (ST s)) Word
as (\Word
x -> Word
x Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` (Int
tz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Int
ix
Mutable v (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v s a
Mutable v (PrimState (ST s)) a
bs (\a
y -> a
y a -> a -> a
`append` Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
fs Int
tz) Int
ix
doPrime Int
p = do
let p' :: Word
p' = Int -> Word
intToWord Int
p
f0 :: a
f0 = Prime Word -> Word -> a
f (Word -> Prime Word
forall a. a -> Prime a
Prime Word
p') Word
1
logp :: Int
logp = Integer -> Integer -> Int
integerLogBase' (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
p) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
highIndex) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
fs :: Vector a
fs = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
logp (\Int
k -> Prime Word -> Word -> a
f (Word -> Prime Word
forall a. a -> Prime a
Prime Word
p') (Int -> Word
intToWord Int
k Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
2))
npLow :: Int
npLow = (Int
lowIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
p
npHigh :: Int
npHigh = Int
highIndex Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
p
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
npLow .. Int
npHigh] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
np -> do
let !(I# Int#
ix#) = Int
np Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lowIndex
(Int
q, Int
r) = Int
np Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
p
if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
then do
MVector (PrimState (ST s)) Word -> (Word -> Word) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector s Word
MVector (PrimState (ST s)) Word
as (Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
p') (Int# -> Int
I# Int#
ix#)
Mutable v (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v s a
Mutable v (PrimState (ST s)) a
bs (a -> a -> a
`append` a
f0) (Int# -> Int
I# Int#
ix#)
else do
let pow :: Word
pow = Int -> Int -> Word
highestPowerDividing Int
p Int
q
MVector (PrimState (ST s)) Word -> (Word -> Word) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector s Word
MVector (PrimState (ST s)) Word
as (\Word
x -> Word
x Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
p' Word -> Word -> Word
forall a b. (Num a, Integral b) => a -> b -> a
^ (Word
pow Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
2)) (Int# -> Int
I# Int#
ix#)
Mutable v (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v s a
Mutable v (PrimState (ST s)) a
bs (\a
y -> a
y a -> a -> a
`append` Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
fs (Word -> Int
wordToInt Word
pow)) (Int# -> Int
I# Int#
ix#)
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
ps Int -> ST s ()
doPrime
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> do
Word
a <- MVector (PrimState (ST s)) Word -> Int -> ST s Word
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector s Word
MVector (PrimState (ST s)) Word
as Int
k
let a' :: Word
a' = Int -> Word
intToWord (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lowIndex)
Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word
a Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
/= Word
a') (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
Mutable v (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v s a
Mutable v (PrimState (ST s)) a
bs (\a
b -> a
b a -> a -> a
`append` Prime Word -> Word -> a
f (Word -> Prime Word
forall a. a -> Prime a
Prime (Word -> Prime Word) -> Word -> Prime Word
forall a b. (a -> b) -> a -> b
$ Word
a' Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
a) Word
1) Int
k
Mutable v (PrimState (ST s)) a -> ST s (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v s a
Mutable v (PrimState (ST s)) a
bs
highestPowerDividing :: Int -> Int -> Word
highestPowerDividing :: Int -> Int -> Word
highestPowerDividing !Int
_ Int
0 = Word
0
highestPowerDividing Int
p Int
n = Word -> Int -> Word
forall p. Num p => p -> Int -> p
go Word
0 Int
n
where
go :: p -> Int -> p
go !p
k Int
m = case Int
m Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
p of
(Int
q, Int
0) -> p -> Int -> p
go (p
k p -> p -> p
forall a. Num a => a -> a -> a
+ p
1) Int
q
(Int, Int)
_ -> p
k
sieveBlockUnboxed
:: U.Unbox a
=> SieveBlockConfig a
-> Word
-> Word
-> U.Vector a
sieveBlockUnboxed :: SieveBlockConfig a -> Word -> Word -> Vector a
sieveBlockUnboxed = SieveBlockConfig a -> Word -> Word -> Vector a
forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Int -> Word -> Word -> U.Vector Int #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Word -> Word -> Word -> U.Vector Word #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Bool -> Word -> Word -> U.Vector Bool #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Moebius -> Word -> Word -> U.Vector Moebius #-}