{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module Math.NumberTheory.ArithmeticFunctions.Moebius
( Moebius(..)
, runMoebius
, sieveBlockMoebius
) where
import Control.Monad (forM_)
import Control.Monad.ST (runST)
import Data.Bits
import Data.Int
import Data.Word
#if __GLASGOW_HASKELL__ < 803
import Data.Semigroup
#endif
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Primitive as P
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import GHC.Exts
import GHC.Integer.GMP.Internals
import Unsafe.Coerce
import Math.NumberTheory.Roots (integerSquareRoot)
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils.FromIntegral
import Math.NumberTheory.Logarithms
data Moebius
= MoebiusN
| MoebiusZ
| MoebiusP
deriving (Moebius -> Moebius -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Moebius -> Moebius -> Bool
$c/= :: Moebius -> Moebius -> Bool
== :: Moebius -> Moebius -> Bool
$c== :: Moebius -> Moebius -> Bool
Eq, Eq Moebius
Moebius -> Moebius -> Bool
Moebius -> Moebius -> Ordering
Moebius -> Moebius -> Moebius
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Moebius -> Moebius -> Moebius
$cmin :: Moebius -> Moebius -> Moebius
max :: Moebius -> Moebius -> Moebius
$cmax :: Moebius -> Moebius -> Moebius
>= :: Moebius -> Moebius -> Bool
$c>= :: Moebius -> Moebius -> Bool
> :: Moebius -> Moebius -> Bool
$c> :: Moebius -> Moebius -> Bool
<= :: Moebius -> Moebius -> Bool
$c<= :: Moebius -> Moebius -> Bool
< :: Moebius -> Moebius -> Bool
$c< :: Moebius -> Moebius -> Bool
compare :: Moebius -> Moebius -> Ordering
$ccompare :: Moebius -> Moebius -> Ordering
Ord, Int -> Moebius -> ShowS
[Moebius] -> ShowS
Moebius -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Moebius] -> ShowS
$cshowList :: [Moebius] -> ShowS
show :: Moebius -> String
$cshow :: Moebius -> String
showsPrec :: Int -> Moebius -> ShowS
$cshowsPrec :: Int -> Moebius -> ShowS
Show)
runMoebius :: Num a => Moebius -> a
runMoebius :: forall a. Num a => Moebius -> a
runMoebius Moebius
m = forall a. Num a => Integer -> a
fromInteger (Int# -> Integer
S# (forall a. a -> Int#
dataToTag# Moebius
m Int# -> Int# -> Int#
-# Int#
1#))
fromMoebius :: Moebius -> Int8
fromMoebius :: Moebius -> Int8
fromMoebius Moebius
m = Int -> Int8
intToInt8 forall a b. (a -> b) -> a -> b
$ Int# -> Int
I# (forall a. a -> Int#
dataToTag# Moebius
m)
{-# INLINE fromMoebius #-}
toMoebius :: Int8 -> Moebius
toMoebius :: Int8 -> Moebius
toMoebius Int8
i = let !(I# Int#
i#) = Int8 -> Int
int8ToInt Int8
i in forall a. Int# -> a
tagToEnum# Int#
i#
{-# INLINE toMoebius #-}
newtype instance U.MVector s Moebius = MV_Moebius (P.MVector s Int8)
newtype instance U.Vector Moebius = V_Moebius (P.Vector Int8)
instance U.Unbox Moebius
instance M.MVector U.MVector Moebius where
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicOverlaps #-}
{-# INLINE basicUnsafeNew #-}
{-# INLINE basicInitialize #-}
{-# INLINE basicUnsafeReplicate #-}
{-# INLINE basicUnsafeRead #-}
{-# INLINE basicUnsafeWrite #-}
{-# INLINE basicClear #-}
{-# INLINE basicSet #-}
{-# INLINE basicUnsafeCopy #-}
{-# INLINE basicUnsafeGrow #-}
basicLength :: forall s. MVector s Moebius -> Int
basicLength (MV_Moebius MVector s Int8
v) = forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.basicLength MVector s Int8
v
basicUnsafeSlice :: forall s. Int -> Int -> MVector s Moebius -> MVector s Moebius
basicUnsafeSlice Int
i Int
n (MV_Moebius MVector s Int8
v) = forall s. MVector s Int8 -> MVector s Moebius
MV_Moebius forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
M.basicUnsafeSlice Int
i Int
n MVector s Int8
v
basicOverlaps :: forall s. MVector s Moebius -> MVector s Moebius -> Bool
basicOverlaps (MV_Moebius MVector s Int8
v1) (MV_Moebius MVector s Int8
v2) = forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> Bool
M.basicOverlaps MVector s Int8
v1 MVector s Int8
v2
basicUnsafeNew :: forall s. Int -> ST s (MVector s Moebius)
basicUnsafeNew Int
n = forall s. MVector s Int8 -> MVector s Moebius
MV_Moebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> * -> *) a s. MVector v a => Int -> ST s (v s a)
M.basicUnsafeNew Int
n
basicInitialize :: forall s. MVector s Moebius -> ST s ()
basicInitialize (MV_Moebius MVector s Int8
v) = forall (v :: * -> * -> *) a s. MVector v a => v s a -> ST s ()
M.basicInitialize MVector s Int8
v
basicUnsafeReplicate :: forall s. Int -> Moebius -> ST s (MVector s Moebius)
basicUnsafeReplicate Int
n Moebius
x = forall s. MVector s Int8 -> MVector s Moebius
MV_Moebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> a -> ST s (v s a)
M.basicUnsafeReplicate Int
n (Moebius -> Int8
fromMoebius Moebius
x)
basicUnsafeRead :: forall s. MVector s Moebius -> Int -> ST s Moebius
basicUnsafeRead (MV_Moebius MVector s Int8
v) Int
i = Int8 -> Moebius
toMoebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> ST s a
M.basicUnsafeRead MVector s Int8
v Int
i
basicUnsafeWrite :: forall s. MVector s Moebius -> Int -> Moebius -> ST s ()
basicUnsafeWrite (MV_Moebius MVector s Int8
v) Int
i Moebius
x = forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> a -> ST s ()
M.basicUnsafeWrite MVector s Int8
v Int
i (Moebius -> Int8
fromMoebius Moebius
x)
basicClear :: forall s. MVector s Moebius -> ST s ()
basicClear (MV_Moebius MVector s Int8
v) = forall (v :: * -> * -> *) a s. MVector v a => v s a -> ST s ()
M.basicClear MVector s Int8
v
basicSet :: forall s. MVector s Moebius -> Moebius -> ST s ()
basicSet (MV_Moebius MVector s Int8
v) Moebius
x = forall (v :: * -> * -> *) a s. MVector v a => v s a -> a -> ST s ()
M.basicSet MVector s Int8
v (Moebius -> Int8
fromMoebius Moebius
x)
basicUnsafeCopy :: forall s. MVector s Moebius -> MVector s Moebius -> ST s ()
basicUnsafeCopy (MV_Moebius MVector s Int8
v1) (MV_Moebius MVector s Int8
v2) = forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> ST s ()
M.basicUnsafeCopy MVector s Int8
v1 MVector s Int8
v2
basicUnsafeMove :: forall s. MVector s Moebius -> MVector s Moebius -> ST s ()
basicUnsafeMove (MV_Moebius MVector s Int8
v1) (MV_Moebius MVector s Int8
v2) = forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> ST s ()
M.basicUnsafeMove MVector s Int8
v1 MVector s Int8
v2
basicUnsafeGrow :: forall s. MVector s Moebius -> Int -> ST s (MVector s Moebius)
basicUnsafeGrow (MV_Moebius MVector s Int8
v) Int
n = forall s. MVector s Int8 -> MVector s Moebius
MV_Moebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> ST s (v s a)
M.basicUnsafeGrow MVector s Int8
v Int
n
instance G.Vector U.Vector Moebius where
{-# INLINE basicUnsafeFreeze #-}
{-# INLINE basicUnsafeThaw #-}
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicUnsafeIndexM #-}
{-# INLINE elemseq #-}
basicUnsafeFreeze :: forall s. Mutable Vector s Moebius -> ST s (Vector Moebius)
basicUnsafeFreeze (MV_Moebius MVector s Int8
v) = Vector Int8 -> Vector Moebius
V_Moebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> *) a s. Vector v a => Mutable v s a -> ST s (v a)
G.basicUnsafeFreeze MVector s Int8
v
basicUnsafeThaw :: forall s. Vector Moebius -> ST s (Mutable Vector s Moebius)
basicUnsafeThaw (V_Moebius Vector Int8
v) = forall s. MVector s Int8 -> MVector s Moebius
MV_Moebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> *) a s. Vector v a => v a -> ST s (Mutable v s a)
G.basicUnsafeThaw Vector Int8
v
basicLength :: Vector Moebius -> Int
basicLength (V_Moebius Vector Int8
v) = forall (v :: * -> *) a. Vector v a => v a -> Int
G.basicLength Vector Int8
v
basicUnsafeSlice :: Int -> Int -> Vector Moebius -> Vector Moebius
basicUnsafeSlice Int
i Int
n (V_Moebius Vector Int8
v) = Vector Int8 -> Vector Moebius
V_Moebius forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.basicUnsafeSlice Int
i Int
n Vector Int8
v
basicUnsafeIndexM :: Vector Moebius -> Int -> Box Moebius
basicUnsafeIndexM (V_Moebius Vector Int8
v) Int
i = Int8 -> Moebius
toMoebius forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (v :: * -> *) a. Vector v a => v a -> Int -> Box a
G.basicUnsafeIndexM Vector Int8
v Int
i
basicUnsafeCopy :: forall s. Mutable Vector s Moebius -> Vector Moebius -> ST s ()
basicUnsafeCopy (MV_Moebius MVector s Int8
mv) (V_Moebius Vector Int8
v) = forall (v :: * -> *) a s.
Vector v a =>
Mutable v s a -> v a -> ST s ()
G.basicUnsafeCopy MVector s Int8
mv Vector Int8
v
elemseq :: forall b. Vector Moebius -> Moebius -> b -> b
elemseq Vector Moebius
_ = seq :: forall a b. a -> b -> b
seq
instance Semigroup Moebius where
Moebius
MoebiusZ <> :: Moebius -> Moebius -> Moebius
<> Moebius
_ = Moebius
MoebiusZ
Moebius
_ <> Moebius
MoebiusZ = Moebius
MoebiusZ
Moebius
MoebiusP <> Moebius
a = Moebius
a
Moebius
a <> Moebius
MoebiusP = Moebius
a
Moebius
_ <> Moebius
_ = Moebius
MoebiusP
instance Monoid Moebius where
mempty :: Moebius
mempty = Moebius
MoebiusP
mappend :: Moebius -> Moebius -> Moebius
mappend = forall a. Semigroup a => a -> a -> a
(<>)
sieveBlockMoebius
:: Word
-> Word
-> U.Vector Moebius
sieveBlockMoebius :: Word -> Word -> Vector Moebius
sieveBlockMoebius Word
_ Word
0 = forall a. Unbox a => Vector a
U.empty
sieveBlockMoebius Word
lowIndex' Word
len'
= (forall a b. a -> b
unsafeCoerce :: U.Vector Word8 -> U.Vector Moebius) forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
MVector s Word8
as <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
len (Word8
0x80 :: Word8)
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
ps forall a b. (a -> b) -> a -> b
$ \Int
p -> do
let offset :: Int
offset = forall a. Num a => a -> a
negate Int
lowIndex forall a. Integral a => a -> a -> a
`mod` Int
p
offset2 :: Int
offset2 = forall a. Num a => a -> a
negate Int
lowIndex forall a. Integral a => a -> a -> a
`mod` (Int
p forall a. Num a => a -> a -> a
* Int
p)
l :: Word8
l :: Word8
l = Int -> Word8
intToWord8 forall a b. (a -> b) -> a -> b
$ Int -> Int
intLog2 Int
p forall a. Bits a => a -> a -> a
.|. Int
1
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
offset, Int
offset forall a. Num a => a -> a -> a
+ Int
p .. Int
len forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector s Word8
as (forall a. Num a => a -> a -> a
+ Word8
l)
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
offset2, Int
offset2 forall a. Num a => a -> a -> a
+ Int
p forall a. Num a => a -> a -> a
* Int
p .. Int
len forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
ix ->
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Word8
as Int
ix Word8
0
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
len forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
ix ->
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector s Word8
as (Int -> Word8 -> Word8
mapper Int
ix) Int
ix
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Word8
as
where
lowIndex :: Int
lowIndex :: Int
lowIndex = Word -> Int
wordToInt Word
lowIndex'
len :: Int
len :: Int
len = Word -> Int
wordToInt Word
len'
highIndex :: Int
highIndex :: Int
highIndex = Int
lowIndex forall a. Num a => a -> a -> a
+ Int
len forall a. Num a => a -> a -> a
- Int
1
ps :: [Int]
ps :: [Int]
ps = 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 Int
2 .. forall a.
(Bits a, Integral a, UniqueFactorisation a) =>
a -> Prime a
precPrime (Int
191 forall a. Ord a => a -> a -> a
`max` forall a. Integral a => a -> a
integerSquareRoot Int
highIndex)]
mapper :: Int -> Word8 -> Word8
mapper :: Int -> Word8 -> Word8
mapper Int
ix Word8
val
| Word8
val forall a. Bits a => a -> a -> a
.&. Word8
0x80 forall a. Eq a => a -> a -> Bool
== Word8
0x00
= Word8
1
| Word8 -> Int
word8ToInt (Word8
val forall a. Bits a => a -> a -> a
.&. Word8
0x7F) forall a. Ord a => a -> a -> Bool
< Int -> Int
intLog2 (Int
ix forall a. Num a => a -> a -> a
+ Int
lowIndex) forall a. Num a => a -> a -> a
- Int
5
forall a. Num a => a -> a -> a
- (if Int
ix forall a. Num a => a -> a -> a
+ Int
lowIndex forall a. Ord a => a -> a -> Bool
>= Int
0x100000 then Int
2 else Int
0)
forall a. Num a => a -> a -> a
- (if Int
ix forall a. Num a => a -> a -> a
+ Int
lowIndex forall a. Ord a => a -> a -> Bool
>= Int
0x10000000 then Int
1 else Int
0)
= (Word8
val forall a. Bits a => a -> a -> a
.&. Word8
1) forall a. Bits a => a -> Int -> a
`shiftL` Int
1
| Bool
otherwise
= ((Word8
val forall a. Bits a => a -> a -> a
.&. Word8
1) forall a. Bits a => a -> a -> a
`xor` Word8
1) forall a. Bits a => a -> Int -> a
`shiftL` Int
1