-- | Table of Conway polynomials
--
-- The data is from <http://www.math.rwth-aachen.de/~Frank.Luebeck/data/ConwayPol/index.html>
--
-- The text file from the above link is converted to a C source file during the 
-- build process, and the resulting object file is then linked.
-- This should be much faster (in compilation time) than encoding it as a Haskell 
-- data structure, because GHC does not like big constants. Also the in-memory representation
-- can be controlled more precisely.
--

{-# LANGUAGE ForeignFunctionInterface, BangPatterns, KindSignatures, GADTs #-}
module Math.FiniteField.Conway.Internal where

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

import Data.Word
import Data.Bits

import GHC.TypeNats (Nat)

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

import Foreign.C
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal
import Foreign.Marshal.Array

import qualified System.IO.Unsafe as Unsafe

--------------------------------------------------------------------------------
-- * The witness (it's here so internal modules can acces the inside of it)

newtype ConwayPoly (p :: Nat) (m :: Nat) where
  ConwayWitness :: Ptr Word32 -> ConwayPoly p m

fromConwayPoly :: ConwayPoly p m -> Ptr Word32
fromConwayPoly :: forall (p :: Nat) (m :: Nat). ConwayPoly p m -> Ptr Word32
fromConwayPoly (ConwayWitness Ptr Word32
ptr) = Ptr Word32
ptr

-- | @(prime,exponent)@
conwayParams_ :: ConwayPoly p m -> (Word64,Int)
conwayParams_ :: forall (p :: Nat) (m :: Nat). ConwayPoly p m -> (Word64, Int)
conwayParams_ (ConwayWitness Ptr Word32
ptr) = forall a. IO a -> a
Unsafe.unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
  (Word64
p,Int
m) <- Ptr Word32 -> IO (Word64, Int)
getConwayEntryParams Ptr Word32
ptr
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
p, forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m)

conwayPrime_ :: ConwayPoly p m -> Word64
conwayPrime_ :: forall (p :: Nat) (m :: Nat). ConwayPoly p m -> Word64
conwayPrime_ (ConwayWitness Ptr Word32
ptr) = forall a. IO a -> a
Unsafe.unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
  (Word64
p,Int
_) <- Ptr Word32 -> IO (Word64, Int)
getConwayEntryParams Ptr Word32
ptr
  forall (m :: * -> *) a. Monad m => a -> m a
return Word64
p

-- | We have some Conway polynomials for @m=1@ too; the roots of 
-- these linear polynomials are primitive roots in @F_p@
lookupConwayPrimRoot_ :: Int -> Maybe Word64
lookupConwayPrimRoot_ :: Int -> Maybe Word64
lookupConwayPrimRoot_ !Int
p = case forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (Int -> Int -> Int
encodePrimeExpo Int
p Int
1) (ConwayTable -> IntMap (Ptr Word32)
fromConwayTable ConwayTable
theConwayTable) of
  Maybe (Ptr Word32)
Nothing   -> forall a. Maybe a
Nothing
  Just Ptr Word32
ptr  -> case (forall a. IO a -> a
Unsafe.unsafePerformIO forall a b. (a -> b) -> a -> b
$ Ptr Word32 -> IO (Word64, Int, [Word64])
marshalConwayEntry Ptr Word32
ptr) of
    (Word64
_,Int
_,[Word64
c,Word64
1])  -> forall a. a -> Maybe a
Just (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
c)
    (Word64, Int, [Word64])
_            -> forall a. HasCallStack => [Char] -> a
error [Char]
"lookupConwayPrimRoot: fatal error (should not happen)" 

--------------------------------------------------------------------------------
-- * The C table

foreign import ccall "get_conway_table_size" c_conway_table_size :: Word32
foreign import ccall "get_conway_table_ptr"  c_conway_table_ptr  :: Ptr Word32

getConwayEntryParams :: Ptr Word32 -> IO (Word64,Int)
getConwayEntryParams :: Ptr Word32 -> IO (Word64, Int)
getConwayEntryParams !Ptr Word32
ptr = do
  Word32
p <- forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
  Word32
m <- forall a. Storable a => Ptr a -> IO a
peek (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
p, forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m)

marshalConwayEntry :: Ptr Word32 -> IO (Word64,Int,[Word64])
marshalConwayEntry :: Ptr Word32 -> IO (Word64, Int, [Word64])
marshalConwayEntry !Ptr Word32
ptr = do
  Word32
p <- forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
  Word32
m <- forall a. Storable a => Ptr a -> IO a
peek (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
  [Word32]
coeffs <- forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m forall a. Num a => a -> a -> a
+ Int
1) (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
8) :: IO [Word32]
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
p , forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m , forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral [Word32]
coeffs)

