-- | Small prime fields, up to @p < 2^31@ (using @Word64@ under the hood).
--
-- This should be faster than the generic implementation which
-- uses @Integer@ under the hood.
--
-- NB: Because the multiplication of two 32 bit integers needs 64 bits,
-- the limit is @2^32@ and not @2^64@. And because I'm lazy right now
-- to check if everything works properly unsigned, the actual limit 
-- is @2^31@ instead :)
--

{-# LANGUAGE BangPatterns, DataKinds, KindSignatures, TypeFamilies #-}
{-# LANGUAGE ExistentialQuantification, StandaloneDeriving #-}

module Math.FiniteField.PrimeField.Small
  ( -- * Witness for the existence of the field
    WitnessFp(..)
  , SomeWitnessFp(..)
  , mkSmallPrimeField
  , unsafeSmallPrimeField
    -- * Field elements
  , Fp
  , primRoot
  )
  where

--------------------------------------------------------------------------------

import Data.Bits
import Data.Int
import Data.Word

import GHC.TypeNats (Nat)

import System.Random ( RandomGen , randomR )

import Math.FiniteField.Primes
import Math.FiniteField.TypeLevel
import Math.FiniteField.Class 
import Math.FiniteField.Conway.Internal ( lookupConwayPrimRoot_ )

import qualified Math.FiniteField.PrimeField.Small.Raw as Raw

--------------------------------------------------------------------------------

-- | A witness for the existence of the prime field @F_p@
newtype WitnessFp (p :: Nat) 
  = WitnessFp { forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp :: IsSmallPrime p }
  deriving Int -> WitnessFp p -> ShowS
forall (p :: Nat). Int -> WitnessFp p -> ShowS
forall (p :: Nat). [WitnessFp p] -> ShowS
forall (p :: Nat). WitnessFp p -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [WitnessFp p] -> ShowS
$cshowList :: forall (p :: Nat). [WitnessFp p] -> ShowS
show :: WitnessFp p -> String
$cshow :: forall (p :: Nat). WitnessFp p -> String
showsPrec :: Int -> WitnessFp p -> ShowS
$cshowsPrec :: forall (p :: Nat). Int -> WitnessFp p -> ShowS
Show

data SomeWitnessFp 
  = forall p. SomeWitnessFp (WitnessFp p) 

deriving instance Show SomeWitnessFp

-- | Note: currently this checks the primality of the input using
-- trial division, so it's only practical for (very) small primes...
--
-- But you can use 'unsafeSmallPrimeField' to cheat.
mkSmallPrimeField :: Int -> Maybe SomeWitnessFp
mkSmallPrimeField :: Int -> Maybe SomeWitnessFp
mkSmallPrimeField Int
p = case Int64 -> SomeSNat64
someSNat64 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p) of
  SomeSNat64 SNat64 n
sp -> (forall (p :: Nat). WitnessFp p -> SomeWitnessFp
SomeWitnessFp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: Nat). IsSmallPrime p -> WitnessFp p
WitnessFp) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (n :: Nat). SNat64 n -> Maybe (IsSmallPrime n)
isSmallPrime SNat64 n
sp 

-- | You are responsible for guaranteeing that the input is a prime.
unsafeSmallPrimeField :: Int -> SomeWitnessFp
unsafeSmallPrimeField :: Int -> SomeWitnessFp
unsafeSmallPrimeField Int
p = case Int64 -> SomeSNat64
someSNat64 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p) of
  SomeSNat64 SNat64 n
sp -> forall (p :: Nat). WitnessFp p -> SomeWitnessFp
SomeWitnessFp (forall (p :: Nat). IsSmallPrime p -> WitnessFp p
WitnessFp (forall (n :: Nat). SNat64 n -> IsSmallPrime n
believeMeItsASmallPrime SNat64 n
sp))

-- instance FieldWitness (WitnessFp p) where
--   type FieldElem    (WitnessFp p) = Fp p
--   type WitnessPrime (WitnessFp p) = p
--   type WitnessDim   (WitnessFp p) = 1

--------------------------------------------------------------------------------

