-- | Low level stuff

{-# LANGUAGE BangPatterns, DataKinds, KindSignatures, GADTs #-}
{-# LANGUAGE ScopedTypeVariables, TypeApplications #-}

module Math.FiniteField.GaloisField.Small.Internal where

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

import Data.Bits
import Data.Int
import Data.Word
import Data.List

import GHC.TypeNats (Nat)

import Control.Monad
import Control.Monad.ST

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

import Data.Vector.Unboxed (Vector, MVector)
import qualified Data.Vector.Unboxed         as Vec
import qualified Data.Vector.Generic.Mutable as MVec

import System.IO.Unsafe as Unsafe
import Control.Monad.ST.Unsafe (unsafeIOToST)      -- debugging only

import Math.FiniteField.Primes
import Math.FiniteField.TypeLevel
import Math.FiniteField.Conway
import Math.FiniteField.Conway.Internal
import Math.FiniteField.Misc
 
import qualified Math.FiniteField.PrimeField.Small.Raw as Raw

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

type P = Word64
type M = Int
type C = Ptr Word32
type F = Vector Word32

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

neg :: P -> F -> F
neg :: P -> F -> F
neg !P
p !F
xs = (Word32 -> Word32) -> F -> F
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
Vec.map Word32 -> Word32
neg1 F
xs where
  neg1 :: Word32 -> Word32
neg1 !Word32
x = P -> Word32
w64toW32 (P -> P -> P
Raw.neg P
p (Word32 -> P
w32toW64 Word32
x))

add :: P -> F -> F -> F
add :: P -> F -> F -> F
add !P
p !F
xs !F
ys = (Word32 -> Word32 -> Word32) -> F -> F -> F
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
Vec.zipWith Word32 -> Word32 -> Word32
add1 F
xs F
ys where 
  add1 :: Word32 -> Word32 -> Word32
add1 !Word32
x !Word32
y = P -> Word32
w64toW32 (P -> P -> P -> P
Raw.add P
p (Word32 -> P
w32toW64 Word32
x) (Word32 -> P
w32toW64 Word32
y))

sub :: P -> F -> F -> F
sub :: P -> F -> F -> F
sub !P
p !F
xs !F
ys = (Word32 -> Word32 -> Word32) -> F -> F -> F
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
Vec.zipWith Word32 -> Word32 -> Word32
sub1 F
xs F
ys where
  sub1 :: Word32 -> Word32 -> Word32
sub1 !Word32
x !Word32
y = P -> Word32
w64toW32 (P -> P -> P -> P
Raw.sub P
p (Word32 -> P
w32toW64 Word32
x) (Word32 -> P
w32toW64 Word32
y))

mul :: C -> F -> F -> F
mul :: C -> F -> F -> F
mul !C
ptr !F
xs !F
ys = IO F -> F
forall a. IO a -> a
Unsafe.unsafePerformIO (C -> F -> F -> IO F
mulIO C
ptr F
xs F
ys)

inv :: P -> C -> F -> F
inv :: P -> C -> F -> F
inv !P
p !C
ptr !F
xs = P -> C -> F -> F
invEuclid P
p C
ptr F
xs

div :: P -> C -> F -> F -> F
div :: P -> C -> F -> F -> F
div !P
p !C
ptr !F
xs !F
ys = C -> F -> F -> F
mul C
ptr F
xs (P -> C -> F -> F
inv P
p C
ptr F
ys)

--------------------------------------------------------------------------------  
-- * multiplication

mulIO :: C -> F -> F -> IO F
mulIO :: C -> F -> F -> IO F
mulIO !C
ptr !F
xs !F
ys = 
  do
    !Word32
p0 <- C -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek C
ptr             :: IO Word32
    !Word32
m0 <- C -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (C -> Int -> C
forall a b. Ptr a -> Int -> Ptr b
plusPtr C
ptr Int
4) :: IO Word32
    let !p :: P
p  = Word32 -> P
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
p0   :: Word64
        !m :: Int
m  = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m0   :: Int
        !m1 :: Int
m1 = Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1             :: Int
    F
cvec <- Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ([Word32] -> F) -> IO [Word32] -> IO F
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> C -> IO [Word32]
forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (C -> Int -> C
forall a b. Ptr a -> Int -> Ptr b
plusPtr C
ptr Int
8)
    MVector RealWorld Word32
mvec <- F -> IO (MVector (PrimState IO) Word32)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
Vec.unsafeThaw (P -> F -> F -> F
mulPoly P
p F
xs F
ys)
    [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Int]
downIndices Int
m1) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k -> do
      Word32
c0 <- MVector (PrimState IO) Word32 -> Int -> IO Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MVec.unsafeRead MVector RealWorld Word32
MVector (PrimState IO) Word32
mvec (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
      -- print (k,c0)
      let !c :: P
c = Word32 -> P
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
c0 :: Word64
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (P
c P -> P -> Bool
forall a. Eq a => a -> a -> Bool
/= P
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
m] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
j -> do     -- NOTE: this could be [0..m1], a bit faster, but while debugging i leave it as is
          let !y :: P
y = Word32 -> P
w32toW64 (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
cvec Int
j)
          MVector (PrimState IO) Word32 -> (Word32 -> Word32) -> Int -> IO ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MVec.unsafeModify MVector RealWorld Word32
MVector (PrimState IO) Word32
mvec (\ !Word32
z -> P -> Word32
w64toW32 (P -> P -> P -> P
Raw.sub P
p (Word32 -> P
w32toW64 Word32
z) (P -> P -> P -> P
Raw.mul P
p P
c P
y))) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
    -- print =<< Vec.freeze mvec
    [Word32]
xs <- [Int] -> (Int -> IO Word32) -> IO [Word32]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0..Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO Word32) -> IO [Word32])
-> (Int -> IO Word32) -> IO [Word32]
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> MVector (PrimState IO) Word32 -> Int -> IO Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MVec.unsafeRead MVector RealWorld Word32
MVector (PrimState IO) Word32
mvec Int
i
    F -> IO F
