{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
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
{ sbcEmpty :: a
, sbcFunctionOnPrimePower :: Prime Word -> Word -> a
, sbcAppend :: a -> a -> a
}
multiplicativeSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig f = SieveBlockConfig
{ sbcEmpty = 1
, sbcFunctionOnPrimePower = f
, sbcAppend = (*)
}
additiveSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig f = SieveBlockConfig
{ sbcEmpty = 0
, sbcFunctionOnPrimePower = f
, sbcAppend = (+)
}
runFunctionOverBlock
:: ArithmeticFunction Word a
-> Word
-> Word
-> V.Vector a
runFunctionOverBlock (ArithmeticFunction f g) = (G.map g .) . sieveBlock SieveBlockConfig
{ sbcEmpty = mempty
, sbcAppend = mappend
, sbcFunctionOnPrimePower = coerce f
}
sieveBlock
:: forall v a.
G.Vector v a
=> SieveBlockConfig a
-> Word
-> Word
-> v a
sieveBlock _ _ 0 = G.empty
sieveBlock (SieveBlockConfig empty f append) !lowIndex' len' = runST $ do
let lowIndex :: Int
lowIndex = wordToInt lowIndex'
len :: Int
len = wordToInt len'
highIndex :: Int
highIndex = lowIndex + len - 1
highIndex' :: Word
highIndex' = intToWord highIndex
ps :: [Int]
ps = if highIndex < 4 then [] else map unPrime [nextPrime 2 .. precPrime (integerSquareRoot highIndex)]
as <- MU.replicate len 1
bs <- MG.replicate len empty
let doPrime 2 = do
let fs = V.generate (wordLog2 highIndex')
(\k -> f (Prime 2) (intToWord k + 1))
npLow = (lowIndex' + 1) `shiftR` 1
npHigh = highIndex' `shiftR` 1
forM_ [npLow .. npHigh] $ \np@(W# np#) -> do
let ix = wordToInt (np `shiftL` 1) - lowIndex :: Int
tz = I# (word2Int# (ctz# np#))
MU.unsafeModify as (\x -> x `shiftL` (tz + 1)) ix
MG.unsafeModify bs (\y -> y `append` V.unsafeIndex fs tz) ix
doPrime p = do
let p' = intToWord p
f0 = f (Prime p') 1
logp = integerLogBase' (toInteger p) (toInteger highIndex) - 1
fs = V.generate logp (\k -> f (Prime p') (intToWord k + 2))
npLow = (lowIndex + p - 1) `quot` p
npHigh = highIndex `quot` p
forM_ [npLow .. npHigh] $ \np -> do
let !(I# ix#) = np * p - lowIndex
(q, r) = np `quotRem` p
if r /= 0
then do
MU.unsafeModify as (\x -> x * p') (I# ix#)
MG.unsafeModify bs (\y -> y `append` f0) (I# ix#)
else do
let pow = highestPowerDividing p q
MU.unsafeModify as (\x -> x * p' ^ (pow + 2)) (I# ix#)
MG.unsafeModify bs (\y -> y `append` V.unsafeIndex fs (wordToInt pow)) (I# ix#)
forM_ ps doPrime
forM_ [0 .. len - 1] $ \k -> do
a <- MU.unsafeRead as k
let a' = intToWord (k + lowIndex)
when (a /= a') $
MG.unsafeModify bs (\b -> b `append` f (Prime $ a' `quot` a) 1) k
G.unsafeFreeze bs
highestPowerDividing :: Int -> Int -> Word
highestPowerDividing !_ 0 = 0
highestPowerDividing p n = go 0 n
where
go !k m = case m `quotRem` p of
(q, 0) -> go (k + 1) q
_ -> k
sieveBlockUnboxed
:: U.Unbox a
=> SieveBlockConfig a
-> Word
-> Word
-> U.Vector a
sieveBlockUnboxed = 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 #-}