-- | An element of the prime field @F_p@
data Fp (p :: Nat) 
  = Fp {-# UNPACK #-} !(IsSmallPrime p) {-# UNPACK #-} !Word64

fpWitness :: Fp p -> WitnessFp p
fpWitness :: forall (p :: Nat). Fp p -> WitnessFp p
fpWitness (Fp IsSmallPrime p
p Word64
_) = forall (p :: Nat). IsSmallPrime p -> WitnessFp p
WitnessFp IsSmallPrime p
p

-- | Constructing elements
fp :: IsSmallPrime p -> Word64 -> Fp p
fp :: forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
fp !IsSmallPrime p
p !Word64
n 
  | Word64
n forall a. Ord a => a -> a -> Bool
>= Word64
0 Bool -> Bool -> Bool
&& Word64
n forall a. Ord a => a -> a -> Bool
< Word64
q  = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p Word64
n
  | Bool
otherwise        = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (forall a. Integral a => a -> a -> a
mod Word64
n Word64
q)
  where
    !q :: Word64
q = forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p

randomFp :: RandomGen gen => IsSmallPrime p -> gen -> (Fp p, gen)
randomFp :: forall gen (p :: Nat).
RandomGen gen =>
IsSmallPrime p -> gen -> (Fp p, gen)
randomFp !IsSmallPrime p
p !gen
gen = case forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Word64
0,Word64
qforall a. Num a => a -> a -> a
-Word64
1) gen
gen of { (Word64
x, gen
gen') -> (forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p Word64
x, gen
gen') } where !q :: Word64
q = forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p

randomInvFp :: RandomGen gen => IsSmallPrime p -> gen -> (Fp p, gen)
randomInvFp :: forall gen (p :: Nat).
RandomGen gen =>
IsSmallPrime p -> gen -> (Fp p, gen)
randomInvFp !IsSmallPrime p
p !gen
gen = case forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Word64
1,Word64
qforall a. Num a => a -> a -> a
-Word64
1) gen
gen of { (Word64
x, gen
gen') -> (forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p Word64
x, gen
gen') } where !q :: Word64
q = forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p

-- | The order of the field
fpOrder :: Fp p -> Word64
fpOrder :: forall (p :: Nat). Fp p -> Word64
fpOrder (Fp IsSmallPrime p
p Word64
_) = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p)

modp :: Word64 -> IsSmallPrime p -> Word64
modp :: forall (p :: Nat). Word64 -> IsSmallPrime p -> Word64
modp !Word64
x !IsSmallPrime p
p = forall a. Integral a => a -> a -> a
mod Word64
x (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p)

modpInteger :: Integer -> IsSmallPrime p -> Word64
modpInteger :: forall (p :: Nat). Integer -> IsSmallPrime p -> Word64
modpInteger Integer
x IsSmallPrime p
p = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Integral a => a -> a -> a
mod Integer
x (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p)))

modpSigned :: Int64 -> IsSmallPrime p -> Int64
modpSigned :: forall (p :: Nat). Int64 -> IsSmallPrime p -> Int64
modpSigned Int64
x IsSmallPrime p
p = forall a. Integral a => a -> a -> a
mod Int64
x (forall (n :: Nat). IsSmallPrime n -> Int64
fromSmallPrimeSigned IsSmallPrime p
p)

primRoot :: IsSmallPrime p -> Fp p 
primRoot :: forall (p :: Nat). IsSmallPrime p -> Fp p
primRoot IsSmallPrime p
p = case Int -> Maybe Word64
lookupConwayPrimRoot_ (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p)) of
  Just Word64
g     -> forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
fp IsSmallPrime p
p (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
g)
  Maybe Word64
Nothing    -> forall a. HasCallStack => String -> a
error String
"PrimeField/Small/Fp/primRoot: primitive generator not found in the Conway table"

-- | Enumarte all field elements
enumerateFp :: IsSmallPrime p -> [Fp p]
enumerateFp :: forall (p :: Nat). IsSmallPrime p -> [Fp p]
enumerateFp IsSmallPrime p
p = [ forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p Word64
k | Word64
k <- [Word64
0..forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
pforall a. Num a => a -> a -> a
-Word64
1] ]

--------------------------------------------------------------------------------

instance Eq (Fp p) where
  == :: Fp p -> Fp p -> Bool
