1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
|
-- |
-- 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 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
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.Num.Integer
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 (Eq, Ord, Show)
-- | Convert to any numeric type.
runMoebius :: Num a => Moebius -> a
runMoebius m = fromInteger (IS (dataToTag# m -# 1#))
fromMoebius :: Moebius -> Int8
fromMoebius m = intToInt8 $ I# (dataToTag# m)
{-# INLINE fromMoebius #-}
toMoebius :: Int8 -> Moebius
toMoebius i = let !(I# i#) = int8ToInt i in tagToEnum# 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 (MV_Moebius v) = M.basicLength v
basicUnsafeSlice i n (MV_Moebius v) = MV_Moebius $ M.basicUnsafeSlice i n v
basicOverlaps (MV_Moebius v1) (MV_Moebius v2) = M.basicOverlaps v1 v2
basicUnsafeNew n = MV_Moebius <$> M.basicUnsafeNew n
basicInitialize (MV_Moebius v) = M.basicInitialize v
basicUnsafeReplicate n x = MV_Moebius <$> M.basicUnsafeReplicate n (fromMoebius x)
basicUnsafeRead (MV_Moebius v) i = toMoebius <$> M.basicUnsafeRead v i
basicUnsafeWrite (MV_Moebius v) i x = M.basicUnsafeWrite v i (fromMoebius x)
basicClear (MV_Moebius v) = M.basicClear v
basicSet (MV_Moebius v) x = M.basicSet v (fromMoebius x)
basicUnsafeCopy (MV_Moebius v1) (MV_Moebius v2) = M.basicUnsafeCopy v1 v2
basicUnsafeMove (MV_Moebius v1) (MV_Moebius v2) = M.basicUnsafeMove v1 v2
basicUnsafeGrow (MV_Moebius v) n = MV_Moebius <$> M.basicUnsafeGrow v n
instance G.Vector U.Vector Moebius where
{-# INLINE basicUnsafeFreeze #-}
{-# INLINE basicUnsafeThaw #-}
{-# INLINE basicLength #-}
{-# INLINE basicUnsafeSlice #-}
{-# INLINE basicUnsafeIndexM #-}
{-# INLINE elemseq #-}
basicUnsafeFreeze (MV_Moebius v) = V_Moebius <$> G.basicUnsafeFreeze v
basicUnsafeThaw (V_Moebius v) = MV_Moebius <$> G.basicUnsafeThaw v
basicLength (V_Moebius v) = G.basicLength v
basicUnsafeSlice i n (V_Moebius v) = V_Moebius $ G.basicUnsafeSlice i n v
basicUnsafeIndexM (V_Moebius v) i = toMoebius <$> G.basicUnsafeIndexM v i
basicUnsafeCopy (MV_Moebius mv) (V_Moebius v) = G.basicUnsafeCopy mv v
elemseq _ = seq
instance Semigroup Moebius where
MoebiusZ <> _ = MoebiusZ
_ <> MoebiusZ = MoebiusZ
MoebiusP <> a = a
a <> MoebiusP = a
_ <> _ = MoebiusP
instance Monoid Moebius where
mempty = MoebiusP
-- | 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 _ 0 = U.empty
sieveBlockMoebius lowIndex' len'
= (unsafeCoerce :: U.Vector Word8 -> U.Vector Moebius) $ runST $ do
as <- MU.replicate len (0x80 :: Word8)
forM_ ps $ \p -> do
let offset = negate lowIndex `mod` p
offset2 = negate lowIndex `mod` (p * p)
l :: Word8
l = intToWord8 $ intLog2 p .|. 1
forM_ [offset, offset + p .. len - 1] $
MU.unsafeModify as (+ l)
forM_ [offset2, offset2 + p * p .. len - 1] $ \ix ->
MU.unsafeWrite as ix 0
forM_ [0 .. len - 1] $ \ix ->
MU.unsafeModify as (mapper ix) ix
U.unsafeFreeze as
where
lowIndex :: Int
lowIndex = wordToInt lowIndex'
len :: Int
len = wordToInt len'
highIndex :: Int
highIndex = lowIndex + len - 1
-- Bit fiddling in 'mapper' is correct only
-- if all sufficiently small (<= 191) primes has been sieved out.
ps :: [Int]
ps = map unPrime [nextPrime 2 .. precPrime (191 `max` integerSquareRoot highIndex)]
mapper :: Int -> Word8 -> Word8
mapper ix val
| val .&. 0x80 == 0x00
= 1
| word8ToInt (val .&. 0x7F) < intLog2 (ix + lowIndex) - 5
- (if ix + lowIndex >= 0x100000 then 2 else 0)
- (if ix + lowIndex >= 0x10000000 then 1 else 0)
= (val .&. 1) `shiftL` 1
| otherwise
= ((val .&. 1) `xor` 1) `shiftL` 1
|