--------------------------------------------------------------------------------
-- * The global table

newtype ConwayTable 
  = ConwayTable { ConwayTable -> IntMap (Ptr Word32)
fromConwayTable :: IntMap (Ptr Word32) }

encodePrimeExpo :: Int -> Int -> Int
encodePrimeExpo :: Int -> Int -> Int
encodePrimeExpo !Int
prime !Int
expo  = Int
prime forall a. Bits a => a -> a -> a
.|. (forall a. Bits a => a -> Int -> a
shiftL Int
expo Int
20)

decodePrimeExpo :: Int -> (Int,Int)
decodePrimeExpo :: Int -> (Int, Int)
decodePrimeExpo !Int
code = (Int
code forall a. Bits a => a -> a -> a
.&. Int
0xfffff , forall a. Bits a => a -> Int -> a
shiftR Int
code Int
20)

{-# NOINLINE theConwayTable #-}
theConwayTable :: ConwayTable
theConwayTable :: ConwayTable
theConwayTable = forall a. IO a -> a
Unsafe.unsafePerformIO IO ConwayTable
readConwayTableIO

{-# NOINLINE lookupConwayEntry #-}
lookupConwayEntry :: Int -> Int -> Maybe (Word64,Int,[Word64])
lookupConwayEntry :: Int -> Int -> Maybe (Word64, Int, [Word64])
lookupConwayEntry Int
p Int
m = case forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (Int -> Int -> Int
encodePrimeExpo Int
p Int
m) (ConwayTable -> IntMap (Ptr Word32)
fromConwayTable ConwayTable
theConwayTable) of
  Maybe (Ptr Word32)
Nothing  -> forall a. Maybe a
Nothing
  Just Ptr Word32
ptr -> forall a. a -> Maybe a
Just (forall a. IO a -> a
Unsafe.unsafePerformIO (Ptr Word32 -> IO (Word64, Int, [Word64])
marshalConwayEntry Ptr Word32
ptr))

--------------------------------------------------------------------------------
-- * Parse the C table into a lookup table

readConwayTableIO :: IO ConwayTable
readConwayTableIO :: IO ConwayTable
readConwayTableIO = 
  do
    [(Word32, Word32, Ptr Word32)]
list <- Word32 -> Ptr Word32 -> IO [(Word32, Word32, Ptr Word32)]
go Word32
c_conway_table_size Ptr Word32
c_conway_table_ptr 
    let f :: (a, a, b) -> (Int, b)
f (a
p,a
m,b
ptr) = (Int -> Int -> Int
encodePrimeExpo (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
p) (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
m) , b
ptr)
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ IntMap (Ptr Word32) -> ConwayTable
ConwayTable forall a b. (a -> b) -> a -> b
$ forall a. [(Int, a)] -> IntMap a
IntMap.fromList (forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a} {b}.
(Integral a, Integral a) =>
(a, a, b) -> (Int, b)
f [(Word32, Word32, Ptr Word32)]
list)
  where
    go :: Word32 -> Ptr Word32 -> IO [(Word32,Word32,Ptr Word32)]
    go :: Word32 -> Ptr Word32 -> IO [(Word32, Word32, Ptr Word32)]
go Word32
0  Ptr Word32
_    = forall (m :: * -> *) a. Monad m => a -> m a
return []
    go !Word32
k !Ptr Word32
ptr = do
      Word32
p <- forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
      Word32
m <- forall a. Storable a => Ptr a -> IO a
peek (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
      let ptr' :: Ptr Word32
ptr' = forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr (Int
8 forall a. Num a => a -> a -> a
+ Int
4forall a. Num a => a -> a -> a
*(forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m forall a. Num a => a -> a -> a
+ Int
1))
      let this :: (Word32, Word32, Ptr Word32)
this = (Word32
p,Word32
m,Ptr Word32
ptr)
      [(Word32, Word32, Ptr Word32)]
rest <- Word32 -> Ptr Word32 -> IO [(Word32, Word32, Ptr Word32)]
go (Word32
kforall a. Num a => a -> a -> a
-Word32
1) Ptr Word32
ptr'
      forall (m :: * -> *) a. Monad m => a -> m a
return ((Word32, Word32, Ptr Word32)
thisforall a. a -> [a] -> [a]
:[(Word32, Word32, Ptr Word32)]
rest)

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