-- | Type class interface to different implementations of finite fields

{-# LANGUAGE BangPatterns, FlexibleContexts, TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables, TypeFamilies, TypeFamilyDependencies #-}
{-# LANGUAGE ExistentialQuantification, StandaloneDeriving #-}
{-# LANGUAGE DataKinds #-}

module Math.FiniteField.Class where

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

import Data.Bits
import Data.List

import GHC.TypeNats (Nat)

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import System.Random ( RandomGen )

import Math.FiniteField.TypeLevel
import Math.FiniteField.TypeLevel.Singleton

--------------------------------------------------------------------------------
-- * Fields

{-
-- NOTE: recursive classes just cause problems

-- | A class for witness types (these witness the existence of a field)
class 
  ( Show w, Field (FieldElem w)
  , WitnessPrime w ~ Prime (FieldElem w) 
  , WitnessDim   w ~ Dim   (FieldElem w) 
  ) => FieldWitness w 
  where
    type FieldElem    w = f | f -> w
    type WitnessPrime w :: Nat
    type WitnessDim   w :: Nat
-}

-- | A class for field element types 
class (Eq f, Ord f, Show f, Num f, Fractional f, Show (Witness f)) => Field f where
    -- | witness for the existence of the field (this is an injective type family!) 
    type Witness f = w | w -> f   
    -- | the characteristic at type level
    type Prime f :: Nat
    -- | the dimension at type level
    type Dim f :: Nat              
    -- | the prime characteristic
    characteristic   :: Witness f -> Integer   
    -- | dimension over the prime field (the exponent @m@ in @q=p^m@)
    dimension        :: Witness f -> Integer    -- TODO: this should be Int
    -- | the size (or order) of the field
    fieldSize        :: Witness f -> Integer   
    -- | The additive identity of the field
    zero             :: Witness f -> f
    -- | The multiplicative identity of the field
    one              :: Witness f -> f
    -- | check for equality with the additive identity
    isZero           :: f -> Bool
    -- | check for equality with the multiplicative identity
    isOne            :: f -> Bool
    -- | an element of the prime field
    embed            :: Witness f -> Integer -> f         
    embedSmall       :: Witness f -> Int     -> f   
    -- | a uniformly random field element
    randomFieldElem  :: RandomGen gen => Witness f -> gen -> (f,gen) 
    -- | a random invertible element
    randomInvertible :: RandomGen gen => Witness f -> gen -> (f,gen) 
    -- | a primitive generator
    primGen          :: Witness f -> f                    
    -- | extract t  he witness from a field element
    witnessOf        :: f -> Witness f                    
    -- | exponentiation 
    power            :: f -> Integer -> f                 
    powerSmall       :: f -> Int     -> f            
    -- | Frobenius automorphism @x -> x^p@
    frobenius        :: f -> f
    -- | list of field elements (of course it's only useful for very small fields)
    enumerate        :: Witness f -> [f]                  
  
    -- default implementations
  
    embedSmall !Witness f
w !Int
x = Witness f -> Integer -> f
forall f. Field f => Witness f -> Integer -> f
embed Witness f
w (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)
    powerSmall !f
x !Int
e = f -> Integer -> f
forall f. Field f => f -> Integer -> f
power f
x (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
e)      -- it's important not to use powerDefault here, because 'power' may be more efficient
    fieldSize  !Witness f
w    = Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic Witness f
w Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Witness f -> Integer
forall f. Field f => Witness f -> Integer
dimension Witness f
w
    power            = f -> Integer -> f
forall f. Field f => f -> Integer -> f
powerDefault 
    frobenius     !f
x = f -> Integer -> f
forall f. Field f => f -> Integer -> f
power f
x (Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic (f -> Witness f
forall f. Field f => f -> Witness f
witnessOf f
x))
  
    zero       !Witness f
w    = Witness f -> Int -> f
forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
0
    one        !Witness f
w    = Witness f -> Int -> f
forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
1
    -- isZero     !x    = (x == zero w)   -- we don't have a witness available here...
    -- isOne      !x    = (x == one  w)
  
    randomInvertible !Witness f
w !gen
g = case Witness f -> gen -> (f, gen)
forall f gen.
(Field f, RandomGen gen) =>
Witness f -> gen -> (f, gen)
randomFieldElem Witness f
w gen
g of 
      (f
x,gen
g') -> if f -> Bool
forall f. Field f => f -> Bool
isZero f
x then Witness f -> gen -> (f, gen)
forall f gen.
(Field f, RandomGen gen) =>
Witness f -> gen -> (f, gen)
randomInvertible Witness f
w gen
g' else (f
x,gen
g')

fieldPrimeSNat :: Field f => Witness f -> SNat (Prime f)
fieldPrimeSNat :: Witness f -> SNat (Prime f)
fieldPrimeSNat Witness f
w = Integer -> SNat (Prime f)
forall (n :: Nat). Integer -> SNat n
SNat (Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic Witness f
w)

fieldPrimeSNat64 :: Field f => Witness f -> SNat64 (Prime f)
fieldPrimeSNat64 :: Witness f -> SNat64 (Prime f)
fieldPrimeSNat64 Witness f
w = Word64 -> SNat64 (Prime f)
forall (n :: Nat). Word64 -> SNat64 n
SNat64 (Integer -> Word64
forall a. Num a => Integer -> a
fromInteger (Integer -> Word64) -> Integer -> Word64
forall a b. (a -> b) -> a -> b
$ Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic Witness f
w)

fieldDimSNat :: Field f => Witness f -> SNat (Dim f)
fieldDimSNat :: Witness f -> SNat (Dim f)
fieldDimSNat Witness f
w = Integer -> SNat (Dim f)
forall (n :: Nat). Integer -> SNat n
SNat (Witness f -> Integer
forall f. Field f => Witness f -> Integer
dimension Witness f
w)

fieldDimSNat64 :: Field f => Witness f -> SNat64 (Dim f)
fieldDimSNat64 :: Witness f -> SNat64 (Dim f)
fieldDimSNat64 Witness f
w = Word64 -> SNat64 (Dim f)
forall (n :: Nat). Word64 -> SNat64 n
SNat64 (Integer -> Word64
forall a. Num a => Integer -> a
fromInteger (Integer -> Word64) -> Integer -> Word64
forall a b. (a -> b) -> a -> b
$ Witness f -> Integer
forall f. Field f => Witness f -> Integer
dimension Witness f
w)

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

data SomeField 
  = forall f. Field f => SomeField (Witness f)

deriving instance Show SomeField

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

-- -- * Subfields
-- 
-- class 
--   ( FieldWitness (AmbientWitness w)
--   , FieldWitness (SubWitness     w) 
--   ) => SubFieldWitness w 
--   where
--     type AmbientWitness w :: *
--     type SubWitness     w :: *
-- 
-- class (Field (AmbientField s) , Field (SubField s)) => SubField s where
--   type AmbientField s :: *
--   type SubField     s :: *
--   type SubFieldWitness 
--   embedSubField :: SubField s -> AmbientField s 
--   mbSubField    :: AmbientField s -> Maybe (SubField s)
--   isSubField    :: AmbientField s -> Bool
--   coordinates   :: AmbientField s -> [SubField s]

--------------------------------------------------------------------------------
-- * Some generic functions

-- | Returns @"GF(p)"@ or @"GF(p^m)"
fieldName :: Field f => Witness f -> String
fieldName :: Witness f -> String
fieldName Witness f
w 
  | Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1     = String
"GF(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
p String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  | Bool
otherwise  = String
"GF(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
p String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
m String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  where
    p :: Integer
p = Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic Witness f
w
    m :: Integer
m = Witness f -> Integer
forall f. Field f => Witness f -> Integer
dimension      Witness f
w

-- | The multiplicate inverse (synonym for 'recip')
inverse :: Field f => f -> f
inverse :: f -> f
inverse = f -> f
forall a. Fractional a => a -> a
recip

-- | Enumerate the elements of the prime field only 
enumPrimeField :: forall f. Field f => Witness f -> [f]
enumPrimeField :: Witness f -> [f]
enumPrimeField Witness f
w = [ Witness f -> Int -> f
forall f. Field f => Witness f -> Int -> f
embedSmall Witness f
w Int
i | Int
i<-[Int
0..Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ] where
  pbig :: Integer
pbig = Witness f -> Integer
forall f. Field f => Witness f -> Integer
characteristic Witness f
w
  p :: Int
p    = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
pbig :: Int

-- | The nonzero elements in cyclic order, starting from the primitive generator
-- (of course this is only useful for very small fields)
multGroup :: Field f => Witness f -> [f]    
multGroup :: Witness f -> [f]
multGroup Witness f
w = (f -> f -> f) -> [f] -> [f]
forall a. (a -> a -> a) -> [a] -> [a]
scanl1 f -> f -> f
forall a. Num a => a -> a -> a
(*) [f]
list where
  g :: f
g    = Witness f -> f
forall f. Field f => Witness f -> f
primGen Witness f
w
  m :: Integer
m    = Witness f -> Integer
forall f. Field f => Witness f -> Integer
fieldSize Witness f
w
  list :: [f]
list = Int -> f -> [f]
forall a. Int -> a -> [a]
replicate (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) f
g

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

-- | Computes a table of discrete logarithms with respect to the primitive 
-- generator. Note: zero (that is, the additive identitiy of the field) 
-- is not present in the resulting map.
discreteLogTable :: forall f. Field f => Witness f -> Map f Int
discreteLogTable :: Witness f -> Map f Int
discreteLogTable Witness f
witness = [(f, Int)] -> Map f Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (Int -> f -> [(f, Int)]
worker Int
0 (Witness f -> f
forall f. Field f => Witness f -> f
one Witness f
witness)) where
  g :: f
g = Witness f -> f
forall f. Field f => Witness f -> f
primGen   Witness f
witness
  q :: Integer
q = Witness f -> Integer
forall f. Field f => Witness f -> Integer
fieldSize Witness f
witness
  qm1 :: Int
qm1 = Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  worker :: Int -> f -> [(f,Int)]
  worker :: Int -> f -> [(f, Int)]
worker !Int
e !f
acc
    | Int
e Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
qm1    = (f
acc,Int
e) (f, Int) -> [(f, Int)] -> [(f, Int)]
forall a. a -> [a] -> [a]
: Int -> f -> [(f, Int)]
worker (Int
eInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (f
accf -> f -> f
forall a. Num a => a -> a -> a
*f
g)
    | Bool
otherwise  = []

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

-- | Generic exponentiation
powerDefault :: forall f. (Field f) => f -> Integer -> f
powerDefault :: f -> Integer -> f
powerDefault !f
z !Integer
e 
  | f -> Bool
forall f. Field f => f -> Bool
isZero f
z  = f
z 
  | Integer
e Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0    = Witness f -> f
forall f. Field f => Witness f -> f
one Witness f
w
  | Integer
e Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0     = f -> Integer -> f
forall f. Field f => f -> Integer -> f
powerDefault (f -> f
forall a. Fractional a => a -> a
recip f
z) (Integer -> Integer
forall a. Num a => a -> a
negate Integer
e)
  | Integer
e Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
pm1  = f -> f -> Integer -> f
go (Witness f -> f
forall f. Field f => Witness f -> f
one Witness f
w) f
z (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod Integer
e Integer
pm1)
  | Bool
otherwise = f -> f -> Integer -> f
go (Witness f -> f
forall f. Field f => Witness f -> f
one Witness f
w) f
z Integer
e
  where
    w :: Witness f
w   = f -> Witness f
forall f. Field f => f -> Witness f
witnessOf f
z
    pm1 :: Integer
pm1 = Witness f -> Integer
forall f. Field f => Witness f -> Integer
fieldSize Witness f
w Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
    go :: f -> f -> Integer -> f
    go :: f -> f -> Integer -> f
go !f
acc !f
y !Integer
e = if Integer
e Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 
      then f
acc
      else case (Integer
e Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
1) of
        Integer
0 -> f -> f -> Integer -> f
go  f
acc    (f
yf -> f -> f
forall a. Num a => a -> a -> a
*f
y) (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
e Int
1)
        Integer
_ -> f -> f -> Integer -> f
go (f
accf -> f -> f
forall a. Num a => a -> a -> a
*f
y) (f
yf -> f -> f
forall a. Num a => a -> a -> a
*f
y) (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
e Int
1)

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