{-# 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.Powers.Squares (integerSquareRoot)
import Math.NumberTheory.Utils (splitOff)
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, _) = splitOff 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
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 #-}