-- | 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 :: ConwayPoly p m -> Ptr Word32
fromConwayPoly (ConwayWitness Ptr Word32
ptr) = Ptr Word32
ptr

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

conwayPrime_ :: ConwayPoly p m -> Word64
conwayPrime_ :: ConwayPoly p m -> Word64
conwayPrime_ (ConwayWitness Ptr Word32
ptr) = IO Word64 -> Word64
forall a. IO a -> a
Unsafe.unsafePerformIO (IO Word64 -> Word64) -> IO Word64 -> Word64
forall a b. (a -> b) -> a -> b
$ do
  (Word64
p,Int
_) <- Ptr Word32 -> IO (Word64, Int)
getConwayEntryParams Ptr Word32
ptr
  Word64 -> IO Word64
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 Int -> IntMap (Ptr Word32) -> Maybe (Ptr Word32)
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   -> Maybe Word64
forall a. Maybe a
Nothing
  Just Ptr Word32
ptr  -> case (IO (Word64, Int, [Word64]) -> (Word64, Int, [Word64])
forall a. IO a -> a
Unsafe.unsafePerformIO (IO (Word64, Int, [Word64]) -> (Word64, Int, [Word64]))
-> IO (Word64, Int, [Word64]) -> (Word64, Int, [Word64])
forall a b. (a -> b) -> a -> b
$ Ptr Word32 -> IO (Word64, Int, [Word64])
marshalConwayEntry Ptr Word32
ptr) of
    (Word64
_,Int
_,[Word64
c,Word64
1])  -> Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
c)
    (Word64, Int, [Word64])
_            -> [Char] -> Maybe 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 <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
  Word32
m <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (Ptr Word32 -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
  (Word64, Int) -> IO (Word64, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
p, Word32 -> Int
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 <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
  Word32
m <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (Ptr Word32 -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
  [Word32]
coeffs <- Int -> Ptr Word32 -> IO [Word32]
forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Ptr Word32 -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
8) :: IO [Word32]
  (Word64, Int, [Word64]) -> IO (Word64, Int, [Word64])
forall (m :: * -> *) a. Monad m => a -> m a
return (Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
p , Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m , (Word32 -> Word64) -> [Word32] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map Word32 -> Word64
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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (Int -> Int -> Int
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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0xfffff , Int -> Int -> Int
forall a. Bits a => a -> Int -> a
shiftR Int
code Int
20)

{-# NOINLINE theConwayTable #-}
theConwayTable :: ConwayTable
theConwayTable :: ConwayTable
theConwayTable = IO ConwayTable -> ConwayTable
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 Int -> IntMap (Ptr Word32) -> Maybe (Ptr Word32)
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  -> Maybe (Word64, Int, [Word64])
forall a. Maybe a
Nothing
  Just Ptr Word32
ptr -> (Word64, Int, [Word64]) -> Maybe (Word64, Int, [Word64])
forall a. a -> Maybe a
Just (IO (Word64, Int, [Word64]) -> (Word64, Int, [Word64])
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 (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
p) (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
m) , b
ptr)
    ConwayTable -> IO ConwayTable
forall (m :: * -> *) a. Monad m => a -> m a
return (ConwayTable -> IO ConwayTable) -> ConwayTable -> IO ConwayTable
forall a b. (a -> b) -> a -> b
$ IntMap (Ptr Word32) -> ConwayTable
ConwayTable (IntMap (Ptr Word32) -> ConwayTable)
-> IntMap (Ptr Word32) -> ConwayTable
forall a b. (a -> b) -> a -> b
$ [(Int, Ptr Word32)] -> IntMap (Ptr Word32)
forall a. [(Int, a)] -> IntMap a
IntMap.fromList (((Word32, Word32, Ptr Word32) -> (Int, Ptr Word32))
-> [(Word32, Word32, Ptr Word32)] -> [(Int, Ptr Word32)]
forall a b. (a -> b) -> [a] -> [b]
map (Word32, Word32, Ptr Word32) -> (Int, Ptr Word32)
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
_    = [(Word32, Word32, Ptr Word32)] -> IO [(Word32, Word32, Ptr Word32)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    go !Word32
k !Ptr Word32
ptr = do
      Word32
p <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek Ptr Word32
ptr             :: IO Word32
      Word32
m <- Ptr Word32 -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (Ptr Word32 -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr Int
4) :: IO Word32
      let ptr' :: Ptr Word32
ptr' = Ptr Word32 -> Int -> Ptr Word32
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word32
ptr (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
*(Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m Int -> Int -> Int
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
kWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
-Word32
1) Ptr Word32
ptr'
      [(Word32, Word32, Ptr Word32)] -> IO [(Word32, Word32, Ptr Word32)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Word32, Word32, Ptr Word32)
this(Word32, Word32, Ptr Word32)
-> [(Word32, Word32, Ptr Word32)] -> [(Word32, Word32, Ptr Word32)]
forall a. a -> [a] -> [a]
:[(Word32, Word32, Ptr Word32)]
rest)

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