content: {-# LANGUAGE LambdaCase #-}
  
  {-# OPTIONS_GHC -fno-warn-type-defaults #-}
  
  module Math.NumberTheory.SieveBlockBench
    ( benchSuite
    ) where
  
  import Test.Tasty.Bench
  import qualified Data.Vector as V
  import qualified Data.Vector.Unboxed as U
  
  import Math.NumberTheory.ArithmeticFunctions.Moebius
  import Math.NumberTheory.ArithmeticFunctions.SieveBlock
  import Math.NumberTheory.Primes
  
  blockLen :: Word
  blockLen = 1000000
  
  totientHelper :: Word -> Word -> Word
  totientHelper p 1 = p - 1
  totientHelper p 2 = (p - 1) * p
  totientHelper p k = (p - 1) * p ^ (k - 1)
  
  totientBlockConfig :: SieveBlockConfig Word
  totientBlockConfig = SieveBlockConfig
    { sbcEmpty                = 1
    , sbcAppend               = (*)
    , sbcFunctionOnPrimePower = totientHelper . unPrime
    }
  
  carmichaelHelper :: Word -> Word -> Word
  carmichaelHelper 2 1 = 1
  carmichaelHelper 2 2 = 2
  carmichaelHelper 2 k = 2 ^ (k - 2)
  carmichaelHelper p 1 = p - 1
  carmichaelHelper p 2 = (p - 1) * p
  carmichaelHelper p k = (p - 1) * p ^ (k - 1)
  
  carmichaelBlockConfig :: SieveBlockConfig Word
  carmichaelBlockConfig = SieveBlockConfig
    { sbcEmpty                = 1
    -- There is a specialized 'gcd' for Word, but not 'lcm'.
    , sbcAppend               = \x y -> (x `quot` gcd x y) * y
    , sbcFunctionOnPrimePower = carmichaelHelper . unPrime
    }
  
  moebiusConfig :: SieveBlockConfig Moebius
  moebiusConfig = SieveBlockConfig
    { sbcEmpty                = MoebiusP
    , sbcAppend               = (<>)
    , sbcFunctionOnPrimePower = const $ \case
        0 -> MoebiusP
        1 -> MoebiusN
        _ -> MoebiusZ
    }
  
  benchSuite :: Benchmark
  benchSuite = bgroup "SieveBlock"
    [ bgroup "totient"
      [ bench "boxed"   $ nf (V.sum . sieveBlock        totientBlockConfig 1) blockLen
      , bench "unboxed" $ nf (U.sum . sieveBlockUnboxed totientBlockConfig 1) blockLen
      ]
    , bgroup "carmichael"
      [ bench "boxed"   $ nf (V.sum . sieveBlock        carmichaelBlockConfig 1) blockLen
      , bench "unboxed" $ nf (U.sum . sieveBlockUnboxed carmichaelBlockConfig 1) blockLen
      ]
    , bgroup "moebius"
      [ bench "boxed"   $ nf (V.sum . V.map runMoebius . sieveBlock        moebiusConfig 1 :: Word -> Int) blockLen
      , bench "unboxed" $ nf (U.sum . U.map runMoebius . sieveBlockUnboxed moebiusConfig 1 :: Word -> Int) blockLen
      , bench "special" $ nf (U.sum . U.map runMoebius . sieveBlockMoebius 1 :: Word -> Int) blockLen
      ]
    ]
  