forall (m :: * -> *) a. Monad m => a -> m a
return (F -> IO F) -> F -> IO F
forall a b. (a -> b) -> a -> b
$ Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN Int
m [Word32]
xs

--------------------------------------------------------------------------------  e
-- * polynomials

-- -- | Multiplication of polynomials. Note: we assume both have length @m@
-- mulVecPolyFix  :: P -> M -> Vector Word32 -> Vector Word32 -> Vector Word32
-- mulVecPolyFIx !p !m !xs !ys = 
--   Vec.create $ do
--     mvec <- MVec.replicate mbig 0
--     forM_ [0..m-1] $ \ !i -> do
--       let !x = w32toW64 (Vec.unsafeIndex xs i)
--       forM_ [0..m-1] $ \ !j -> do
--         let !y = w32toW64 (Vec.unsafeIndex ys j)
--         MVec.unsafeModify mvec (\ !z -> w64toW32 (Raw.add p (w32toW64 z) (Raw.mul p x y))) (i+j)
--     return mvec
--   where
--     !mbig = m+m-1

resizePoly :: Int -> Vector Word32 -> Vector Word32
resizePoly :: Int -> F -> F
resizePoly Int
m F
vec = (forall s. ST s (MVector s Word32)) -> F
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
Vec.create ((forall s. ST s (MVector s Word32)) -> F)
-> (forall s. ST s (MVector s Word32)) -> F
forall a b. (a -> b) -> a -> b
$ do
  let !l :: Int
l = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
m (F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
vec) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  MVector s Word32
mvec <- Int -> Word32 -> ST s (MVector (PrimState (ST s)) Word32)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MVec.replicate Int
m Word32
0
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
l] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> MVector (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MVec.unsafeWrite MVector s Word32
MVector (PrimState (ST s)) Word32
mvec Int
i (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
vec Int
i)
  MVector s Word32 -> ST s (MVector s Word32)
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Word32
mvec

-- | The vectors can have different lengths
addPoly :: P -> Vector Word32 -> Vector Word32 -> Vector Word32
addPoly :: P -> F -> F -> F
addPoly !P
p !F
xs !F
ys 
  | F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
ys = (Word32 -> Word32 -> Word32) -> F -> F -> F
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
Vec.zipWith Word32 -> Word32 -> Word32
add1 F
xs F
ys 
  | Bool
otherwise = [Word32] -> F
forall a. Unbox a => [a] -> Vector a
Vec.fromList (Word32
-> Word32
-> (Word32 -> Word32 -> Word32)
-> [Word32]
-> [Word32]
-> [Word32]
forall a b c. a -> b -> (a -> b -> c) -> [a] -> [b] -> [c]
longZipWith Word32
0 Word32
0 Word32 -> Word32 -> Word32
add1 (F -> [Word32]
forall a. Unbox a => Vector a -> [a]
Vec.toList F
xs) (F -> [Word32]
forall a. Unbox a => Vector a -> [a]
Vec.toList F
ys))
  where
    add1 :: Word32 -> Word32 -> Word32
add1 !Word32
x !Word32
y = P -> Word32
w64toW32 (P -> P -> P -> P
Raw.add P
p (Word32 -> P
w32toW64 Word32
x) (Word32 -> P
w32toW64 Word32
y))

-- | The vectors can have different lengths
subPoly :: P -> Vector Word32 -> Vector Word32 -> Vector Word32
subPoly :: P -> F -> F -> F
subPoly !P
p !F
xs !F
ys 
  | F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
ys = (Word32 -> Word32 -> Word32) -> F -> F -> F
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
Vec.zipWith Word32 -> Word32 -> Word32
sub1 F
xs F
ys 
  | Bool
otherwise = [Word32] -> F
forall a. Unbox a => [a] -> Vector a
Vec.fromList (Word32
-> Word32
-> (Word32 -> Word32 -> Word32)
-> [Word32]
-> [Word32]
-> [Word32]
forall a b c. a -> b -> (a -> b -> c) -> [a] -> [b] -> [c]
longZipWith Word32
0 Word32
0 Word32 -> Word32 -> Word32
sub1 (F -> [Word32]
forall a. Unbox a => Vector a -> [a]
Vec.toList F
xs) (F -> [Word32]
forall a. Unbox a => Vector a -> [a]
Vec.toList F
ys))
  where
    sub1 :: Word32 -> Word32 -> Word32
sub1 !Word32
x !Word32
y = P -> Word32
w64toW32 (P -> P -> P -> P
Raw.sub P
p (Word32 -> P
w32toW64 Word32
x) (Word32 -> P
w32toW64 Word32
y))

scalePoly :: P -> Word32 -> Vector Word32 -> Vector Word32
scalePoly :: P -> Word32 -> F -> F
scalePoly !P
p !Word32
s !F
ys = (Word32 -> Word32) -> F -> F
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
Vec.map Word32 -> Word32
mul1 F
ys where
  mul1 :: Word32 -> Word32
mul1 !Word32
y = P -> Word32
w64toW32 (P -> P -> P -> P
Raw.mul P
p P
s64 (Word32 -> P
w32toW64 Word32
y))
  s64 :: P
s64 = (Word32 -> P
w32toW64 Word32
s)

mulPoly  :: P -> Vector Word32 -> Vector Word32 -> Vector Word32
mulPoly :: P -> F -> F -> F
mulPoly !P
p !F
xs !F
ys =
  (forall s. ST s (MVector s Word32)) -> F
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
Vec.create ((forall s. ST s (MVector s Word32)) -> F)
-> (forall s. ST s (MVector s Word32)) -> F
forall a b. (a -> b) -> a -> b
$ do
    MVector s Word32
mvec <- Int -> Word32 -> ST s (MVector (PrimState (ST s)) Word32)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MVec.replicate Int
mbig Word32
0
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
m1Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> do
      let !x :: P
x = Word32 -> P
w32toW64 (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
xs Int
i)
      [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
m2Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
j -> do
        let !y :: P
y = Word32 -> P
w32toW64 (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
ys Int
j)
        MVector (PrimState (ST s)) Word32
-> (Word32 -> Word32) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MVec.unsafeModify MVector s Word32
MVector (PrimState (ST s)) Word32
mvec (\ !Word32
z -> P -> Word32
w64toW32 (P -> P -> P -> P
Raw.add P
p (Word32 -> P
w32toW64 Word32
z) (P -> P -> P -> P
Raw.mul P
p P
x P
y))) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
j)
    MVector s Word32 -> ST s (MVector s Word32)
forall (m :: * -> *) a. Monad m => a -> m a
return MVector s Word32
mvec
  where
    !m1 :: Int
m1 = F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
xs
    !m2 :: Int
m2 = F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
ys
    !mbig :: Int
mbig = Int
m1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
m2Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1

zeroPoly :: Int -> Vector Word32 
zeroPoly :: Int -> F
zeroPoly Int
n = Int -> Word32 -> F
forall a. Unbox a => Int -> a -> Vector a
Vec.replicate Int
n Word32
0

onePoly :: Int -> Vector Word32
onePoly :: Int -> F
onePoly Int
n = Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN Int
n (Word32
1 Word32 -> [Word32] -> [Word32]
forall a. a -> [a] -> [a]
: Int -> Word32 -> [Word32]
forall a. Int -> a -> [a]
replicate (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Word32
0)

isZeroPoly :: Vector Word32 -> Bool
isZeroPoly :: F -> Bool
isZeroPoly F
v = (Word32 -> Bool) -> [Word32] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
==Word32
0) (F -> [Word32]
forall a. Unbox a => Vector a -> [a]
Vec.toList F
v)

polyDegree :: Vector Word32 -> Int
polyDegree :: F -> Int
polyDegree F
v = Int -> Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) where
  !n :: Int
