{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.GaloisField where

import Data.Int
import Data.Proxy
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import GHC.Exts
import GHC.Real (Ratio (..))
import GHC.TypeLits

newtype GF (p :: Nat) = GF {forall (p :: Nat). GF p -> Int
unGF :: Int}
  deriving (GF p -> GF p -> Bool
(GF p -> GF p -> Bool) -> (GF p -> GF p -> Bool) -> Eq (GF p)
forall (p :: Nat). GF p -> GF p -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (p :: Nat). GF p -> GF p -> Bool
== :: GF p -> GF p -> Bool
$c/= :: forall (p :: Nat). GF p -> GF p -> Bool
/= :: GF p -> GF p -> Bool
Eq)
  deriving newtype (ReadPrec [GF p]
ReadPrec (GF p)
Int -> ReadS (GF p)
ReadS [GF p]
(Int -> ReadS (GF p))
-> ReadS [GF p]
-> ReadPrec (GF p)
-> ReadPrec [GF p]
-> Read (GF p)
forall (p :: Nat). ReadPrec [GF p]
forall (p :: Nat). ReadPrec (GF p)
forall (p :: Nat). Int -> ReadS (GF p)
forall (p :: Nat). ReadS [GF p]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall (p :: Nat). Int -> ReadS (GF p)
readsPrec :: Int -> ReadS (GF p)
$creadList :: forall (p :: Nat). ReadS [GF p]
readList :: ReadS [GF p]
$creadPrec :: forall (p :: Nat). ReadPrec (GF p)
readPrec :: ReadPrec (GF p)
$creadListPrec :: forall (p :: Nat). ReadPrec [GF p]
readListPrec :: ReadPrec [GF p]
Read, Int -> GF p -> ShowS
[GF p] -> ShowS
GF p -> String
(Int -> GF p -> ShowS)
-> (GF p -> String) -> ([GF p] -> ShowS) -> Show (GF p)
forall (p :: Nat). Int -> GF p -> ShowS
forall (p :: Nat). [GF p] -> ShowS
forall (p :: Nat). GF p -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (p :: Nat). Int -> GF p -> ShowS
showsPrec :: Int -> GF p -> ShowS
$cshow :: forall (p :: Nat). GF p -> String
show :: GF p -> String
$cshowList :: forall (p :: Nat). [GF p] -> ShowS
showList :: [GF p] -> ShowS
Show)

pattern GF# :: Int# -> GF p
pattern $mGF# :: forall {r} {p :: Nat}. GF p -> (Int# -> r) -> ((# #) -> r) -> r
$bGF# :: forall (p :: Nat). Int# -> GF p
GF# x# = GF (I# x#)
{-# COMPLETE GF# #-}

mkGF :: forall p. (KnownNat p) => Int -> GF p
mkGF :: forall (p :: Nat). KnownNat p => Int -> GF p
mkGF Int
x = Int -> GF p
forall (p :: Nat). Int -> GF p
GF (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)))

validateGF :: forall p. (KnownNat p) => GF p -> Bool
validateGF :: forall (p :: Nat). KnownNat p => GF p -> Bool
validateGF (GF Int
x) = Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
x Bool -> Bool -> Bool
&& Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Proxy p -> Int
forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Int
natValAsInt (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @p)

