    -- |
    -- 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               = a -> a -> a
    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               = a -> a -> a
    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) = ((m -> a) -> Vector m -> Vector a
    forall (v :: * -> *) a b.
    (Vector v a, Vector v b) =>
    (a -> b) -> v a -> v b
    G.map m -> a
    g (Vector m -> Vector a) -> (Word -> Vector m) -> Word -> Vector a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    .) ((Word -> Vector m) -> Word -> Vector a)
    -> (Word -> Word -> Vector m) -> Word -> Word -> Vector a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SieveBlockConfig m -> Word -> Word -> Vector m
    forall (v :: * -> *) a.
    Vector v a =>
    SieveBlockConfig a -> Word -> Word -> v a
    sieveBlock SieveBlockConfig
      { sbcEmpty :: m
    sbcEmpty                = m
    forall a. Monoid a => a
    mempty
      , sbcAppend :: m -> m -> m
    sbcAppend               = m -> m -> m
    forall a. Semigroup a => a -> a -> a
    (<>)
      , sbcFunctionOnPrimePower :: Prime Word -> Word -> m
    sbcFunctionOnPrimePower = (Prime Word -> Word -> m) -> Prime Word -> Word -> m
    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 = v a
    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 s. ST s (v a)) -> v a
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
    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 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    len Int -> Int -> Int
    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 Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    4 then [] else (Prime Int -> Int) -> [Prime Int] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Prime Int -> Int
    forall a. Prime a -> a
    unPrime [Int -> Prime Int
    forall a.
    (Bits a, Integral a, UniqueFactorisation a) =>
    a -> Prime a
    nextPrime Int
    2 .. Int -> Prime Int
    forall a.
    (Bits a, Integral a, UniqueFactorisation a) =>
    a -> Prime a
    precPrime (Int -> Int
    forall a. Integral a => a -> a
    integerSquareRoot Int
    highIndex)]

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

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

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

        forM_ ps doPrime

        forM_ [0 .. len - 1] $ \Int
    k -> do
          a <- MVector (PrimState (ST s)) Word -> Int -> ST s Word
    forall (m :: * -> *) a.
    (PrimMonad m, Unbox a) =>
    MVector (PrimState m) a -> Int -> m a
    MU.unsafeRead MVector s Word
    MVector (PrimState (ST s)) Word
    as Int
    k
          let a' = Int -> Word
    intToWord (Int
    k Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    lowIndex)
          when (a /= a') $
            MG.unsafeModify bs (\a
    b -> a
    b a -> a -> a
    `append` Prime Word -> Word -> a
    f (Word -> Prime Word
    forall a. a -> Prime a
    Prime (Word -> Prime Word) -> Word -> Prime Word
    forall a b. (a -> b) -> a -> b
    $ Word
    a' Word -> Word -> Word
    forall a. Integral a => a -> a -> a
    `quot` Word
    a) Word
    1) k

        G.unsafeFreeze 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 = Word -> Int -> Word
    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 Int -> Int -> (Int, Int)
    forall a. Integral a => a -> a -> (a, a)
    `quotRem` Int
    p of
          (Int
    q, Int
    0) -> t -> Int -> t
    go (t
    k t -> t -> t
    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 = SieveBlockConfig a -> Word -> Word -> Vector a
    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 #-}