-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.SieveBlock
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Bulk evaluation of arithmetic functions over continuous intervals
-- without factorisation.
--

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.ArithmeticFunctions.SieveBlock
  ( runFunctionOverBlock
  , SieveBlockConfig(..)
  , multiplicativeSieveBlockConfig
  , additiveSieveBlockConfig
  , sieveBlock
  , sieveBlockUnboxed
  , sieveBlockMoebius
  ) where

import Control.Monad (forM_, when)
import Control.Monad.ST (runST)
import Data.Bits
import Data.Coerce
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import GHC.Exts

import Math.NumberTheory.ArithmeticFunctions.Class
import Math.NumberTheory.ArithmeticFunctions.Moebius (Moebius, sieveBlockMoebius)
import Math.NumberTheory.Logarithms (wordLog2, integerLogBase')
import Math.NumberTheory.Primes
import Math.NumberTheory.Primes.Types
import Math.NumberTheory.Roots (integerSquareRoot)
import Math.NumberTheory.Utils.FromIntegral (wordToInt, intToWord)

-- | A record, which specifies a function to evaluate over a block.
--
-- For example, here is a configuration for the totient function:
--
-- > SieveBlockConfig
-- >   { sbcEmpty                = 1
-- >   , sbcFunctionOnPrimePower = \p a -> (unPrime p - 1) * unPrime p ^ (a - 1)
-- >   , sbcAppend               = (*)
-- >   }
data SieveBlockConfig a = SieveBlockConfig
  { forall a. SieveBlockConfig a -> a
sbcEmpty                :: a
    -- ^ value of a function on 1
  , forall a. SieveBlockConfig a -> Prime Word -> Word -> a
sbcFunctionOnPrimePower :: Prime Word -> Word -> a
    -- ^ how to evaluate a function on prime powers
  , forall a. SieveBlockConfig a -> a -> a -> a
sbcAppend               :: a -> a -> a
    -- ^ how to combine values of a function on coprime arguments
  }

-- | Create a config for a multiplicative function from its definition on prime powers.
multiplicativeSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig :: forall a. Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
multiplicativeSieveBlockConfig Prime Word -> Word -> a
f = SieveBlockConfig
  { sbcEmpty :: a
sbcEmpty                = a
1
  , sbcFunctionOnPrimePower :: Prime Word -> Word -> a
sbcFunctionOnPrimePower = Prime Word -> Word -> a
f
  , sbcAppend :: a -> a -> a
sbcAppend               = forall a. Num a => a -> a -> a
(*)
  }

-- | Create a config for an additive function from its definition on prime powers.
additiveSieveBlockConfig :: Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig :: forall a. Num a => (Prime Word -> Word -> a) -> SieveBlockConfig a
additiveSieveBlockConfig Prime Word -> Word -> a
f = SieveBlockConfig
  { sbcEmpty :: a
sbcEmpty                = a
0
  , sbcFunctionOnPrimePower :: Prime Word -> Word -> a
sbcFunctionOnPrimePower = Prime Word -> Word -> a
f
  , sbcAppend :: a -> a -> a
sbcAppend               = forall a. Num a => a -> a -> a
(+)
  }

-- | 'runFunctionOverBlock' @f@ @x@ @l@ evaluates an arithmetic function
-- for integers between @x@ and @x+l-1@ and returns a vector of length @l@.
-- It completely avoids factorisation, so it is asymptotically faster than
-- pointwise evaluation of @f@.
--
-- Value of @f@ at 0, if zero falls into block, is undefined.
--
-- Beware that for underlying non-commutative monoids the results may potentially
-- differ from pointwise application via 'runFunction'.
--
-- This is a thin wrapper over 'sieveBlock', read more details there.
--
-- >>> import Math.NumberTheory.ArithmeticFunctions
-- >>> runFunctionOverBlock carmichaelA 1 10
-- [1,1,2,2,4,2,6,2,6,4]
runFunctionOverBlock
  :: ArithmeticFunction Word a
  -> Word
  -> Word
  -> V.Vector a
runFunctionOverBlock :: forall a. ArithmeticFunction Word a -> Word -> Word -> Vector a
runFunctionOverBlock (ArithmeticFunction Prime Word -> Word -> m
f m -> a
g) = (forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map m -> a
g forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock SieveBlockConfig
  { sbcEmpty :: m
sbcEmpty                = forall a. Monoid a => a
mempty
  , sbcAppend :: m -> m -> m
sbcAppend               = forall a. Monoid a => a -> a -> a
mappend
  , sbcFunctionOnPrimePower :: Prime Word -> Word -> m
sbcFunctionOnPrimePower = coerce :: forall a b. Coercible a b => a -> b
coerce Prime Word -> Word -> m
f
  }

-- | Evaluate a function over a block in accordance to provided configuration.
-- Value of @f@ at 0, if zero falls into block, is undefined.
--
-- Based on Algorithm M of <https://arxiv.org/pdf/1305.1639.pdf Parity of the number of primes in a given interval and algorithms of the sublinear summation> by A. V. Lelechenko. See Lemma 2 on p. 5 on its algorithmic complexity. For the majority of use-cases its time complexity is O(x^(1+ε)).
--
-- For example, following code lists smallest prime factors:
--
-- >>> sieveBlock (SieveBlockConfig maxBound (\p _ -> unPrime p) min) 2 10 :: Data.Vector.Vector Word
-- [2,3,2,5,2,7,2,3,2,11]
--
-- And this is how to factorise all numbers in a block:
--
-- >>> sieveBlock (SieveBlockConfig [] (\p k -> [(unPrime p, k)]) (++)) 2 10 :: Data.Vector.Vector [(Word, Word)]
-- [[(2,1)],[(3,1)],[(2,2)],[(5,1)],[(2,1),(3,1)],[(7,1)],[(2,3)],[(3,2)],[(2,1),(5,1)],[(11,1)]]
sieveBlock
  :: forall v a.
     G.Vector v a
  => SieveBlockConfig a
  -> Word
  -> Word
  -> v a
sieveBlock :: forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock SieveBlockConfig a
_ Word
_ Word
0 = forall (v :: * -> *) a. Vector v a => v a
G.empty
sieveBlock (SieveBlockConfig a
empty Prime Word -> Word -> a
f a -> a -> a
append) !Word
lowIndex' Word
len' = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do

    let 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

        highIndex' :: Word
        highIndex' :: Word
highIndex' = Int -> Word
intToWord Int
highIndex

        ps :: [Int]
        ps :: [Int]
ps = if Int
highIndex forall a. Ord a => a -> a -> Bool
< Int
4 then [] else 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 (forall a. Integral a => a -> a
integerSquareRoot Int
highIndex)]

    MVector (PrimState (ST s)) Word
as <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate Int
len Word
1
    Mutable v (PrimState (ST s)) a
bs <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MG.replicate Int
len a
empty

    let doPrime :: Int -> ST s ()
doPrime Int
2 = do
          let fs :: Vector a
fs = forall a. Int -> (Int -> a) -> Vector a
V.generate (Word -> Int
wordLog2 Word
highIndex')
                (\Int
k -> Prime Word -> Word -> a
f (forall a. a -> Prime a
Prime Word
2) (Int -> Word
intToWord Int
k forall a. Num a => a -> a -> a
+ Word
1))
              npLow :: Word
npLow  = (Word
lowIndex' forall a. Num a => a -> a -> a
+ Word
1) forall a. Bits a => a -> Int -> a
`shiftR` Int
1
              npHigh :: Word
npHigh = Word
highIndex'      forall a. Bits a => a -> Int -> a
`shiftR` Int
1
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word
npLow .. Word
npHigh] forall a b. (a -> b) -> a -> b
$ \np :: Word
np@(W# Word#
np#) -> do
            let ix :: Int
ix = Word -> Int
wordToInt (Word
np forall a. Bits a => a -> Int -> a
`shiftL` Int
1) forall a. Num a => a -> a -> a
- Int
lowIndex :: Int
                tz :: Int
tz = Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
ctz# Word#
np#))
            forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector (PrimState (ST s)) Word
as (\Word
x -> Word
x forall a. Bits a => a -> Int -> a
`shiftL` (Int
tz forall a. Num a => a -> a -> a
+ Int
1)) Int
ix
            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v (PrimState (ST s)) a
bs (\a
y -> a
y a -> a -> a
`append` forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
fs Int
tz) Int
ix

        doPrime Int
p = do
          let p' :: Word
p' = Int -> Word
intToWord Int
p
              f0 :: a
f0 = Prime Word -> Word -> a
f (forall a. a -> Prime a
Prime Word
p') Word
1
              logp :: Int
logp = Integer -> Integer -> Int
integerLogBase' (forall a. Integral a => a -> Integer
toInteger Int
p) (forall a. Integral a => a -> Integer
toInteger Int
highIndex) forall a. Num a => a -> a -> a
- Int
1
              fs :: Vector a
fs = forall a. Int -> (Int -> a) -> Vector a
V.generate Int
logp (\Int
k -> Prime Word -> Word -> a
f (forall a. a -> Prime a
Prime Word
p') (Int -> Word
intToWord Int
k forall a. Num a => a -> a -> a
+ Word
2))
              npLow :: Int
npLow  = (Int
lowIndex forall a. Num a => a -> a -> a
+ Int
p forall a. Num a => a -> a -> a
- Int
1) forall a. Integral a => a -> a -> a
`quot` Int
p
              npHigh :: Int
npHigh = Int
highIndex          forall a. Integral a => a -> a -> a
`quot` Int
p

          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
npLow .. Int
npHigh] forall a b. (a -> b) -> a -> b
$ \Int
np -> do
            let !(I# Int#
ix#) = Int
np forall a. Num a => a -> a -> a
* Int
p forall a. Num a => a -> a -> a
- Int
lowIndex
                (Int
q, Int
r) = Int
np forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
p
            if Int
r forall a. Eq a => a -> a -> Bool
/= Int
0
            then do
              forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector (PrimState (ST s)) Word
as (forall a. Num a => a -> a -> a
* Word
p')        (Int# -> Int
I# Int#
ix#)
              forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v (PrimState (ST s)) a
bs (a -> a -> a
`append` a
f0) (Int# -> Int
I# Int#
ix#)
            else do
              let pow :: Word
pow = Int -> Int -> Word
highestPowerDividing Int
p Int
q
              forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MU.unsafeModify MVector (PrimState (ST s)) Word
as (\Word
x -> Word
x forall a. Num a => a -> a -> a
* Word
p' forall a b. (Num a, Integral b) => a -> b -> a
^ (Word
pow forall a. Num a => a -> a -> a
+ Word
2))                          (Int# -> Int
I# Int#
ix#)
              forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v (PrimState (ST s)) a
bs (\a
y -> a
y a -> a -> a
`append` forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
fs (Word -> Int
wordToInt Word
pow)) (Int# -> Int
I# Int#
ix#)

    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int]
ps Int -> ST s ()
doPrime

    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
k -> do
      Word
a <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MU.unsafeRead MVector (PrimState (ST s)) Word
as Int
k
      let a' :: Word
a' = Int -> Word
intToWord (Int
k forall a. Num a => a -> a -> a
+ Int
lowIndex)
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word
a forall a. Eq a => a -> a -> Bool
/= Word
a') forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MG.unsafeModify Mutable v (PrimState (ST s)) a
bs (\a
b -> a
b a -> a -> a
`append` Prime Word -> Word -> a
f (forall a. a -> Prime a
Prime forall a b. (a -> b) -> a -> b
$ Word
a' forall a. Integral a => a -> a -> a
`quot` Word
a) Word
1) Int
k

    forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v (PrimState (ST s)) a
bs

-- This is a variant of 'Math.NumberTheory.Utils.splitOff',
-- specialized for better performance.
highestPowerDividing :: Int -> Int -> Word
highestPowerDividing :: Int -> Int -> Word
highestPowerDividing !Int
_ Int
0 = Word
0
highestPowerDividing Int
p Int
n = forall {t}. Num t => t -> Int -> t
go Word
0 Int
n
  where
    go :: t -> Int -> t
go !t
k Int
m = case Int
m forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
p of
      (Int
q, Int
0) -> t -> Int -> t
go (t
k forall a. Num a => a -> a -> a
+ t
1) Int
q
      (Int, Int)
_      -> t
k

-- | This is 'sieveBlock' specialized to unboxed vectors.
--
-- >>> sieveBlockUnboxed (SieveBlockConfig 1 (\_ a -> a + 1) (*)) 1 10
-- [1,2,2,3,2,4,2,4,3,4]
sieveBlockUnboxed
  :: U.Unbox a
  => SieveBlockConfig a
  -> Word
  -> Word
  -> U.Vector a
sieveBlockUnboxed :: forall a. Unbox a => SieveBlockConfig a -> Word -> Word -> Vector a
sieveBlockUnboxed = forall (v :: * -> *) a.
Vector v a =>
SieveBlockConfig a -> Word -> Word -> v a
sieveBlock

{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Int  -> Word -> Word -> U.Vector Int  #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Word -> Word -> Word -> U.Vector Word #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Bool -> Word -> Word -> U.Vector Bool #-}
{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Moebius -> Word -> Word -> U.Vector Moebius #-}