-- | Type level naturals, singletons, prime witnesses and stuff.
--
-- This would be much simpler in a dependently typed language;
-- alas, we are in Haskell, so for now we have to feel the pain.
--
-- How to use this, the over-complicated way:
--
-- * create a 'SomeSNat' from an integer using the function 'someSNat'
-- 
-- * pattern match on it with a @case@ expression. Within the matched
--   scope, the type parameter is \"instantiated\". So all your program
--   will be \"inside\" this case branch (of course you can call out
--   to functions)
--
-- * create witnesses for this being prime and\/or being small using
--   the functions 'isPrime' and 'fits31Bits' 
--
-- * if you want small primes, create a \"small prime witness\" using
--   'mkSmallPrime'
--
-- * now you are ready to use the resulting witness (of type @IsPrime n@
--   or @IsSmallPrime n@) to create finite fields.
--
-- Or you can just the functions provided with each field implementation
-- to create the fields directly (in form of existantials, ie. sigma types,
-- so you still have to do the @case _ of@ thing).
--

{-# LANGUAGE DataKinds, KindSignatures, PolyKinds, ScopedTypeVariables #-}
{-# LANGUAGE GADTs, ExistentialQuantification, StandaloneDeriving #-}

module Math.FiniteField.TypeLevel
  ( -- * Singleton types
    SNat   , fromSNat   , proxyOfSNat   , proxyToSNat   , unsafeSNat
  , SNat64 , fromSNat64 , proxyOfSNat64 , proxyToSNat64 , unsafeSNat64
    -- * Creating singleton types
  , SomeSNat(..) , someSNat , SomeSNat64(..) , someSNat64 , someSNat64_
    -- * Small numbers
  , Fits31Bits , from31Bit , from31BitSigned , from31Bit' , fits31Bits 
    -- * Primes
  , IsPrime , fromPrime , fromPrime' 
  , isPrime , believeMeItsPrime
    -- * Small primes
  , IsSmallPrime , fromSmallPrime , fromSmallPrimeSigned , fromSmallPrimeInteger , fromSmallPrime' 
  , isSmallPrime , believeMeItsASmallPrime
  , smallPrimeIsPrime , smallPrimeIsSmall , mkSmallPrime
    -- * Divisors
  , Divides , _dividend , _divisor , _quotient 
  , dividendSNat , divisorSNat 
  , divides
  , Divisor(..) , constructDivisor , divisors
    -- * Proxy
  , proxyOf, proxyOf1
    -- * Sanity checking
  , checkSomeSNat , checkSomeSNat64
  ) 
  where

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

import Data.Int
import Data.Word
import Data.Proxy
import Data.List

import GHC.TypeNats
import Data.Proxy

import qualified Math.FiniteField.Primes as Primes
import Math.FiniteField.TypeLevel.Singleton

--------------------------------------------------------------------------------
-- * Primes

-- | Prime witness
newtype IsPrime (n :: Nat) where
  PrimeWitness :: SNat n -> IsPrime n

deriving instance Show (IsPrime n)

fromPrime' :: IsPrime n -> SNat n
fromPrime' :: IsPrime n -> SNat n
fromPrime' (PrimeWitness SNat n
sn) = SNat n
sn

fromPrime :: IsPrime n -> Integer
fromPrime :: IsPrime n -> Integer
fromPrime (PrimeWitness SNat n
sn) = SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
sn

-- | Prime testing.
--
-- Note: this uses trial division at the moment, 
-- so it's only good for small numbers for now
--
isPrime :: SNat n -> Maybe (IsPrime n)
isPrime :: SNat n -> Maybe (IsPrime n)
isPrime SNat n
sn = if (SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
sn Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1) Bool -> Bool -> Bool
&& Integer -> Bool
Primes.isPrimeTrialDivision (SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
sn)
  then IsPrime n -> Maybe (IsPrime n)
forall a. a -> Maybe a
Just (SNat n -> IsPrime n
forall (n :: Nat). SNat n -> IsPrime n
PrimeWitness SNat n
sn)
  else Maybe (IsPrime n)
forall a. Maybe a
Nothing

-- | Escape hatch
believeMeItsPrime :: SNat n -> IsPrime n
believeMeItsPrime :: SNat n -> IsPrime n
believeMeItsPrime SNat n
sn = SNat n -> IsPrime n
forall (n :: Nat). SNat n -> IsPrime n
PrimeWitness SNat n
sn

--------------------------------------------------------------------------------
-- * Small numbers

newtype Fits31Bits (n :: Nat) where
  Witness31 :: SNat64 n -> Fits31Bits n

deriving instance Show (Fits31Bits n)

from31Bit' :: Fits31Bits n -> SNat64 n
from31Bit' :: Fits31Bits n -> SNat64 n
from31Bit' (Witness31 SNat64 n
sn) = SNat64 n
sn

from31Bit :: Fits31Bits n -> Word64
from31Bit :: Fits31Bits n -> Word64
from31Bit (Witness31 SNat64 n
sn) = SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn

from31BitSigned :: Fits31Bits n -> Int64
from31BitSigned :: Fits31Bits n -> Int64
from31BitSigned (Witness31 SNat64 n
sn) = Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn)

-- | Creating a witness for a number being small (less than @2^31@)
fits31Bits :: SNat64 n -> Maybe (Fits31Bits n)
fits31Bits :: SNat64 n -> Maybe (Fits31Bits n)
fits31Bits sn :: SNat64 n
sn@(SNat64 Word64
n) 
  | Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
0 Bool -> Bool -> Bool
&& Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
2Word64 -> Integer -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
31   = Fits31Bits n -> Maybe (Fits31Bits n)
forall a. a -> Maybe a
Just (SNat64 n -> Fits31Bits n
forall (n :: Nat). SNat64 n -> Fits31Bits n
Witness31 SNat64 n
sn)
  | Bool
otherwise            = Maybe (Fits31Bits n)
forall a. Maybe a
Nothing

--------------------------------------------------------------------------------
-- * Small primes

newtype IsSmallPrime (n :: Nat) where
  SmallPrimeWitness :: SNat64 n -> IsSmallPrime n

deriving instance Show (IsSmallPrime n)

fromSmallPrime' :: IsSmallPrime n -> SNat64 n
fromSmallPrime' :: IsSmallPrime n -> SNat64 n
fromSmallPrime' (SmallPrimeWitness SNat64 n
sn) = SNat64 n
sn

fromSmallPrime :: IsSmallPrime n -> Word64
fromSmallPrime :: IsSmallPrime n -> Word64
fromSmallPrime (SmallPrimeWitness SNat64 n
sn) = SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn

fromSmallPrimeInteger :: IsSmallPrime n -> Integer
fromSmallPrimeInteger :: IsSmallPrime n -> Integer
fromSmallPrimeInteger (SmallPrimeWitness SNat64 n
sn) = Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn)

fromSmallPrimeSigned :: IsSmallPrime n -> Int64
fromSmallPrimeSigned :: IsSmallPrime n -> Int64
fromSmallPrimeSigned (SmallPrimeWitness SNat64 n
sn) = Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn)

-- | Prime testing.
--
-- Note: this uses trial division at the moment, 
-- so it's only good for small numbers for now
--
isSmallPrime :: SNat64 n -> Maybe (IsSmallPrime n)
isSmallPrime :: SNat64 n -> Maybe (IsSmallPrime n)
isSmallPrime SNat64 n
sn = 
  if (Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
1) Bool -> Bool -> Bool
&& (Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
2Word64 -> Integer -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
31) Bool -> Bool -> Bool
&& Integer -> Bool
Primes.isPrimeTrialDivision (Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n)
    then IsSmallPrime n -> Maybe (IsSmallPrime n)
forall a. a -> Maybe a
Just (SNat64 n -> IsSmallPrime n
forall (n :: Nat). SNat64 n -> IsSmallPrime n
SmallPrimeWitness SNat64 n
sn)
    else Maybe (IsSmallPrime n)
forall a. Maybe a
Nothing
  where
    n :: Word64
n = SNat64 n -> Word64
forall (n :: Nat). SNat64 n -> Word64
fromSNat64 SNat64 n
sn

smallPrimeIsPrime :: IsSmallPrime n -> IsPrime n
smallPrimeIsPrime :: IsSmallPrime n -> IsPrime n
smallPrimeIsPrime (SmallPrimeWitness (SNat64 Word64
n)) = SNat n -> IsPrime n
forall (n :: Nat). SNat n -> IsPrime n
PrimeWitness (Integer -> SNat n
forall (n :: Nat). Integer -> SNat n
SNat (Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n))

smallPrimeIsSmall :: IsSmallPrime n -> Fits31Bits n
smallPrimeIsSmall :: IsSmallPrime n -> Fits31Bits n
smallPrimeIsSmall (SmallPrimeWitness SNat64 n
sn) = SNat64 n -> Fits31Bits n
forall (n :: Nat). SNat64 n -> Fits31Bits n
Witness31 SNat64 n
sn

-- | Creating small primes
mkSmallPrime :: IsPrime p -> Fits31Bits p -> IsSmallPrime p
mkSmallPrime :: IsPrime p -> Fits31Bits p -> IsSmallPrime p
mkSmallPrime IsPrime p
_ (Witness31 SNat64 p
sn) = SNat64 p -> IsSmallPrime p
forall (n :: Nat). SNat64 n -> IsSmallPrime n
SmallPrimeWitness SNat64 p
sn

-- | Escape hatch
believeMeItsASmallPrime :: SNat64 n -> IsSmallPrime n
believeMeItsASmallPrime :: SNat64 n -> IsSmallPrime n
believeMeItsASmallPrime SNat64 n
sn = SNat64 n -> IsSmallPrime n
forall (n :: Nat). SNat64 n -> IsSmallPrime n
SmallPrimeWitness SNat64 n
sn

--------------------------------------------------------------------------------
-- * Divisors