n = F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
v
  go :: Int -> Int
go !Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 
    then (-Int
1) 
    else if F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
v Int
i Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word32
0 
      then Int
i
      else Int -> Int
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

polyLongDivision :: P -> Vector Word32 -> Vector Word32 -> (Vector Word32, Vector Word32)
polyLongDivision :: P -> F -> F -> (F, F)
polyLongDivision !P
p !F
as !F
bs = (forall s. ST s (F, F)) -> (F, F)
forall a. (forall s. ST s a) -> a
runST forall s. ST s (F, F)
action where
  action :: forall s. ST s (Vector Word32, Vector Word32) 
  action :: ST s (F, F)
action = do
    let !m :: Int
m = F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
as
    let !d :: Int
d = F -> Int
polyDegree F
bs
        !lcf0 :: Word32
lcf0 = F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
bs Int
d   -- leading coefficient
        !lcf :: P
lcf  = Word32 -> P
w32toW64 Word32
lcf0
    MVector s Word32
rem  <- F -> ST s (MVector (PrimState (ST s)) Word32)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
Vec.thaw F
as
    MVector s Word32
quot <- Int -> Word32 -> ST s (MVector (PrimState (ST s)) Word32)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
MVec.replicate @(ST s) @MVector Int
m (Word32
0 :: Word32) 
   
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Int]
downIndices (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
d)) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
k -> do
      Word32
c0 <- MVector (PrimState (ST s)) Word32 -> Int -> ST s Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MVec.unsafeRead MVector s Word32
MVector (PrimState (ST s)) Word32
rem (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
      -- print (k,c0)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word32
c0 Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word32
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        let !s :: P
s = P -> P -> P -> P
Raw.div P
p (Word32 -> P
w32toW64 Word32
c0) P
lcf
        -- Vec.freeze rem >>= \r -> unsafeIOToST $ print (k,c0,s,r)
        MVector (PrimState (ST s)) Word32 -> Int -> Word32 -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MVec.unsafeWrite MVector s Word32
MVector (PrimState (ST s)) Word32
quot (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (P -> Word32
w64toW32 P
s)
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
d] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
j -> do     -- NOTE: this could be [0..d-1], a bit faster, but while debugging i leave it as is
          let !y :: P
y = Word32 -> P
w32toW64 (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
bs Int
j)
          MVector (PrimState (ST s)) Word32
-> (Word32 -> Word32) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
MVec.unsafeModify MVector s Word32
MVector (PrimState (ST s)) Word32
rem (\ !Word32
z -> P -> Word32
w64toW32 (P -> P -> P -> P
Raw.sub P
p (Word32 -> P
w32toW64 Word32
z) (P -> P -> P -> P
Raw.mul P
p P
s P
y))) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  
    [Word32]
fquot <- [Int] -> (Int -> ST s Word32) -> ST s [Word32]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0..Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s Word32) -> ST s [Word32])
-> (Int -> ST s Word32) -> ST s [Word32]
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> MVector (PrimState (ST s)) Word32 -> Int -> ST s Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MVec.unsafeRead MVector s Word32
MVector (PrimState (ST s)) Word32
quot Int
i
    [Word32]
