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