-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.Moebius
-- Copyright:   (c) 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Values of <https://en.wikipedia.org/wiki/Möbius_function Möbius function>.
--

{-# 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

-- | Represents three possible values of <https://en.wikipedia.org/wiki/Möbius_function Möbius function>.
data Moebius
  = MoebiusN -- ^ -1
  | MoebiusZ -- ^  0
  | MoebiusP -- ^  1
  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)

-- | Convert to any numeric type.
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
(<>)

-- | Evaluate the Möbius function over a block.
-- Value of @f@ at 0, if zero falls into block, is undefined.
--
-- Based on the sieving algorithm from p. 3 of <https://arxiv.org/pdf/1610.08551.pdf Computations of the Mertens function and improved bounds on the Mertens conjecture> by G. Hurst. It is approximately 5x faster than 'Math.NumberTheory.ArithmeticFunctions.SieveBlock.sieveBlockUnboxed'.
--
-- >>> sieveBlockMoebius 1 10
-- [MoebiusP,MoebiusN,MoebiusN,MoebiusZ,MoebiusN,MoebiusP,MoebiusN,MoebiusZ,MoebiusZ,MoebiusP]
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

    -- Bit fiddling in 'mapper' is correct only
    -- if all sufficiently small (<= 191) primes has been sieved out.
    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