-- |
-- Module:      Math.NumberTheory.Primes.Types
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- This is an internal module, defining types for primes.
-- Should not be exposed to users.
--

{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Math.NumberTheory.Primes.Types
  ( Prime(..)
  , toPrimeIntegral
  ) where

import Data.Bits
import GHC.Generics
import Control.DeepSeq
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed as U

import Math.NumberTheory.Utils.FromIntegral

-- $setup
-- >>> import Math.NumberTheory.Primes (nextPrime, precPrime)

-- | Wrapper for prime elements of @a@. It is supposed to be constructed
-- by 'Math.NumberTheory.Primes.nextPrime' / 'Math.NumberTheory.Primes.precPrime'.
-- and eliminated by 'unPrime'.
--
-- One can leverage 'Enum' instance to generate lists of primes.
-- Here are some examples.
--
-- *  Generate primes from the given interval:
--
--    >>> :set -XFlexibleContexts
--    >>> [nextPrime 101 .. precPrime 130]
--    [Prime 101,Prime 103,Prime 107,Prime 109,Prime 113,Prime 127]
--
-- *  Generate an infinite list of primes:
--
--    > [nextPrime 101 ..]
--    > [Prime 101,Prime 103,Prime 107,Prime 109,Prime 113,Prime 127...
--
-- *  Generate primes from the given interval of form p = 6k+5:
--
--    >>> [nextPrime 101, nextPrime 107 .. precPrime 150]
--    [Prime 101,Prime 107,Prime 113,Prime 131,Prime 137,Prime 149]
--
-- *  Get next prime:
--
--    >>> succ (nextPrime 101)
--    Prime 103
--
-- *  Get previous prime:
--
--    >>> pred (nextPrime 101)
--    Prime 97
--
-- *  Count primes less than a given number (cf. 'Math.NumberTheory.Primes.Counting.approxPrimeCount'):
--
--    >>> fromEnum (precPrime 100)
--    25
--
-- *  Get 25-th prime number (cf. 'Math.NumberTheory.Primes.Counting.nthPrimeApprox'):
--
--    >>> toEnum 25 :: Prime Int
--    Prime 97
--
newtype Prime a = Prime
  { Prime a -> a
unPrime :: a -- ^ Unwrap prime element.
  }
  deriving (Prime a -> Prime a -> Bool
(Prime a -> Prime a -> Bool)
-> (Prime a -> Prime a -> Bool) -> Eq (Prime a)
forall a. Eq a => Prime a -> Prime a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Prime a -> Prime a -> Bool
$c/= :: forall a. Eq a => Prime a -> Prime a -> Bool
== :: Prime a -> Prime a -> Bool
$c== :: forall a. Eq a => Prime a -> Prime a -> Bool
Eq, Eq (Prime a)
Eq (Prime a)
-> (Prime a -> Prime a -> Ordering)
-> (Prime a -> Prime a -> Bool)
-> (Prime a -> Prime a -> Bool)
-> (Prime a -> Prime a -> Bool)
-> (Prime a -> Prime a -> Bool)
-> (Prime a -> Prime a -> Prime a)
-> (Prime a -> Prime a -> Prime a)
-> Ord (Prime a)
Prime a -> Prime a -> Bool
Prime a -> Prime a -> Ordering
Prime a -> Prime a -> Prime a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Prime a)
forall a. Ord a => Prime a -> Prime a -> Bool
forall a. Ord a => Prime a -> Prime a -> Ordering
forall a. Ord a => Prime a -> Prime a -> Prime a
min :: Prime a -> Prime a -> Prime a
$cmin :: forall a. Ord a => Prime a -> Prime a -> Prime a
max :: Prime a -> Prime a -> Prime a
$cmax :: forall a. Ord a => Prime a -> Prime a -> Prime a
>= :: Prime a -> Prime a -> Bool
$c>= :: forall a. Ord a => Prime a -> Prime a -> Bool
> :: Prime a -> Prime a -> Bool
$c> :: forall a. Ord a => Prime a -> Prime a -> Bool
<= :: Prime a -> Prime a -> Bool
$c<= :: forall a. Ord a => Prime a -> Prime a -> Bool
< :: Prime a -> Prime a -> Bool
$c< :: forall a. Ord a => Prime a -> Prime a -> Bool
compare :: Prime a -> Prime a -> Ordering
$ccompare :: forall a. Ord a => Prime a -> Prime a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (Prime a)
Ord, (forall x. Prime a -> Rep (Prime a) x)
-> (forall x. Rep (Prime a) x -> Prime a) -> Generic (Prime a)
forall x. Rep (Prime a) x -> Prime a
forall x. Prime a -> Rep (Prime a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Prime a) x -> Prime a
forall a x. Prime a -> Rep (Prime a) x
$cto :: forall a x. Rep (Prime a) x -> Prime a
$cfrom :: forall a x. Prime a -> Rep (Prime a) x
Generic)

instance NFData a => NFData (Prime a)

instance Show a => Show (Prime a) where
  showsPrec :: Int -> Prime a -> ShowS
showsPrec Int
d (Prime a
p) String
r = (if Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10 then String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")" else String
s) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
r
    where
      s :: String
s = String
"Prime " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
p

newtype instance U.MVector s (Prime a) = MV_Prime (U.MVector s a)
newtype instance U.Vector    (Prime a) = V_Prime  (U.Vector    a)

instance U.Unbox a => U.Unbox (Prime a)

instance M.MVector U.MVector a => M.MVector U.MVector (Prime a) where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicInitialize #-}
  {-# INLINE basicUnsafeReplicate #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicClear #-}
  {-# INLINE basicSet #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE basicUnsafeGrow #-}
  basicLength :: MVector s (Prime a) -> Int
basicLength (MV_Prime v) = MVector s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
M.basicLength MVector s a
v
  basicUnsafeSlice :: Int -> Int -> MVector s (Prime a) -> MVector s (Prime a)
basicUnsafeSlice Int
i Int
n (MV_Prime v) = MVector s a -> MVector s (Prime a)
forall s a. MVector s a -> MVector s (Prime a)
MV_Prime (MVector s a -> MVector s (Prime a))
-> MVector s a -> MVector s (Prime a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> MVector s a -> MVector s a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
M.basicUnsafeSlice Int
i Int
n MVector s a
v
  basicOverlaps :: MVector s (Prime a) -> MVector s (Prime a) -> Bool
basicOverlaps (MV_Prime v1) (MV_Prime v2) = MVector s a -> MVector s a -> Bool
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> Bool
M.basicOverlaps MVector s a
v1 MVector s a
v2
  basicUnsafeNew :: Int -> ST s (MVector s (Prime a))
basicUnsafeNew Int
n = MVector s a -> MVector s (Prime a)
forall s a. MVector s a -> MVector s (Prime a)
MV_Prime (MVector s a -> MVector s (Prime a))
-> ST s (MVector s a) -> ST s (MVector s (Prime a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> ST s (MVector s a)
forall (v :: * -> * -> *) a s. MVector v a => Int -> ST s (v s a)
M.basicUnsafeNew Int
n
  basicInitialize :: MVector s (Prime a) -> ST s ()
basicInitialize (MV_Prime v) = MVector s a -> ST s ()
forall (v :: * -> * -> *) a s. MVector v a => v s a -> ST s ()
M.basicInitialize MVector s a
v
  basicUnsafeReplicate :: Int -> Prime a -> ST s (MVector s (Prime a))
basicUnsafeReplicate Int
n Prime a
x = MVector s a -> MVector s (Prime a)
forall s a. MVector s a -> MVector s (Prime a)
MV_Prime (MVector s a -> MVector s (Prime a))
-> ST s (MVector s a) -> ST s (MVector s (Prime a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> ST s (MVector s a)
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> a -> ST s (v s a)
M.basicUnsafeReplicate Int
n (Prime a -> a
forall a. Prime a -> a
unPrime Prime a
x)
  basicUnsafeRead :: MVector s (Prime a) -> Int -> ST s (Prime a)
basicUnsafeRead (MV_Prime v) Int
i = a -> Prime a
forall a. a -> Prime a
Prime (a -> Prime a) -> ST s a -> ST s (Prime a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector s a -> Int -> ST s a
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> ST s a
M.basicUnsafeRead MVector s a
v Int
i
  basicUnsafeWrite :: MVector s (Prime a) -> Int -> Prime a -> ST s ()
basicUnsafeWrite (MV_Prime v) Int
i Prime a
x = MVector s a -> Int -> a -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> a -> ST s ()
M.basicUnsafeWrite MVector s a
v Int
i (Prime a -> a
forall a. Prime a -> a
unPrime Prime a
x)
  basicClear :: MVector s (Prime a) -> ST s ()
basicClear (MV_Prime v) = MVector s a -> ST s ()
forall (v :: * -> * -> *) a s. MVector v a => v s a -> ST s ()
M.basicClear MVector s a
v
  basicSet :: MVector s (Prime a) -> Prime a -> ST s ()
basicSet (MV_Prime v) Prime a
x = MVector s a -> a -> ST s ()
forall (v :: * -> * -> *) a s. MVector v a => v s a -> a -> ST s ()
M.basicSet MVector s a
v (Prime a -> a
forall a. Prime a -> a
unPrime Prime a
x)
  basicUnsafeCopy :: MVector s (Prime a) -> MVector s (Prime a) -> ST s ()
basicUnsafeCopy (MV_Prime v1) (MV_Prime v2) = MVector s a -> MVector s a -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> ST s ()
M.basicUnsafeCopy MVector s a
v1 MVector s a
v2
  basicUnsafeMove :: MVector s (Prime a) -> MVector s (Prime a) -> ST s ()
basicUnsafeMove (MV_Prime v1) (MV_Prime v2) = MVector s a -> MVector s a -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> ST s ()
M.basicUnsafeMove MVector s a
v1 MVector s a
v2
  basicUnsafeGrow :: MVector s (Prime a) -> Int -> ST s (MVector s (Prime a))
basicUnsafeGrow (MV_Prime v) Int
n = MVector s a -> MVector s (Prime a)
forall s a. MVector s a -> MVector s (Prime a)
MV_Prime (MVector s a -> MVector s (Prime a))
-> ST s (MVector s a) -> ST s (MVector s (Prime a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector s a -> Int -> ST s (MVector s a)
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> ST s (v s a)
M.basicUnsafeGrow MVector s a
v Int
n

instance G.Vector U.Vector a => G.Vector U.Vector (Prime a) where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw #-}
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicUnsafeIndexM #-}
  {-# INLINE elemseq #-}
  basicUnsafeFreeze :: Mutable Vector s (Prime a) -> ST s (Vector (Prime a))
basicUnsafeFreeze (MV_Prime v) = Vector a -> Vector (Prime a)
forall a. Vector a -> Vector (Prime a)
V_Prime (Vector a -> Vector (Prime a))
-> ST s (Vector a) -> ST s (Vector (Prime a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable Vector s a -> ST s (Vector a)
forall (v :: * -> *) a s. Vector v a => Mutable v s a -> ST s (v a)
G.basicUnsafeFreeze MVector s a
Mutable Vector s a
v
  basicUnsafeThaw :: Vector (Prime a) -> ST s (Mutable Vector s (Prime a))
basicUnsafeThaw (V_Prime v) = MVector s a -> MVector s (Prime a)
forall s a. MVector s a -> MVector s (Prime a)
MV_Prime (MVector s a -> MVector s (Prime a))
-> ST s (MVector s a) -> ST s (MVector s (Prime a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector a -> ST s (Mutable Vector s a)
forall (v :: * -> *) a s. Vector v a => v a -> ST s (Mutable v s a)
G.basicUnsafeThaw Vector a
v
  basicLength :: Vector (Prime a) -> Int
basicLength (V_Prime v) = Vector a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.basicLength Vector a
v
  basicUnsafeSlice :: Int -> Int -> Vector (Prime a) -> Vector (Prime a)
basicUnsafeSlice Int
i Int
n (V_Prime v) = Vector a -> Vector (Prime a)
forall a. Vector a -> Vector (Prime a)
V_Prime (Vector a -> Vector (Prime a)) -> Vector a -> Vector (Prime a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector a -> Vector a
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.basicUnsafeSlice Int
i Int
n Vector a
v
  basicUnsafeIndexM :: Vector (Prime a) -> Int -> Box (Prime a)
basicUnsafeIndexM (V_Prime v) Int
i = a -> Prime a
forall a. a -> Prime a
Prime (a -> Prime a) -> Box a -> Box (Prime a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector a -> Int -> Box a
forall (v :: * -> *) a. Vector v a => v a -> Int -> Box a
G.basicUnsafeIndexM Vector a
v Int
i
  basicUnsafeCopy :: Mutable Vector s (Prime a) -> Vector (Prime a) -> ST s ()
basicUnsafeCopy (MV_Prime mv) (V_Prime v) = Mutable Vector s a -> Vector a -> ST s ()
forall (v :: * -> *) a s.
Vector v a =>
Mutable v s a -> v a -> ST s ()
G.basicUnsafeCopy MVector s a
Mutable Vector s a
mv Vector a
v
  elemseq :: Vector (Prime a) -> Prime a -> b -> b
elemseq Vector (Prime a)
_ = Prime a -> b -> b
seq

-- | Convert between primes of different types, similar in spirit to 'toIntegralSized'.
--
-- A simpler version of this function is:
--
-- > toPrimeIntegral :: (Integral a, Integral b) => a -> Maybe b
-- > toPrimeIntegral (Prime a)
-- >   | toInteger a == b = Just (Prime (fromInteger b))
-- >   | otherwise        = Nothing
-- >   where
-- >     b = toInteger a
--
-- The point of 'toPrimeIntegral' is to avoid redundant conversions and conditions,
-- when it is safe to do so, determining type sizes statically with 'bitSizeMaybe'.
-- For example, 'toPrimeIntegral' from 'Prime' 'Int' to 'Prime' 'Word' boils down to
-- 'Just' . 'fromIntegral'.
--
toPrimeIntegral :: (Integral a, Integral b, Bits a, Bits b) => Prime a -> Maybe (Prime b)
toPrimeIntegral :: Prime a -> Maybe (Prime b)
toPrimeIntegral (Prime a
a) = case b -> Maybe Int
forall a. Bits a => a -> Maybe Int
unsignedWidth b
b of
  Maybe Int
Nothing -> Maybe (Prime b)
res
  Just Int
bW -> case a -> Maybe Int
forall a. Bits a => a -> Maybe Int
unsignedWidth a
a of
    Just Int
aW
      | Int
aW Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bW -> Maybe (Prime b)
res
    Maybe Int
_
      | a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> a
forall a. Bits a => Int -> a
bit Int
bW a -> a -> a
forall a. Num a => a -> a -> a
- a
1 -> Maybe (Prime b)
res
      | Bool
otherwise       -> Maybe (Prime b)
forall a. Maybe a
Nothing
  where
    b :: b
b = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral' a
a
    res :: Maybe (Prime b)
res = Prime b -> Maybe (Prime b)
forall a. a -> Maybe a
Just (b -> Prime b
forall a. a -> Prime a
Prime b
b)
{-# INLINE toPrimeIntegral #-}

unsignedWidth :: Bits a => a -> Maybe Int
unsignedWidth :: a -> Maybe Int
unsignedWidth a
t
  | a -> Bool
forall a. Bits a => a -> Bool
isSigned a
t = Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1 (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Maybe Int
forall a. Bits a => a -> Maybe Int
bitSizeMaybe a
t
  | Bool
otherwise  =                a -> Maybe Int
forall a. Bits a => a -> Maybe Int
bitSizeMaybe a
t
{-# INLINE unsignedWidth #-}