{-# LANGUAGE BangPatterns #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  NumberTheory.Sieve.Phi
-- Copyright   :  (c) Leon P Smith 2009
-- License     :  BSD3
--
-- Maintainer  :  leon@melding-monads.com
-- Stability   :  experimental
-- Portability :  portable
--
-- This is a sieve that calculates Euler's Totient,  also know as Euler's
-- Phi function,  for every number up to a given limit.  @phi(n)@ is defined
-- as the number of positive integers less than @n@ that are relatively
-- prime to @n@,  i.e. @gcd n x == 1@.   Given the prime factorization of
-- a number,  it can be calculated efficiently from the formulas:
--
-- > phi (p ^ k) = (p-1)*p^(k-1)    if p is prime
-- > phi (x * y) = phi x * phi y    if x and y are relatively prime
--
-- The second case says that @phi@ is a /multiplicative/ function.  Since
-- any multiplicative function can be calculated from the prime factorization,
-- 'NumberTheory.Sieve.Factor' can also be applied,  however this variant
-- avoids a great deal of integer division,  and so is significantly faster
-- for calculating @phi@ for large quantities of different values.
--
-- This sieve does not represent even numbers directly,  and the maximum
-- number that can currently be sieved is 2^32.   This means that the sieve
-- requires two bytes per number sieved on average,  and thus sieving up to
-- 2^32 requires 8 GB of storage.
--
-----------------------------------------------------------------------------


module NumberTheory.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)

-- |  The upper limit of the sieve
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

-- |  Calculate Euler's Totient for every integer up to the given limit
sieve :: Integral a => a -> PhiSieve
sieve limit
  | not (0 <= intlim && intlim <= fromIntegral (maxBound :: Word32))
    = error ("Data.Numbers.Sieves.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 (p-1)
               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 (p-1) i
              else mult a ((p-1)*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 (p-1)
         loop p pows (n+1)
       else do
         writeArray pows n (x-1)

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

-- |  Retrieves the totient from the sieve
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)

-- |  Is the given number prime?
isPrime :: Integral a => PhiSieve -> a -> Bool
isPrime (PhiSieve a) n
   | even n     =  n == 2
   | otherwise  =  a ! (ix nn) == nn - 1
     where
       nn = fromIntegral n