frem  <- [Int] -> (Int -> ST s Word32) -> ST s [Word32]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0..Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s Word32) -> ST s [Word32])
-> (Int -> ST s Word32) -> ST s [Word32]
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> MVector (PrimState (ST s)) Word32 -> Int -> ST s Word32
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MVec.unsafeRead MVector s Word32
MVector (PrimState (ST s)) Word32
rem  Int
i
    (F, F) -> ST s (F, F)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN Int
m [Word32]
fquot, Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN Int
d [Word32]
frem)

downIndices :: Int -> [Int]
downIndices :: Int -> [Int]
downIndices !Int
k = if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 then [] else Int
k Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> [Int]
downIndices (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

polyQuotient :: P -> Vector Word32 -> Vector Word32 -> Vector Word32
polyQuotient :: P -> F -> F -> F
polyQuotient P
p F
xs F
ys = (F, F) -> F
forall a b. (a, b) -> a
fst (P -> F -> F -> (F, F)
polyLongDivision P
p F
xs F
ys)

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

getCPoly :: Ptr Word32 -> Vector Word32
getCPoly :: C -> F
getCPoly C
ptr = IO F -> F
forall a. IO a -> a
Unsafe.unsafePerformIO (C -> IO F
getCPolyIO C
ptr)

getCPolyIO :: Ptr Word32 -> IO (Vector Word32)
getCPolyIO :: C -> IO F
getCPolyIO C
ptr = do
  !Word32
m0 <- C -> IO Word32
forall a. Storable a => Ptr a -> IO a
peek (C -> Int -> C
forall a b. Ptr a -> Int -> Ptr b
plusPtr C
ptr Int
4) :: IO Word32
  let !m :: Int
m  = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
m0   :: Int
  Int -> [Word32] -> F
forall a. Unbox a => Int -> [a] -> Vector a
Vec.fromListN (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ([Word32] -> F) -> IO [Word32] -> IO F
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> C -> IO [Word32]
forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (C -> Int -> C
forall a b. Ptr a -> Int -> Ptr b
plusPtr C
ptr Int
8)

-- | Based on <https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm#Modular_integers>
invEuclid :: P -> C -> Vector Word32 -> Vector Word32
invEuclid :: P -> C -> F -> F
invEuclid !P
p !C
ptr !F
a = F -> F -> F -> F -> F
worker (Int -> F
zeroPoly Int
n) (Int -> F
onePoly Int
n) (C -> F
getCPoly C
ptr) F
a where

  n :: Int
n = F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
a

  worker :: Vector Word32 -> Vector Word32 -> Vector Word32 -> Vector Word32 -> Vector Word32
  worker :: F -> F -> F -> F -> F
worker !F
t !F
newt !F
r !F
newr = case F -> Bool
isZeroPoly F
newr of

    Bool
False -> let quot :: F
quot = P -> F -> F -> F
polyQuotient P
p F
r F
newr 
                 r' :: F
r'   = P -> F -> F -> F
subPoly P
p F
r (P -> F -> F -> F
mulPoly P
p F
quot F
newr)
                 t' :: F
t'   = P -> F -> F -> F
subPoly P
p F
t (P -> F -> F -> F
mulPoly P
p F
quot F
newt)
             in  F -> F -> F -> F -> F
worker F
newt F
t' F
newr F
r'

    Bool
True  -> if (F -> Int
polyDegree F
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
      then Int -> F
zeroPoly (F -> Int
forall a. Unbox a => Vector a -> Int
Vec.length F
a)
      else let r0 :: P
r0 = Word32 -> P
w32toW64 (F -> Int -> Word32
forall a. Unbox a => Vector a -> Int -> a
Vec.unsafeIndex F
r Int
0)
           in  Int -> F -> F
resizePoly Int
n (P -> Word32 -> F -> F
scalePoly P
p (P -> Word32
w64toW32 (P -> P -> P
Raw.inv P
p P
r0)) F
t)

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