natValAsInt :: (KnownNat n) => proxy n -> Int
natValAsInt :: forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Int
natValAsInt = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> (proxy n -> Integer) -> proxy n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal
{-# INLINE natValAsInt #-}

natValAsWord :: (KnownNat n) => proxy n -> Word
natValAsWord :: forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Word
natValAsWord = Integer -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Word) -> (proxy n -> Integer) -> proxy n -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal
{-# INLINE natValAsWord #-}

{- |
>>> reifyNat (998244353 :: Int) natValAsInt
998244353
-}
reifyNat :: (Integral i) => i -> (forall n. (KnownNat n) => Proxy n -> a) -> a
reifyNat :: forall i a.
Integral i =>
i -> (forall (n :: Nat). KnownNat n => Proxy n -> a) -> a
reifyNat i
n forall (n :: Nat). KnownNat n => Proxy n -> a
f = case Integer -> Maybe SomeNat
someNatVal (i -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
n) of
  Just (SomeNat Proxy n
proxy) -> Proxy n -> a
forall (n :: Nat). KnownNat n => Proxy n -> a
f Proxy n
proxy
  Maybe SomeNat
Nothing -> String -> a
forall a. HasCallStack => String -> a
error String
"reifyNat failed"
{-# INLINE reifyNat #-}

asGFOf :: GF p -> Proxy p -> GF p
asGFOf :: forall (p :: Nat). GF p -> Proxy p -> GF p
asGFOf = GF p -> Proxy p -> GF p
forall a b. a -> b -> a
const
{-# INLINE asGFOf #-}

instance (KnownNat p) => Bounded (GF p) where
  minBound :: GF p
minBound = Int -> GF p
forall (p :: Nat). Int -> GF p
GF Int
0
  maxBound :: GF p
maxBound = Int -> GF p
forall (p :: Nat). Int -> GF p
GF (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

instance (KnownNat p) => Num (GF p) where
  {-# SPECIALIZE instance Num (GF 998244353) #-}
  (GF# Int#
x#) + :: GF p -> GF p -> GF p
+ (GF# Int#
y#) = case Int#
x# Int# -> Int# -> Int#
+# Int#
y# of
    Int#
xy# -> Int# -> GF p
forall (p :: Nat). Int# -> GF p
GF# (Int#
xy# Int# -> Int# -> Int#
-# ((Int#
xy# Int# -> Int# -> Int#
>=# Int#
m#) Int# -> Int# -> Int#
*# Int#
m#))
    where
      !(I# Int#
m#) = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)
  {-# INLINE (+) #-}
  (GF# Int#
x#) - :: GF p -> GF p -> GF p
- (GF# Int#
y#) = case Int#
x# Int# -> Int# -> Int#
-# Int#
y# of
    Int#
xy# -> Int# -> GF p
forall (p :: Nat). Int# -> GF p
GF# (Int#
xy# Int# -> Int# -> Int#
+# ((Int#
xy# Int# -> Int# -> Int#
<# Int#
0#) Int# -> Int# -> Int#
*# Int#
m#))
    where
      !(I# Int#
m#) = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)
  {-# INLINE (-) #-}
  (GF# Int#
x#) * :: GF p -> GF p -> GF p
* (GF# Int#
y#) = case Word# -> Word# -> Word#
timesWord# (Int# -> Word#
int2Word# Int#
x#) (Int# -> Word#
int2Word# Int#
y#) of
    Word#
z# -> case Word# -> Word# -> (# Word#, Word# #)
timesWord2# Word#
z# Word#
im# of
      (# Word#
q#, Word#
_ #) -> case Word# -> Word# -> Word#
minusWord# Word#
z# (Word# -> Word# -> Word#
timesWord# Word#
q# Word#
m#) of
        Word#
v#
          | Int# -> Bool
isTrue# (Word# -> Word# -> Int#
geWord# Word#
v# Word#
m#) -> Int# -> GF p
forall (p :: Nat). Int# -> GF p
GF# (Word# -> Int#
word2Int# (Word# -> Word# -> Word#
plusWord# Word#
v# Word#
m#))
          | Bool
otherwise -> Int# -> GF p
forall (p :: Nat). Int# -> GF p
GF# (Word# -> Int#
word2Int# Word#
v#)
    where
      !(W# Word#
m#) = Integer -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Word) -> Integer -> Word
forall a b. (a -> b) -> a -> b
$ Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)
      im# :: Word#
im# = Word# -> Word# -> Word#
plusWord# (Word# -> Word# -> Word#
quotWord# Word#
0xffffffffffffffff## Word#
m#) Word#
1##
  {-# INLINE (*) #-}
  abs :: GF p -> GF p
abs = GF p -> GF p
forall a. a -> a
id
  signum :: GF p -> GF p
signum = GF p -> GF p -> GF p
forall a b. a -> b -> a
const (Int -> GF p
forall (p :: Nat). Int -> GF p
GF Int
1)
  fromInteger :: Integer -> GF p
fromInteger Integer
x = Int -> GF p
forall (p :: Nat). Int -> GF p
GF (Int -> GF p) -> (Integer -> Int) -> Integer -> GF p
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> GF p) -> Integer -> GF p
forall a b. (a -> b) -> a -> b
$ Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
    where
      m :: Integer
m = Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p)

{- |
>>> recip @(GF 998244353) 0
0
>>> recip @(GF 2) 0
0
-}
instance (KnownNat p) => Fractional (GF p) where
  {-# SPECIALIZE instance Fractional (GF 998244353) #-}
  recip :: GF p -> GF p
recip = case Proxy# p -> Integer
forall (n :: Nat). KnownNat n => Proxy# n -> Integer
natVal' (forall (a :: Nat). Proxy# a
forall {k} (a :: k). Proxy# a
proxy# @p) of
    Integer
2 -> GF p -> GF p
forall a. a -> a
id
    Integer
p -> (GF p -> Int -> GF p
forall a b. (Num a, Integral b) => a -> b -> a
^ (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 :: Int))
  {-# INLINE recip #-}
  fromRational :: Rational -> GF p
fromRational (Integer
p :% Integer
q) = Integer -> GF p
forall a. Num a => Integer -> a
fromInteger Integer
p GF p -> GF p -> GF p
forall a. Fractional a => a -> a -> a
/ Integer -> GF p
forall a. Num a => Integer -> a
fromInteger Integer
q

instance U.IsoUnbox (GF p) Int32 where
  toURepr :: GF p -> Int32
toURepr = (Int -> Int32) -> GF p -> Int32
forall a b. Coercible a b => a -> b
coerce (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Int32)
  {-# INLINE toURepr #-}
  fromURepr :: Int32 -> GF p
fromURepr = (Int32 -> Int) -> Int32 -> GF p
forall a b. Coercible a b => a -> b
coerce (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int32 @Int)
  {-# INLINE fromURepr #-}

newtype instance UM.MVector s (GF p) = MV_GF (UM.MVector s Int32)
newtype instance U.Vector (GF p) = V_GF (U.Vector Int32)
deriving via (GF p `U.As` Int32) instance GM.MVector U.MVector (GF p)
deriving via (GF p `U.As` Int32) instance G.Vector U.Vector (GF p)
instance U.Unbox (GF p)