module Math.Sieve.Phi
( sieve
, phi
, isPrime
, sieveBound
) where
import Control.Monad
import Control.Monad.ST
import Data.Array.Unboxed
import Data.Array.ST
import Data.Bits
import Data.Word
newtype PhiSieve = PhiSieve (UArray Word32 Word32)
instance Show PhiSieve where
show fs = "<<PhiSieve " ++ show (shiftL (sieveBound' fs) 1 + 2) ++ ">>"
instance Eq PhiSieve where
a == b = sieveBound' a == sieveBound' b
instance Ord PhiSieve where
compare a b = compare (sieveBound' a) (sieveBound' b)
sieveBound :: Integral a => PhiSieve -> a
sieveBound (PhiSieve arr) = fromIntegral (shiftL (snd (bounds arr)) 1 + 2)
sieveBound' :: PhiSieve -> Word32
sieveBound' (PhiSieve arr) = (snd (bounds arr))
ix :: Word32 -> Word32
ix n = shiftR n 1
sieve :: Integral a => a -> PhiSieve
sieve limit
| not (0 <= intlim && intlim <= fromIntegral (maxBound :: Word32))
= error ("Math.Sieve.Phi: out of bounds: " ++ show limit)
| otherwise
= PhiSieve $ runSTUArray $ do
a <- newArray (0, lim) 1 :: ST s (STUArray s Word32 Word32)
pows <- newArray (0, 19) 0 :: ST s (STUArray s Word32 Word32)
sieve a 3 pows
where
intlim = (fromIntegral limit) :: Integer
lim = ix (fromIntegral limit 1) :: Word32
mult a p i = do
x <- readArray a i
writeArray a i (p * x)
sieve a p pows
| ix p > lim = return a
| otherwise = do
phi <- readArray a (ix p)
when (phi == 1) $ do
initArray pows (p1)
adjustPhi a p (ix p) pows
sieve a (p+2) pows
adjustPhi !a !p !i !pows
| i > lim = return ()
| otherwise = do
k <- zeros pows
if k == 0
then mult a (p1) i
else mult a ((p1)*p^k) i
decpows p pows
decpows p pows
adjustPhi a p (i + p) pows
initArray :: STUArray s Word32 Word32 -> Word32 -> ST s ()
initArray !a !val = loop a val 0
where
loop a val n
| n <= 19 = writeArray a n val >> loop a val (n+1)
| otherwise = return ()
decpows :: Word32 -> STUArray s Word32 Word32 -> ST s ()
decpows !p !pows = loop p pows 0
where
loop p pows n = do
x <- readArray pows n
if x == 0
then do
writeArray pows n (p1)
loop p pows (n+1)
else do
writeArray pows n (x1)
zeros :: STUArray s Word32 Word32 -> ST s Word32
zeros !pows = loop pows 0
where
loop pows n = do
x <- readArray pows n
if x == 0
then loop pows (n+1)
else return n
phi :: (Integral a, Integral b) => PhiSieve -> a -> b
phi (PhiSieve a) nn
| even n0 = loop 0 (shiftR n0 1)
| otherwise = fromIntegral (a ! (ix n0))
where
n0 = fromIntegral nn
loop !t n
| even n = loop (t+1) (shiftR n 1)
| otherwise = fromIntegral (shiftL (a ! (ix n)) t)
isPrime :: Integral a => PhiSieve -> a -> Bool
isPrime (PhiSieve a) n
| even n = n == 2
| otherwise = a ! (ix nn) == nn 1
where
nn = fromIntegral n