-- | Arbitrary prime fields, naive implementation

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

module Math.FiniteField.PrimeField.Generic 
  ( -- * Witness for the existence of the field
    WitnessFp(..)
  , SomeWitnessFp(..)
  , mkPrimeField
  , unsafePrimeField
    -- * Field elements
  , Fp
  )
  where

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

import Data.Bits

import GHC.TypeNats (Nat)

import System.Random ( RandomGen , randomR )

import Math.FiniteField.Primes
import Math.FiniteField.TypeLevel
import Math.FiniteField.Class

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

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

-- | A witness for the existence of the prime field @F_p@
newtype WitnessFp (p :: Nat) 
  = WitnessFp { WitnessFp p -> IsPrime p
fromWitnessFp :: IsPrime p }
  deriving Int -> WitnessFp p -> ShowS
[WitnessFp p] -> ShowS
WitnessFp p -> String
(Int -> WitnessFp p -> ShowS)
-> (WitnessFp p -> String)
-> ([WitnessFp p] -> ShowS)
-> Show (WitnessFp p)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (p :: Nat). Int -> WitnessFp p -> ShowS
forall (p :: Nat). [WitnessFp p] -> ShowS
forall (p :: Nat). WitnessFp p -> String
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 not really practical...
--
-- But you can use 'unsafePrimeField' to cheat.
mkPrimeField :: Integer -> Maybe SomeWitnessFp
mkPrimeField :: Integer -> Maybe SomeWitnessFp
mkPrimeField Integer
p = case Integer -> SomeSNat
someSNat Integer
p of
  SomeSNat SNat n
sp -> (WitnessFp n -> SomeWitnessFp
forall (p :: Nat). WitnessFp p -> SomeWitnessFp
SomeWitnessFp (WitnessFp n -> SomeWitnessFp)
-> (IsPrime n -> WitnessFp n) -> IsPrime n -> SomeWitnessFp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IsPrime n -> WitnessFp n
forall (p :: Nat). IsPrime p -> WitnessFp p
WitnessFp) (IsPrime n -> SomeWitnessFp)
-> Maybe (IsPrime n) -> Maybe SomeWitnessFp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SNat n -> Maybe (IsPrime n)
forall (n :: Nat). SNat n -> Maybe (IsPrime n)
isPrime SNat n
sp 

-- | You are responsible for guaranteeing that the input is a prime.
unsafePrimeField :: Integer -> SomeWitnessFp
unsafePrimeField :: Integer -> SomeWitnessFp
unsafePrimeField Integer
p = case Integer -> SomeSNat
someSNat Integer
p of
  SomeSNat SNat n