(==) (Fp IsSmallPrime p
_ Word64
x) (Fp IsSmallPrime p
_ Word64
y) = Word64
x forall a. Eq a => a -> a -> Bool
== Word64
y

-- | Note: the @Ord@ instance is present only so that you can use 'GF' as key
-- in @Maps@ - the ordering is kind of arbitrary! 
instance Ord (Fp p) where
  compare :: Fp p -> Fp p -> Ordering
compare (Fp IsSmallPrime p
_ Word64
x) (Fp IsSmallPrime p
_ Word64
y) = forall a. Ord a => a -> a -> Ordering
compare Word64
x Word64
y

instance Show (Fp p) where
  show :: Fp p -> String
show (Fp IsSmallPrime p
p Word64
k) = String
"(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Word64
k forall a. [a] -> [a] -> [a]
++ String
" mod " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) forall a. [a] -> [a] -> [a]
++ String
")"

instance Num (Fp p) where
  fromInteger :: Integer -> Fp p
fromInteger = forall a. HasCallStack => String -> a
error String
"PrimeField/Small/Fp/fromInteger: cannot be defined; use `embed` instead" 
  negate :: Fp p -> Fp p
negate (Fp IsSmallPrime p
p Word64
x) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64
Raw.neg (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x)
  + :: Fp p -> Fp p -> Fp p
(+) (Fp IsSmallPrime p
p Word64
x) (Fp IsSmallPrime p
_ Word64
y) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.add (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Word64
y)
  (-) (Fp IsSmallPrime p
p Word64
x) (Fp IsSmallPrime p
_ Word64
y) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.sub (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Word64
y)
  * :: Fp p -> Fp p -> Fp p
(*) (Fp IsSmallPrime p
p Word64
x) (Fp IsSmallPrime p
_ Word64
y) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.mul (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Word64
y)
  abs :: Fp p -> Fp p
abs = forall a. a -> a
id
  signum :: Fp p -> Fp p
signum (Fp IsSmallPrime p
p Word64
_) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p Word64
1

instance Fractional (Fp p) where
  fromRational :: Rational -> Fp p
fromRational = forall a. HasCallStack => String -> a
error String
"PrimeField/Small/Fp/fromRational: cannot be defined; use `embed` instead" 
  recip :: Fp p -> Fp p
recip (Fp IsSmallPrime p
p Word64
x)          = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64
Raw.inv (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x)
  / :: Fp p -> Fp p -> Fp p
(/)   (Fp IsSmallPrime p
p Word64
x) (Fp IsSmallPrime p
_ Word64
y) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.div (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Word64
y)

instance Field (Fp p) where
  type Witness (Fp p) = WitnessFp p
  type Prime   (Fp p) = p
  type Dim     (Fp p) = 1
  characteristic :: Witness (Fp p) -> Integer
characteristic     Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp IsSmallPrime p
p -> forall (n :: Nat). IsSmallPrime n -> Integer
fromSmallPrimeInteger IsSmallPrime p
p }
  dimension :: Witness (Fp p) -> Integer
dimension          Witness (Fp p)
_ = Integer
1
  fieldSize :: Witness (Fp p) -> Integer
fieldSize          Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp IsSmallPrime p
p -> forall (n :: Nat). IsSmallPrime n -> Integer
fromSmallPrimeInteger IsSmallPrime p
p }
  enumerate :: Witness (Fp p) -> [Fp p]
enumerate          Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp IsSmallPrime p
p -> let q :: Word64
q = forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p in [ forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
fp IsSmallPrime p
p Word64
k | Word64
k<-[Word64
0..Word64
qforall a. Num a => a -> a -> a
-Word64
1] ] }
  embed :: Witness (Fp p) -> Integer -> Fp p
embed            Witness (Fp p)
w Integer
x = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
fp (forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp Witness (Fp p)
w) (forall a. Num a => Integer -> a
fromInteger  Integer
x)
  embedSmall :: Witness (Fp p) -> Int -> Fp p
embedSmall       Witness (Fp p)
w Int
x = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
fp (forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp Witness (Fp p)
w) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)
  primGen :: Witness (Fp p) -> Fp p
