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

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

module Math.FiniteField.Class where

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

import Data.Bits
import Data.List

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

import System.Random ( RandomGen )

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

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 = r | r -> f                 
  -- | the prime characteristic
  characteristic   :: Witness f -> Integer   
  -- | dimension over the prime field (the exponent @m@ in @q=p^m@)
  dimension        :: Witness f -> Integer   
  -- | 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            
  -- | 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)
  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

  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')

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

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

deriving instance Show SomeField

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

-- | 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 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)

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