-- |
-- Module      : Math.NumberTheory.Prime.Count
-- Copyright   : 2021 Preetham Gujjula
-- License     : BSD-3-Clause
-- Maintainer  : libraries@mail.preetham.io
-- Stability   : experimental
--
-- This module provides a high-level, polymorphic interface to the primecount
-- library. For a lower-level interface, see
-- "Math.NumberTheory.Prime.Count.FFI#".
module Math.NumberTheory.Prime.Count
  ( primePi,
    primePiMaxBound,
    nthPrime,
    nthPrimeMaxBound,
    primePhi,
    getNumPrimecountThreads,
    setNumPrimecountThreads,
    primecountVersion,
  )
where

import Data.Int (Int64)
import Foreign.C.String (CString, peekCString, withCString)
import Foreign.C.Types (CSize)
import Foreign.Marshal.Array (allocaArray)
import Math.NumberTheory.Prime.Count.FFI
import System.IO.Unsafe (unsafePerformIO)
import Text.Read (readMaybe)

-- | The number of primes less than or equal to @n@. Throws an error if @n@
--   is larger than 'primePiMaxBound'. Also might throw an error if there's not
--   enough memory available to compute the result, which can happen even if @n@
--   is smaller than @primePiMaxBound@.
primePi :: Integral a => a -> a
primePi :: forall a. Integral a => a -> a
primePi a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = a
0
  | Integer
n' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
bound = Int64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int64
primecount_pi (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
n'))
  | Integer
n' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
primePiMaxBound = Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer
primePiStr Integer
n')
  | Bool
otherwise = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"primePi: input larger than primePiMaxBound"
  where
    bound :: Integer
    bound :: Integer
bound = Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger (Int64
forall a. Bounded a => a
maxBound :: Int64)

    n' :: Integer
    n' :: Integer
n' = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n

-- | The largest input supported by 'primePi'.
--
-- * 64-bit CPUs: @10^31@
-- * 32-bit CPUs: @2^63 - 1@
primePiMaxBound :: Integer
primePiMaxBound :: Integer
primePiMaxBound = [Char] -> Integer
forall a. Read a => [Char] -> a
read ([Char] -> Integer) -> [Char] -> Integer
forall a b. (a -> b) -> a -> b
$ IO [Char] -> [Char]
forall a. IO a -> a
unsafePerformIO (IO [Char] -> [Char]) -> IO [Char] -> [Char]
forall a b. (a -> b) -> a -> b
$ CString -> IO [Char]
peekCString CString
primecount_get_max_x
{-# NOINLINE primePiMaxBound #-}

primePiStr :: Integer -> Integer
primePiStr :: Integer -> Integer
primePiStr Integer
n = IO Integer -> Integer
forall a. IO a -> a
unsafePerformIO (IO Integer -> Integer) -> IO Integer -> Integer
forall a b. (a -> b) -> a -> b
$ do
  [Char] -> (CString -> IO Integer) -> IO Integer
forall a. [Char] -> (CString -> IO a) -> IO a
withCString (Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n) ((CString -> IO Integer) -> IO Integer)
-> (CString -> IO Integer) -> IO Integer
forall a b. (a -> b) -> a -> b
$ \CString
nString -> do
    let len :: CSize
len = CSize
32 :: CSize
    Int -> (CString -> IO Integer) -> IO Integer
forall a b. Storable a => Int -> (Ptr a -> IO b) -> IO b
allocaArray (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
len) ((CString -> IO Integer) -> IO Integer)
-> (CString -> IO Integer) -> IO Integer
forall a b. (a -> b) -> a -> b
$ \(CString
res :: CString) -> do
      Int
ret <- CString -> CString -> CSize -> IO Int
primecount_pi_str CString
nString CString
res CSize
len
      if Int
ret Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
ret Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
len
        then [Char] -> IO Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"primePi: call to primecount_pi_str failed"
        else do
          [Char]
answer <- CString -> IO [Char]
peekCString CString
res
          IO Integer
-> (Integer -> IO Integer) -> Maybe Integer -> IO Integer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            ([Char] -> IO Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"primePi: couldn't parse result of primecount_pi_str")
            Integer -> IO Integer
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ([Char] -> Maybe Integer
forall a. Read a => [Char] -> Maybe a
readMaybe [Char]
answer)
{-# NOINLINE primePiStr #-}

-- | The nth prime, starting at @nthPrime 1 == 2@.
--
--    * Throws an error if the input is less than 1.
--    * Throws an error if the input is larger than 'nthPrimeMaxBound`.
nthPrime :: Integral a => a -> a
nthPrime :: forall a. Integral a => a -> a
nthPrime a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1 = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"nthPrime: n must be >= 1"
  | Integer
n' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
bound = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"nthPrime: answer cannot be packed into a 64-bit int"
  | Bool
otherwise = Int64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int64
primecount_nth_prime (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
n'))
  where
    bound :: Integer
    bound :: Integer
bound = Integer
216289611853439384

    n' :: Integer
    n' :: Integer
n' = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n

-- | The largest input supported by 'nthPrime', equal to
--   @primePi ('maxBound' :: 'Int64') == 216289611853439384@.
nthPrimeMaxBound :: Integer
nthPrimeMaxBound :: Integer
nthPrimeMaxBound = Integer
216289611853439384

-- | @primePhi n a@ counts the number of positive integers @<= n@ that are not
--    divisible by any of the first @a@ primes. Throws an error if @n@ is larger
--    than @'maxBound' :: 'Int64'@.
primePhi :: Integral a => a -> a -> a
primePhi :: forall a. Integral a => a -> a -> a
primePhi a
n a
a
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 = a
0
  | a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 = a
n
  | Integer
n' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
bound = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"primePhi: input cannot be packed into a 64-bit int"
  | Bool
otherwise = Int64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int64 -> Int64
primecount_phi (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
n') (Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
a'))
  where
    bound :: Integer
    bound :: Integer
bound = Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger (Int64
forall a. Bounded a => a
maxBound :: Int64)

    n' :: Integer
    n' :: Integer
n' = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
n

    a' :: Integer
    a' :: Integer
a' = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
bound (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
a)

-- | Get the currently set number of threads used by @libprimecount@.
getNumPrimecountThreads :: IO Int
getNumPrimecountThreads :: IO Int
getNumPrimecountThreads = IO Int
primecount_get_num_threads

-- | Set the number of threads used by @libprimecount@. If the input is not
--   positive, the thread count is set to 1. If the input is greater than the
--   number of CPUs available, the thread count is set to the number of CPUs
--   available.
setNumPrimecountThreads :: Int -> IO ()
setNumPrimecountThreads :: Int -> IO ()
setNumPrimecountThreads = Int -> IO ()
primecount_set_num_threads

-- | Get the @libprimecount@ version number, in the form @"i.j"@
primecountVersion :: String
primecountVersion :: [Char]
primecountVersion = IO [Char] -> [Char]
forall a. IO a -> a
unsafePerformIO (CString -> IO [Char]
peekCString CString
primecount_version)
{-# NOINLINE primecountVersion #-}