sp -> WitnessFp n -> SomeWitnessFp
forall (p :: Nat). WitnessFp p -> SomeWitnessFp
SomeWitnessFp (IsPrime n -> WitnessFp n
forall (p :: Nat). IsPrime p -> WitnessFp p
WitnessFp (SNat n -> IsPrime n
forall (n :: Nat). SNat n -> IsPrime n
believeMeItsPrime SNat 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 !(IsPrime p) !Integer

fpWitness :: Fp p -> WitnessFp p
fpWitness :: Fp p -> WitnessFp p
fpWitness (Fp IsPrime p
p Integer
_) = IsPrime p -> WitnessFp p
forall (p :: Nat). IsPrime p -> WitnessFp p
WitnessFp IsPrime p
p

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

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

randomInvFp :: RandomGen gen => IsPrime p -> gen -> (Fp p, gen)
randomInvFp :: IsPrime p -> gen -> (Fp p, gen)
randomInvFp !IsPrime p
p !gen
gen = case (Integer, Integer) -> gen -> (Integer, gen)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Integer
1,Integer
qInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) gen
gen of { (Integer
x, gen
gen') -> (IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p Integer
x, gen
gen') } where !q :: Integer
q = IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p

-- | The order of the field
fpOrder :: Fp p -> Integer
fpOrder :: Fp p -> Integer
fpOrder (Fp IsPrime p
p Integer
_) = IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p

modp :: Integer -> IsPrime p -> Integer
modp :: Integer -> IsPrime p -> Integer
modp !Integer
x !IsPrime p
p = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod Integer
x (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p)

instance Eq (Fp p) where
  == :: Fp p -> Fp p -> Bool
(==) (Fp IsPrime p
_ Integer
x) (Fp IsPrime p
_ Integer
y) = Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
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 IsPrime p
_ Integer
x) (Fp IsPrime p
_ Integer
y) = Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Integer
x Integer
y

instance Show (Fp p) where
  show :: Fp p -> String
show (Fp IsPrime p
p Integer
k) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" mod " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

instance Num (Fp p) where
  fromInteger :: Integer -> Fp p
fromInteger = String -> Integer -> Fp p
forall a. HasCallStack => String -> a
error String
"PrimeField/Generic/Fp/fromInteger: not defined; use `fp` instead" 
  negate :: Fp p -> Fp p
negate (Fp IsPrime p
p Integer
x) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer
Raw.neg (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x)
  + :: Fp p -> Fp p -> Fp p
(+) (Fp IsPrime p
p Integer
x) (Fp IsPrime p
_ Integer
y) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer -> Integer
Raw.add (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x Integer
y)
  (-) (Fp IsPrime p
p Integer
x) (Fp IsPrime p
_ Integer
y) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer -> Integer
Raw.sub (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x Integer
y)
  * :: Fp p -> Fp p -> Fp p
(*) (Fp IsPrime p
p Integer
x) (Fp IsPrime p
_ Integer
y) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer -> Integer
Raw.mul (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x Integer
y)
  abs :: Fp p -> Fp p
abs = Fp p -> Fp p
forall a. a -> a
id
  signum :: Fp p -> Fp p
signum (Fp IsPrime p
p Integer
_) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p Integer
1

instance Fractional (Fp p) where
  fromRational :: Rational -> Fp p
fromRational = String -> Rational -> Fp p
forall a. HasCallStack => String -> a
error String
"PrimeField/Generic/Fp/fromRational: not defined; use `fp` instead" 
  recip :: Fp p -> Fp p
recip (Fp IsPrime p
p Integer
x)          = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer
Raw.inv (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x)
  / :: Fp p -> Fp p -> Fp p
(/)   (Fp IsPrime p
p Integer
x) (Fp IsPrime p
_ Integer
y) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer -> Integer
Raw.div (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x Integer
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 p -> IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime 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 p -> IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p }
  enumerate :: Witness (Fp p) -> [Fp p]
enumerate         Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp p -> let q :: Integer
q = IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p in [ IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
fp IsPrime p
p Integer
k | Integer
k<-[Integer
0..Integer
qInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1] ] }
  embed :: Witness (Fp p) -> Integer -> Fp p
embed           Witness (Fp p)
w Integer
x = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
fp (WitnessFp p -> IsPrime p
forall (p :: Nat). WitnessFp p -> IsPrime p
fromWitnessFp Witness (Fp p)
WitnessFp p
w) Integer
x
  primGen :: Witness (Fp p) -> Fp p
primGen             = String -> WitnessFp p -> Fp p
forall a. HasCallStack => String -> a
error String
"PrimeField/Generic/Fp: primGen: not implemented"
  witnessOf :: Fp p -> Witness (Fp p)
witnessOf           = Fp p -> Witness (Fp p)
forall (p :: Nat). Fp p -> WitnessFp p
fpWitness
  power :: Fp p -> Integer -> Fp p
power               = Fp p -> Integer -> Fp p
forall (p :: Nat). Fp p -> Integer -> Fp p
fpPow
  frobenius :: Fp p -> Fp p
frobenius           = Fp p -> Fp p
forall a. a -> a
id
  randomFieldElem :: Witness (Fp p) -> gen -> (Fp p, gen)
randomFieldElem   Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp p -> IsPrime p -> gen -> (Fp p, gen)
forall gen (p :: Nat).
RandomGen gen =>
IsPrime p -> gen -> (Fp p, gen)
randomFp    IsPrime p
p } 
  randomInvertible :: Witness (Fp p) -> gen -> (Fp p, gen)
randomInvertible  Witness (Fp p)
w = case Witness (Fp p)
w of { WitnessFp p -> IsPrime p -> gen -> (Fp p, gen)
forall gen (p :: Nat).
RandomGen gen =>
IsPrime p -> gen -> (Fp p, gen)
randomInvFp IsPrime p
p }
  zero :: Witness (Fp p) -> Fp p
zero Witness (Fp p)
w = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp (WitnessFp p -> IsPrime p
forall (p :: Nat). WitnessFp p -> IsPrime p
fromWitnessFp Witness (Fp p)
WitnessFp p
w) Integer
0
  one :: Witness (Fp p) -> Fp p
one  Witness (Fp p)
w = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp (WitnessFp p -> IsPrime p
forall (p :: Nat). WitnessFp p -> IsPrime p
fromWitnessFp Witness (Fp p)
WitnessFp p
w) Integer
1
  isZero :: Fp p -> Bool
isZero (Fp IsPrime p
_ Integer
x) = (Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0)
  isOne :: Fp p -> Bool
isOne  (Fp IsPrime p
_ Integer
x) = (Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1)

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

-- | Powers
fpPow :: Fp p -> Integer -> Fp p
fpPow :: Fp p -> Integer -> Fp p
fpPow (Fp IsPrime p
p Integer
x) Integer
e = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer -> Integer
Raw.pow (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
x Integer
e)

-- | Inversion (using Euclid's algorithm)
fpInv :: Fp p -> Fp p
fpInv :: Fp p -> Fp p
fpInv (Fp IsPrime p
p Integer
a) = IsPrime p -> Integer -> Fp p
forall (p :: Nat). IsPrime p -> Integer -> Fp p
Fp IsPrime p
p (Integer -> Integer -> Integer
Raw.inv (IsPrime p -> Integer
forall (n :: Nat). IsPrime n -> Integer
fromPrime IsPrime p
p) Integer
a)

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

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

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