-- | A proof that @k|n@
data Divides (k :: Nat) (n :: Nat) = Divides 
  { Divides k n -> Word64
_dividend :: {-# UNPACK #-} !Word64    -- ^ @n@
  , Divides k n -> Word64
_divisor  :: {-# UNPACK #-} !Word64    -- ^ @k@
  , Divides k n -> Word64
_quotient :: {-# UNPACK #-} !Word64    -- ^ @q=n/k@
  } 

dividendSNat :: Divides k n -> SNat64 n
dividendSNat :: Divides k n -> SNat64 n
dividendSNat (Divides Word64
n Word64
_ Word64
_) = Word64 -> SNat64 n
forall (n :: Nat). Word64 -> SNat64 n
SNat64 Word64
n

divisorSNat :: Divides k n -> SNat64 k
divisorSNat :: Divides k n -> SNat64 k
divisorSNat (Divides Word64
_ Word64
k Word64
_) = Word64 -> SNat64 k
forall (n :: Nat). Word64 -> SNat64 n
SNat64 Word64
k

divides :: SNat64 k -> SNat64 n -> Maybe (Divides k n)
divides :: SNat64 k -> SNat64 n -> Maybe (Divides k n)
divides (SNat64 Word64
k) (SNat64 Word64
n) = case Word64 -> Word64 -> (Word64, Word64)
forall a. Integral a => a -> a -> (a, a)
divMod Word64
n Word64
k of
  (Word64
q,Word64
r) -> if Word64
r Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0 then Divides k n -> Maybe (Divides k n)
forall a. a -> Maybe a
Just (Word64 -> Word64 -> Word64 -> Divides k n
forall (k :: Nat) (n :: Nat).
Word64 -> Word64 -> Word64 -> Divides k n
Divides Word64
n Word64
k Word64
q) else Maybe (Divides k n)
forall a. Maybe a
Nothing

instance Show (Divides k n) where
  show :: Divides k n -> String
show (Divides Word64
n Word64
k Word64
q) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall a. Show a => a -> String
show Word64
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"|" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word64 -> String
forall a. Show a => a -> String
show Word64
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

data Divisor (n :: Nat) 
  = forall k. Divisor (Divides k n)

deriving instance Show (Divisor n)

constructDivisor :: SNat64 n -> SNat64 k -> Maybe (Divisor n)
constructDivisor :: SNat64 n -> SNat64 k -> Maybe (Divisor n)
constructDivisor SNat64 n
sn SNat64 k
sk = case SNat64 k -> SNat64 n -> Maybe (Divides k n)
forall (k :: Nat) (n :: Nat).
SNat64 k -> SNat64 n -> Maybe (Divides k n)
divides SNat64 k
sk SNat64 n
sn of 
  Maybe (Divides k n)
Nothing -> Maybe (Divisor n)
forall a. Maybe a
Nothing
  Just Divides k n
d  -> Divisor n -> Maybe (Divisor n)
forall a. a -> Maybe a
Just (Divides k n -> Divisor n
forall (n :: Nat) (k :: Nat). Divides k n -> Divisor n
Divisor Divides k n
d)

divisors :: forall n. SNat64 n -> [Divisor n]
divisors :: SNat64 n -> [Divisor n]
divisors sn :: SNat64 n
sn@(SNat64 Word64
nn) = (Word64 -> Divisor n) -> [Word64] -> [Divisor n]
forall a b. (a -> b) -> [a] -> [b]
map Word64 -> Divisor n
worker [Word64]
ds where
    ds :: [Word64]
ds = [Word64] -> [Word64]
forall a. Ord a => [a] -> [a]
sort ((Integer -> Word64) -> [Integer] -> [Word64]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Integer] -> [Word64]) -> [Integer] -> [Word64]
forall a b. (a -> b) -> a -> b
$ Integer -> [Integer]
Primes.divisors (Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
nn)) :: [Word64]
    worker :: Word64 -> Divisor n
    worker :: Word64 -> Divisor n
worker Word64
d = case Word64 -> SomeSNat64
someSNat64_ Word64
d of
     SomeSNat64 SNat64 n
sd -> case SNat64 n -> SNat64 n -> Maybe (Divisor n)
forall (n :: Nat) (k :: Nat).
SNat64 n -> SNat64 k -> Maybe (Divisor n)
constructDivisor SNat64 n
sn SNat64 n
sd of
      Just Divisor n
proof     -> Divisor n
proof
      Maybe (Divisor n)
Nothing        -> String -> Divisor n
forall a. HasCallStack => String -> a
error String
"divisors: fatal error, should not happen"

--------------------------------------------------------------------------------
-- * Proxy

proxyOf :: a -> Proxy a
proxyOf :: a -> Proxy a
proxyOf a
_ = Proxy a
forall k (t :: k). Proxy t
Proxy

proxyOf1 :: f a -> Proxy a
proxyOf1 :: f a -> Proxy a
proxyOf1 f a
_ = Proxy a
forall k (t :: k). Proxy t
Proxy

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