{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
module Math.NumberTheory.ArithmeticFunctions.SieveBlock.Unboxed
( SieveBlockConfig(..)
, multiplicativeSieveBlockConfig
, additiveSieveBlockConfig
, sieveBlockUnboxed
) where
import Control.Monad (forM_)
import Control.Monad.ST (runST)
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as MV
import GHC.Exts
import Math.NumberTheory.ArithmeticFunctions.Moebius (Moebius)
import Math.NumberTheory.Logarithms (integerLogBase')
import Math.NumberTheory.Primes.Sieve (primes)
import Math.NumberTheory.Primes.Types (Prime(..))
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 = (+)
}
sieveBlockUnboxed
:: V.Unbox a
=> SieveBlockConfig a
-> Word
-> Word
-> V.Vector a
sieveBlockUnboxed _ _ 0 = V.empty
sieveBlockUnboxed (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
let lowIndex :: Int
lowIndex = wordToInt lowIndex'
len :: Int
len = wordToInt len'
as <- V.unsafeThaw $ V.enumFromN lowIndex' len
bs <- MV.replicate len empty
let highIndex :: Int
highIndex = lowIndex + len - 1
ps :: [Int]
ps = takeWhile (<= integerSquareRoot highIndex) $ map unPrime primes
forM_ ps $ \p -> do
let p# :: Word#
!p'@(W# p#) = intToWord p
fs = V.generate
(integerLogBase' (toInteger p) (toInteger highIndex))
(\k -> f (Prime p') (intToWord k + 1))
offset :: Int
offset = negate lowIndex `mod` p
forM_ [offset, offset + p .. len - 1] $ \ix -> do
W# a# <- MV.unsafeRead as ix
let !(# pow#, a'# #) = splitOff# p# (a# `quotWord#` p#)
MV.unsafeWrite as ix (W# a'#)
MV.unsafeModify bs (\y -> y `append` V.unsafeIndex fs (I# (word2Int# pow#))) ix
forM_ [0 .. len - 1] $ \k -> do
a <- MV.unsafeRead as k
MV.unsafeModify bs (\b -> if a /= 1 then b `append` f (Prime a) 1 else b) k
V.unsafeFreeze bs
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Int -> Word -> Word -> V.Vector Int #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Word -> Word -> Word -> V.Vector Word #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Bool -> Word -> Word -> V.Vector Bool #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Moebius -> Word -> Word -> V.Vector Moebius #-}