-- |-- Module: Math.NumberTheory.ArithmeticFunctions.SieveBlock-- Copyright: (c) 2017 Andrew Lelechenko-- Licence: MIT-- Maintainer: Andrew Lelechenko <andrew.lelechenko@gmail.com>---- Bulk evaluation of arithmetic functions over continuous intervals-- without factorisation.--{-# LANGUAGE BangPatterns #-}{-# LANGUAGE MagicHash #-}{-# LANGUAGE ScopedTypeVariables #-}moduleMath.NumberTheory.ArithmeticFunctions.SieveBlock(runFunctionOverBlock ,SieveBlockConfig (..),multiplicativeSieveBlockConfig ,additiveSieveBlockConfig ,sieveBlock ,sieveBlockUnboxed ,sieveBlockMoebius )whereimportControl.Monad(forM_,when)importControl.Monad.ST(runST)importData.BitsimportData.CoerceimportqualifiedData.Vector.GenericasGimportqualifiedData.Vector.Generic.MutableasMGimportqualifiedData.VectorasVimportqualifiedData.Vector.UnboxedasUimportqualifiedData.Vector.Unboxed.MutableasMUimportGHC.ExtsimportMath.NumberTheory.ArithmeticFunctions.Class importMath.NumberTheory.ArithmeticFunctions.Moebius (Moebius ,sieveBlockMoebius )importMath.NumberTheory.Logarithms(wordLog2,integerLogBase')importMath.NumberTheory.Primes importMath.NumberTheory.Primes.Types importMath.NumberTheory.Roots(integerSquareRoot)importMath.NumberTheory.Utils.FromIntegral (wordToInt ,intToWord )-- | A record, which specifies a function to evaluate over a block.---- For example, here is a configuration for the totient function:---- > SieveBlockConfig-- > { sbcEmpty = 1-- > , sbcFunctionOnPrimePower = \p a -> (unPrime p - 1) * unPrime p ^ (a - 1)-- > , sbcAppend = (*)-- > }dataSieveBlockConfig a =SieveBlockConfig {forall a. SieveBlockConfig a -> a
sbcEmpty ::a -- ^ value of a function on 1,forall a. SieveBlockConfig a -> Prime Word -> Word -> a
sbcFunctionOnPrimePower ::Prime Word->Word->a -- ^ how to evaluate a function on prime powers,forall a. SieveBlockConfig a -> a -> a -> a
sbcAppend ::a ->a ->a -- ^ how to combine values of a function on coprime arguments}-- | Create a config for a multiplicative function from its definition on prime powers.multiplicativeSieveBlockConfig ::Numa =>(Prime Word->Word->a )->SieveBlockConfig a multiplicativeSieveBlockConfig :: forall a. Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig Prime Word -> Word -> a
f =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
(*)}-- | Create a config for an additive function from its definition on prime powers.additiveSieveBlockConfig ::Numa =>(Prime Word->Word->a )->SieveBlockConfig a additiveSieveBlockConfig :: forall a. Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig Prime Word -> Word -> a
f =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' @f@ @x@ @l@ evaluates an arithmetic function-- for integers between @x@ and @x+l-1@ and returns a vector of length @l@.-- It completely avoids factorisation, so it is asymptotically faster than-- pointwise evaluation of @f@.---- Value of @f@ at 0, if zero falls into block, is undefined.---- Beware that for underlying non-commutative monoids the results may potentially-- differ from pointwise application via 'runFunction'.---- This is a thin wrapper over 'sieveBlock', read more details there.---- >>> import Math.NumberTheory.ArithmeticFunctions-- >>> runFunctionOverBlock carmichaelA 1 10-- [1,1,2,2,4,2,6,2,6,4]runFunctionOverBlock ::ArithmeticFunction Worda ->Word->Word->V.Vectora runFunctionOverBlock :: forall a. 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.mapm -> 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 {sbcEmpty :: m
sbcEmpty =m
forall a. Monoid a => a
mempty,sbcAppend :: m -> m -> m
sbcAppend =m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>),sbcFunctionOnPrimePower :: Prime Word -> Word -> m
sbcFunctionOnPrimePower =(Prime Word -> Word -> m) -> Prime Word -> Word -> m
forall a b. Coercible a b => a -> b
coercePrime Word -> Word -> m
f }-- | Evaluate a function over a block in accordance to provided configuration.-- Value of @f@ at 0, if zero falls into block, is undefined.---- Based on Algorithm M of <https://arxiv.org/pdf/1305.1639.pdf Parity of the number of primes in a given interval and algorithms of the sublinear summation> by A. V. Lelechenko. See Lemma 2 on p. 5 on its algorithmic complexity. For the majority of use-cases its time complexity is O(x^(1+ε)).---- For example, following code lists smallest prime factors:---- >>> sieveBlock (SieveBlockConfig maxBound (\p _ -> unPrime p) min) 2 10 :: Data.Vector.Vector Word-- [2,3,2,5,2,7,2,3,2,11]---- And this is how to factorise all numbers in a block:---- >>> sieveBlock (SieveBlockConfig [] (\p k -> [(unPrime p, k)]) (++)) 2 10 :: Data.Vector.Vector [(Word, Word)]-- [[(2,1)],[(3,1)],[(2,2)],[(5,1)],[(2,1),(3,1)],[(7,1)],[(2,3)],[(3,2)],[(2,1),(5,1)],[(11,1)]]sieveBlock ::forallv a .G.Vectorv a =>SieveBlockConfig a ->Word->Word->v a sieveBlock :: forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock SieveBlockConfig a
_Word
_Word
0=v a
forall (v :: * -> *) a. Vector v a => v a
G.emptysieveBlock (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
$doletlowIndex ::IntlowIndex :: Int
lowIndex =Word -> Int
wordToInt Word
lowIndex' len ::Intlen :: Int
len =Word -> Int
wordToInt Word
len' highIndex ::InthighIndex :: 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
1highIndex' ::WordhighIndex' :: Word
highIndex' =Int -> Word
intToWord Int
highIndex ps ::[Int]ps :: [Int]
ps =ifInt
highIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
4then[]else(Prime Int -> Int) -> [Prime Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
mapPrime 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
integerSquareRootInt
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.replicateInt
len Word
1Mutable 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.replicateInt
len a
empty letdoPrime :: Int -> m ()
doPrime Int
2=doletfs :: Vector a
fs =Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate(Word -> Int
wordLog2Word
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
1npHigh :: Word
npHigh =Word
highIndex' Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftR`Int
1[Word] -> (Word -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_[Word
npLow ..Word
npHigh ]((Word -> m ()) -> m ()) -> (Word -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$\np :: Word
np @(W#Word#
np# )->doletix :: 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 ::Inttz :: Int
tz =Int# -> Int
I#(Word# -> Int#
word2Int#(Word# -> Word#
ctz#Word#
np# ))MVector (PrimState m) Word -> (Word -> Word) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModifyMVector s Word
MVector (PrimState m) 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 m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModifyMutable v s a
Mutable v (PrimState m) a
bs (\a
y ->a
y a -> a -> a
`append` Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndexVector a
fs Int
tz )Int
ix doPrime Int
p =doletp' :: 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
1logp :: Int
logp =Integer -> Integer -> Int
integerLogBase'(Int -> Integer
forall a. Integral a => a -> Integer
toIntegerInt
p )(Int -> Integer
forall a. Integral a => a -> Integer
toIntegerInt
highIndex )Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1fs :: Vector a
fs =Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generateInt
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 -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_[Int
npLow ..Int
npHigh ]((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$\Int
np ->dolet!(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 ifInt
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/=Int
0thendoMVector (PrimState m) Word -> (Word -> Word) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModifyMVector s Word
MVector (PrimState m) Word
as (Word -> Word -> Word
forall a. Num a => a -> a -> a
*Word
p' )(Int# -> Int
I#Int#
ix# )Mutable v (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModifyMutable v s a
Mutable v (PrimState m) a
bs (a -> a -> a
`append` a
f0 )(Int# -> Int
I#Int#
ix# )elsedoletpow :: Word
pow =Int -> Int -> Word
highestPowerDividing Int
p Int
q MVector (PrimState m) Word -> (Word -> Word) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModifyMVector s Word
MVector (PrimState m) 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 m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModifyMutable v s a
Mutable v (PrimState m) a
bs (\a
y ->a
y a -> a -> a
`append` Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndexVector 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 ()
forall {m :: * -> *}. (PrimState m ~ s, PrimMonad m) => Int -> m ()
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 ->doWord
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.unsafeReadMVector s Word
MVector (PrimState (ST s)) Word
as Int
k leta' :: 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.unsafeModifyMutable 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.unsafeFreezeMutable v s a
Mutable v (PrimState (ST s)) a
bs -- This is a variant of 'Math.NumberTheory.Utils.splitOff',-- specialized for better performance.highestPowerDividing ::Int->Int->WordhighestPowerDividing :: Int -> Int -> Word
highestPowerDividing !Int
_Int
0=Word
0highestPowerDividing Int
p Int
n =Word -> Int -> Word
forall {t}. Num t => t -> Int -> t
go Word
0Int
n wherego :: t -> Int -> t
go !t
k Int
m =caseInt
m Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem`Int
p of(Int
q ,Int
0)->t -> Int -> t
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
+t
1)Int
q (Int, Int)
_->t
k -- | This is 'sieveBlock' specialized to unboxed vectors.---- >>> sieveBlockUnboxed (SieveBlockConfig 1 (\_ a -> a + 1) (*)) 1 10-- [1,2,2,3,2,4,2,4,3,4]sieveBlockUnboxed ::U.Unboxa =>SieveBlockConfig a ->Word->Word->U.Vectora sieveBlockUnboxed :: forall a. Unbox a => 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 {-# SPECIALIZEsieveBlockUnboxed ::SieveBlockConfig Int->Word->Word->U.VectorInt#-}{-# SPECIALIZEsieveBlockUnboxed ::SieveBlockConfig Word->Word->Word->U.VectorWord#-}{-# SPECIALIZEsieveBlockUnboxed ::SieveBlockConfig Bool->Word->Word->U.VectorBool#-}{-# SPECIALIZEsieveBlockUnboxed ::SieveBlockConfig Moebius ->Word->Word->U.VectorMoebius #-}