{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Math.NumberTheory.ArithmeticFunctions.NFreedom
( nFrees
, nFreesBlock
, sieveBlockNFree
) where
import Control.Monad (forM_)
import Control.Monad.ST (runST)
import Data.Bits (Bits)
import Data.List (scanl')
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import Math.NumberTheory.Roots
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils.FromIntegral
sieveBlockNFree
:: forall a. (Integral a, Enum (Prime a), Bits a, UniqueFactorisation a)
=> Word
-> a
-> Word
-> U.Vector Bool
sieveBlockNFree :: forall a.
(Integral a, Enum (Prime a), Bits a, UniqueFactorisation a) =>
Word -> a -> Word -> Vector Bool
sieveBlockNFree Word
_ a
_ Word
0 = forall a. Unbox a => Vector a
U.empty
sieveBlockNFree Word
n a
lowIndex Word
len'
= forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
MVector s Bool
as <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate (Word -> Int
wordToInt Word
len') Bool
True
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
ps forall a b. (a -> b) -> a -> b
$ \a
p -> do
let pPow :: a
pPow :: a
pPow = a
p forall a b. (Num a, Integral b) => a -> b -> a
^ Word
n
offset :: a
offset :: a
offset = forall a. Num a => a -> a
negate a
lowIndex forall a. Integral a => a -> a -> a
`mod` a
pPow
indices :: [a]
indices :: [a]
indices = [a
offset, a
offset forall a. Num a => a -> a -> a
+ a
pPow .. a
len forall a. Num a => a -> a -> a
- a
1]
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
indices forall a b. (a -> b) -> a -> b
$ \a
ix ->
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.write MVector s Bool
as (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
ix) Bool
False
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.freeze MVector s Bool
as
where
len :: a
len :: a
len = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len'
highIndex :: a
highIndex :: a
highIndex = a
lowIndex forall a. Num a => a -> a -> a
+ a
len forall a. Num a => a -> a -> a
- a
1
ps :: [a]
ps :: [a]
ps = if a
highIndex forall a. Ord a => a -> a -> Bool
< a
4 then [] else forall a b. (a -> b) -> [a] -> [b]
map forall a. Prime a -> a
unPrime [forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
nextPrime a
2 .. forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
precPrime (forall a. Integral a => a -> a
integerSquareRoot a
highIndex)]
nFrees
:: forall a. (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
=> Word
-> [a]
nFrees :: forall a.
(Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
Word -> [a]
nFrees Word
0 = [a
1]
nFrees Word
1 = [a
1]
nFrees Word
n = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a.
(Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
Word -> a -> Word -> [a]
nFreesBlock Word
n)) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [a]
bounds [Word]
strides
where
strides :: [Word]
strides :: [Word]
strides = forall a. Int -> [a] -> [a]
take Int
55 (forall a. (a -> a) -> a -> [a]
iterate (Word
2 forall a. Num a => a -> a -> a
*) Word
256) forall a. [a] -> [a] -> [a]
++ forall a. a -> [a]
repeat (Int -> Word
intToWord (forall a. Bounded a => a
maxBound :: Int))
bounds :: [a]
bounds :: [a]
bounds = forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' forall a. Num a => a -> a -> a
(+) a
1 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word]
strides
nFreesBlock
:: forall a . (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
=> Word
-> a
-> Word
-> [a]
nFreesBlock :: forall a.
(Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
Word -> a -> Word -> [a]
nFreesBlock Word
0 a
lo Word
_ = forall a. Integral a => a -> [a]
help a
lo
nFreesBlock Word
1 a
lo Word
_ = forall a. Integral a => a -> [a]
help a
lo
nFreesBlock Word
n a
lowIndex Word
len =
let
len' :: Int
len' :: Int
len' = Word -> Int
wordToInt Word
len
len'' :: a
len'' :: a
len'' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len
bs :: Vector Bool
bs = forall a.
(Integral a, Enum (Prime a), Bits a, UniqueFactorisation a) =>
Word -> a -> Word -> Vector Bool
sieveBlockNFree Word
n a
lowIndex Word
len
in forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
.
forall a. (a -> Bool) -> [a] -> [a]
filter ((Vector Bool
bs forall a. Unbox a => Vector a -> Int -> a
U.!) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 .. Int
len' forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ [a
lowIndex .. a
lowIndex forall a. Num a => a -> a -> a
+ a
len'']
{-# INLINE nFreesBlock #-}
help :: Integral a => a -> [a]
help :: forall a. Integral a => a -> [a]
help a
1 = [a
1]
help a
_ = []
{-# INLINE help #-}