primGen            Witness (Fp p)
w = forall (p :: Nat). IsSmallPrime p -> Fp p
primRoot (forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp Witness (Fp p)
w)
  witnessOf :: Fp p -> Witness (Fp p)
witnessOf            = forall (p :: Nat). Fp p -> WitnessFp p
fpWitness
  power :: Fp p -> Integer -> Fp p
power                = forall (p :: Nat). Fp p -> Integer -> Fp p
fpPow
  powerSmall :: Fp p -> Int -> Fp p
powerSmall       Fp p
x Int
e = forall (p :: Nat). Fp p -> Int64 -> Fp p
fpPow_ Fp p
x (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
e)
  frobenius :: Fp p -> Fp p
frobenius            = forall a. a -> a
id
  randomFieldElem :: forall gen. RandomGen gen => Witness (Fp p) -> gen -> (Fp p, gen)
randomFieldElem    Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp IsSmallPrime p
p -> forall gen (p :: Nat).
RandomGen gen =>
IsSmallPrime p -> gen -> (Fp p, gen)
randomFp    IsSmallPrime p
p } 
  randomInvertible :: forall gen. RandomGen gen => Witness (Fp p) -> gen -> (Fp p, gen)
randomInvertible   Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp IsSmallPrime p
p -> forall gen (p :: Nat).
RandomGen gen =>
IsSmallPrime p -> gen -> (Fp p, gen)
randomInvFp IsSmallPrime p
p }
  zero :: Witness (Fp p) -> Fp p
zero Witness (Fp p)
w = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp (forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp Witness (Fp p)
w) Word64
0
  one :: Witness (Fp p) -> Fp p
one  Witness (Fp p)
w = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp (forall (p :: Nat). WitnessFp p -> IsSmallPrime p
fromWitnessFp Witness (Fp p)
w) Word64
1
  isZero :: Fp p -> Bool
isZero (Fp IsSmallPrime p
_ Word64
x) = (Word64
x forall a. Eq a => a -> a -> Bool
== Word64
0)
  isOne :: Fp p -> Bool
isOne  (Fp IsSmallPrime p
_ Word64
x) = (Word64
x forall a. Eq a => a -> a -> Bool
== Word64
1)

--------------------------------------------------------------------------------
-- * Nontrivial operations

-- | Powers
fpPow_ :: Fp p -> Int64 -> Fp p
fpPow_ :: forall (p :: Nat). Fp p -> Int64 -> Fp p
fpPow_ (Fp IsSmallPrime p
p Word64
x) Int64
e = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Int64 -> Word64
Raw.pow (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Int64
e)

fpPow :: Fp p -> Integer -> Fp p
fpPow :: forall (p :: Nat). Fp p -> Integer -> Fp p
fpPow (Fp IsSmallPrime p
p Word64
x) Integer
e = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Integer -> Word64
Raw.pow' (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
x Integer
e)

-- | Inversion (using Euclid's algorithm)
fpInv :: Fp p -> Fp p
fpInv :: forall (p :: Nat). Fp p -> Fp p
fpInv (Fp IsSmallPrime p
p Word64
a) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64
Raw.inv (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
a)

-- | Division via Euclid's algorithm
fpDiv :: Fp p -> Fp p -> Fp p
fpDiv :: forall (p :: Nat). Fp p -> Fp p -> Fp p
fpDiv (Fp IsSmallPrime p
p Word64
a) (Fp IsSmallPrime p
_ Word64
b) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.div (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
a Word64
b)

-- | Division via multiplying by the inverse
fpDiv2 :: Fp p -> Fp p -> Fp p
fpDiv2 :: forall (p :: Nat). Fp p -> Fp p -> Fp p
fpDiv2 (Fp IsSmallPrime p
p Word64
a) (Fp IsSmallPrime p
_ Word64
b) = forall (p :: Nat). IsSmallPrime p -> Word64 -> Fp p
Fp IsSmallPrime p
p (Word64 -> Word64 -> Word64 -> Word64
Raw.div2 (forall (n :: Nat). IsSmallPrime n -> Word64
fromSmallPrime IsSmallPrime p
p) Word64
a Word64
b)

--------------------------------------------------------------------------------