-- |
-- Module:      Math.NumberTheory.Moduli.Internal
-- Copyright:   (c) 2020 Bhavik Mehta
-- Licence:     MIT
-- Maintainer:  Bhavik Mehta <bhavikmehta8@gmail.com>
--
-- Multiplicative groups of integers modulo m.
--

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Math.NumberTheory.Moduli.Internal
  ( isPrimitiveRoot'
  , discreteLogarithmPP
  ) where

import qualified Data.Map as M
import Data.Maybe
import Data.Mod
import Data.Proxy
import GHC.TypeNats (SomeNat(..), someNatVal)
import GHC.Num.Integer
import Numeric.Natural

import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Equations
import Math.NumberTheory.Moduli.Singleton
import Math.NumberTheory.Primes
import Math.NumberTheory.Roots
import Math.NumberTheory.Utils.FromIntegral

-- https://en.wikipedia.org/wiki/Primitive_root_modulo_n#Finding_primitive_roots
isPrimitiveRoot'
  :: (Integral a, UniqueFactorisation a)
  => CyclicGroup a m
  -> a
  -> Bool
isPrimitiveRoot' :: forall a (m :: Natural).
(Integral a, UniqueFactorisation a) =>
CyclicGroup a m -> a -> Bool
isPrimitiveRoot' CyclicGroup a m
cg a
r =
  case CyclicGroup a m
cg of
    CyclicGroup a m
CG2                       -> a
r forall a. Eq a => a -> a -> Bool
== a
1
    CyclicGroup a m
CG4                       -> a
r forall a. Eq a => a -> a -> Bool
== a
3
    CGOddPrimePower Prime a
p Word
k       -> forall {a} {a}.
(UniqueFactorisation a, Num a, Eq a, Integral a) =>
a -> a -> a -> Bool
oddPrimePowerTest (forall a. Prime a -> a
unPrime Prime a
p) Word
k a
r
    CGDoubleOddPrimePower Prime a
p Word
k -> forall {a} {a}.
(Integral a, UniqueFactorisation a, Num a, Eq a) =>
a -> a -> a -> Bool
doubleOddPrimePowerTest (forall a. Prime a -> a
unPrime Prime a
p) Word
k a
r
  where
    oddPrimePowerTest :: a -> a -> a -> Bool
oddPrimePowerTest a
p a
1 a
g       = forall {a}. (UniqueFactorisation a, Integral a) => a -> a -> Bool
oddPrimeTest a
p (a
g forall a. Integral a => a -> a -> a
`mod` a
p)
    oddPrimePowerTest a
p a
_ a
g       = forall {a}. (UniqueFactorisation a, Integral a) => a -> a -> Bool
oddPrimeTest a
p (a
g forall a. Integral a => a -> a -> a
`mod` a
p) Bool -> Bool -> Bool
&& case Natural -> SomeNat
someNatVal (forall a b. (Integral a, Num b) => a -> b
fromIntegral' (a
p forall a. Num a => a -> a -> a
* a
p)) of
      SomeNat (Proxy n
_ :: Proxy pp) -> forall a b. (Integral a, Num b) => a -> b
fromIntegral a
g forall a b. (Num a, Integral b) => a -> b -> a
^ (a
p forall a. Num a => a -> a -> a
- a
1) forall a. Eq a => a -> a -> Bool
/= (Mod n
1 :: Mod pp)

    doubleOddPrimePowerTest :: a -> a -> a -> Bool
doubleOddPrimePowerTest a
p a
k a
g = forall a. Integral a => a -> Bool
odd a
g Bool -> Bool -> Bool
&& forall {a} {a}.
(UniqueFactorisation a, Num a, Eq a, Integral a) =>
a -> a -> a -> Bool
oddPrimePowerTest a
p a
k a
g

    oddPrimeTest :: a -> a -> Bool
oddPrimeTest a
p a
g = a
g forall a. Eq a => a -> a -> Bool
/= a
0 Bool -> Bool -> Bool
&& forall a. Integral a => a -> a -> a
gcd a
g a
p forall a. Eq a => a -> a -> Bool
== a
1 Bool -> Bool -> Bool
&& case Natural -> SomeNat
someNatVal (forall a b. (Integral a, Num b) => a -> b
fromIntegral' a
p) of
      SomeNat (Proxy n
_ :: Proxy p) -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\a
x -> forall a b. (Integral a, Num b) => a -> b
fromIntegral a
g forall a b. (Num a, Integral b) => a -> b -> a
^ a
x forall a. Eq a => a -> a -> Bool
/= (Mod n
1 :: Mod p)) [a]
pows
      where
        pows :: [a]
pows = forall a b. (a -> b) -> [a] -> [b]
map (\(Prime a
q, Word
_) -> (a
p forall a. Num a => a -> a -> a
- a
1) forall a. Integral a => a -> a -> a
`quot` forall a. Prime a -> a
unPrime Prime a
q) (forall a. UniqueFactorisation a => a -> [(Prime a, Word)]
factorise (a
p forall a. Num a => a -> a -> a
- a
1))

-- Implementation of Bach reduction (https://www2.eecs.berkeley.edu/Pubs/TechRpts/1984/CSD-84-186.pdf)
{-# INLINE discreteLogarithmPP #-}
discreteLogarithmPP :: Integer -> Word -> Integer -> Integer -> Natural
discreteLogarithmPP :: Integer -> Word -> Integer -> Integer -> Natural
discreteLogarithmPP Integer
p Word
1 Integer
a Integer
b = Integer -> Integer -> Integer -> Natural
discreteLogarithmPrime Integer
p Integer
a Integer
b
discreteLogarithmPP Integer
p Word
k Integer
a Integer
b = forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ if Integer
result forall a. Ord a => a -> a -> Bool
< Integer
0 then Integer
result forall a. Num a => a -> a -> a
+ Integer
pkMinusPk1 else Integer
result
  where
    baseSol :: Integer
baseSol    = forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> Natural
discreteLogarithmPrime Integer
p (Integer
a forall a. Integral a => a -> a -> a
`rem` Integer
p) (Integer
b forall a. Integral a => a -> a -> a
`rem` Integer
p)
    thetaA :: Integer
thetaA     = Integer -> Integer -> Integer -> Integer
theta Integer
p Integer
pkMinusOne Integer
a
    thetaB :: Integer
thetaB     = Integer -> Integer -> Integer -> Integer
theta Integer
p Integer
pkMinusOne Integer
b
    pkMinusOne :: Integer
pkMinusOne = Integer
pforall a b. (Num a, Integral b) => a -> b -> a
^(Word
kforall a. Num a => a -> a -> a
-Word
1)
    c :: Integer
c          = (forall a. Integral a => a -> Integer
toInteger Natural
t forall a. Num a => a -> a -> a
* Integer
thetaB) forall a. Integral a => a -> a -> a
`rem` Integer
pkMinusOne
      where
        (# Natural
t | #) = Integer -> Natural -> (# Natural | () #)
integerRecipMod# Integer
thetaA (forall a. Num a => Integer -> a
fromInteger Integer
pkMinusOne)
    (Integer
result, Integer
pkMinusPk1) = forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall a.
(Eq a, Ring a, Euclidean a) =>
(a, a) -> (a, a) -> Maybe (a, a)
chinese (Integer
baseSol, Integer
pforall a. Num a => a -> a -> a
-Integer
1) (Integer
c, Integer
pkMinusOne)

-- compute the homomorphism theta given in https://math.stackexchange.com/a/1864495/418148
{-# INLINE theta #-}
theta :: Integer -> Integer -> Integer -> Integer
theta :: Integer -> Integer -> Integer -> Integer
theta Integer
p Integer
pkMinusOne Integer
a = (Integer
numerator forall a. Integral a => a -> a -> a
`quot` Integer
pk) forall a. Integral a => a -> a -> a
`rem` Integer
pkMinusOne
  where
    pk :: Integer
pk           = Integer
pkMinusOne forall a. Num a => a -> a -> a
* Integer
p
    p2kMinusOne :: Integer
p2kMinusOne  = Integer
pkMinusOne forall a. Num a => a -> a -> a
* Integer
pk
    numerator :: Integer
numerator    = (forall a. Integral a => a -> Integer
toInteger Natural
t  forall a. Num a => a -> a -> a
- Integer
1) forall a. Integral a => a -> a -> a
`rem` Integer
p2kMinusOne
      where
        (# Natural
t | #) = Integer -> Integer -> Natural -> (# Natural | () #)
integerPowMod# Integer
a (Integer
pk forall a. Num a => a -> a -> a
- Integer
pkMinusOne) (forall a. Num a => Integer -> a
fromInteger Integer
p2kMinusOne)

-- TODO: Use Pollig-Hellman to reduce the problem further into groups of prime order.
-- While Bach reduction simplifies the problem into groups of the form (Z/pZ)*, these
-- have non-prime order, and the Pollig-Hellman algorithm can reduce the problem into
-- smaller groups of prime order.
-- In addition, the gcd check before solveLinear is applied in Pollard below will be
-- made redundant, since n would be prime.
discreteLogarithmPrime :: Integer -> Integer -> Integer -> Natural
discreteLogarithmPrime :: Integer -> Integer -> Integer -> Natural
discreteLogarithmPrime Integer
p Integer
a Integer
b
  | Integer
p forall a. Ord a => a -> a -> Bool
< Integer
100000000 = Int -> Natural
intToNatural forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Int
discreteLogarithmPrimeBSGS (forall a. Num a => Integer -> a
fromInteger Integer
p) (forall a. Num a => Integer -> a
fromInteger Integer
a) (forall a. Num a => Integer -> a
fromInteger Integer
b)
  | Bool
otherwise     = Integer -> Integer -> Integer -> Natural
discreteLogarithmPrimePollard Integer
p Integer
a Integer
b

discreteLogarithmPrimeBSGS :: Int -> Int -> Int -> Int
discreteLogarithmPrimeBSGS :: Int -> Int -> Int -> Int
discreteLogarithmPrimeBSGS Int
p Int
a Int
b =
  case [Int
iforall a. Num a => a -> a -> a
*Int
m forall a. Num a => a -> a -> a
+ Int
j | (Int
v,Int
i) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
giants [Int
0..Int
mforall a. Num a => a -> a -> a
-Int
1], Int
j <- forall a. Maybe a -> [a]
maybeToList (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Map Int Int
table)] of
    [] -> forall a. HasCallStack => [Char] -> a
error ([Char]
"discreteLogarithmPrimeBSGS: failed, please report this as a bug. Inputs: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show [Int
p,Int
a,Int
b])
    Int
hd : [Int]
_ -> Int
hd
  where
    m :: Int
    m :: Int
m        = forall a. Integral a => a -> a
integerSquareRoot (Int
p forall a. Num a => a -> a -> a
- Int
2) forall a. Num a => a -> a -> a
+ Int
1 -- simple way of ceiling (sqrt (p-1))

    babies :: [Int]
    babies :: [Int]
babies   = forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
.* Int
a) Int
1

    table :: M.Map Int Int
    table :: Map Int Int
table    = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
babies [Int
0..Int
mforall a. Num a => a -> a -> a
-Int
1])

    aInv :: Integer
    aInv :: Integer
aInv     = forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
ap
      where
        (# Natural
ap | #) = Integer -> Natural -> (# Natural | () #)
integerRecipMod# (forall a. Integral a => a -> Integer
toInteger Int
a) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p)

    bigGiant :: Int
    bigGiant :: Int
bigGiant = forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
aInvmp
      where
        (# Natural
aInvmp | #) = Integer -> Integer -> Natural -> (# Natural | () #)
integerPowMod# Integer
aInv (forall a. Integral a => a -> Integer
toInteger Int
m) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p)

    giants :: [Int]
    giants :: [Int]
giants   = forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
.* Int
bigGiant) Int
b

    (.*) :: Int -> Int -> Int
    Int
x .* :: Int -> Int -> Int
.* Int
y   = Int
x forall a. Num a => a -> a -> a
* Int
y forall a. Integral a => a -> a -> a
`rem` Int
p

-- TODO: Use more advanced walks, in order to reduce divisions, cf
-- https://maths-people.anu.edu.au/~brent/pd/rpb231.pdf
-- This will slightly improve the expected time to collision, and can reduce the
-- number of divisions performed.
discreteLogarithmPrimePollard :: Integer -> Integer -> Integer -> Natural
discreteLogarithmPrimePollard :: Integer -> Integer -> Integer -> Natural
discreteLogarithmPrimePollard Integer
p Integer
a Integer
b =
  case forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Integer, Integer) -> [Integer]
runPollard [(Integer
x,Integer
y) | Integer
x <- [Integer
0..Integer
n], Integer
y <- [Integer
0..Integer
n]] of
    (Integer
t:[Integer]
_)  -> forall a. Num a => Integer -> a
fromInteger Integer
t
    []     -> forall a. HasCallStack => [Char] -> a
error ([Char]
"discreteLogarithm: pollard's rho failed, please report this as a bug. Inputs: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show [Integer
p,Integer
a,Integer
b])
  where
    n :: Integer
n                 = Integer
pforall a. Num a => a -> a -> a
-Integer
1 -- order of the cyclic group
    halfN :: Integer
halfN             = Integer
n forall a. Integral a => a -> a -> a
`quot` Integer
2
    mul2 :: Integer -> Integer
mul2 Integer
m            = if Integer
m forall a. Ord a => a -> a -> Bool
< Integer
halfN then Integer
m forall a. Num a => a -> a -> a
* Integer
2 else Integer
m forall a. Num a => a -> a -> a
* Integer
2 forall a. Num a => a -> a -> a
- Integer
n
    sqrtN :: Integer
sqrtN             = forall a. Integral a => a -> a
integerSquareRoot Integer
n
    step :: (Integer, Integer, Integer) -> (Integer, Integer, Integer)
step (Integer
xi,!Integer
ai,!Integer
bi) = case Integer
xi forall a. Integral a => a -> a -> a
`rem` Integer
3 of
                          Integer
0 -> (Integer
xiforall a. Num a => a -> a -> a
*Integer
xi forall a. Integral a => a -> a -> a
`rem` Integer
p, Integer -> Integer
mul2 Integer
ai, Integer -> Integer
mul2 Integer
bi)
                          Integer
1 -> ( Integer
aforall a. Num a => a -> a -> a
*Integer
xi forall a. Integral a => a -> a -> a
`rem` Integer
p,    Integer
aiforall a. Num a => a -> a -> a
+Integer
1,      Integer
bi)
                          Integer
_ -> ( Integer
bforall a. Num a => a -> a -> a
*Integer
xi forall a. Integral a => a -> a -> a
`rem` Integer
p,      Integer
ai,    Integer
biforall a. Num a => a -> a -> a
+Integer
1)
    initialise :: (Integer, Integer) -> (Integer, Integer, Integer)
initialise (Integer
x,Integer
y)  = (forall a. Integral a => a -> Integer
toInteger Natural
axn forall a. Num a => a -> a -> a
* forall a. Integral a => a -> Integer
toInteger Natural
byn forall a. Integral a => a -> a -> a
`rem` Integer
n, Integer
x, Integer
y)
      where
        (# Natural
axn | #) = Integer -> Integer -> Natural -> (# Natural | () #)
integerPowMod# Integer
a Integer
x (forall a. Num a => Integer -> a
fromInteger Integer
n)
        (# Natural
byn | #) = Integer -> Integer -> Natural -> (# Natural | () #)
integerPowMod# Integer
b Integer
y (forall a. Num a => Integer -> a
fromInteger Integer
n)
    begin :: (Integer, Integer, Integer) -> [Integer]
begin (Integer, Integer, Integer)
t           = (Integer, Integer, Integer)
-> (Integer, Integer, Integer) -> [Integer]
go ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step (Integer, Integer, Integer)
t) ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step (Integer, Integer, Integer)
t))
    check :: Integer -> Bool
check Integer
t           = case Integer -> Integer -> Natural -> (# Natural | () #)
integerPowMod# Integer
a Integer
t (forall a. Num a => Integer -> a
fromInteger Integer
p) of
      (# Natural
atp | #) -> forall a. Integral a => a -> Integer
toInteger Natural
atp forall a. Eq a => a -> a -> Bool
== Integer
b
      (# | ()
_ #) -> Bool
False
    go :: (Integer, Integer, Integer)
-> (Integer, Integer, Integer) -> [Integer]
go tort :: (Integer, Integer, Integer)
tort@(Integer
xi,Integer
ai,Integer
bi) hare :: (Integer, Integer, Integer)
hare@(Integer
x2i,Integer
a2i,Integer
b2i)
      | Integer
xi forall a. Eq a => a -> a -> Bool
== Integer
x2i, forall a. Integral a => a -> a -> a
gcd (Integer
bi forall a. Num a => a -> a -> a
- Integer
b2i) Integer
n forall a. Ord a => a -> a -> Bool
< Integer
sqrtN = case Natural -> SomeNat
someNatVal (forall a. Num a => Integer -> a
fromInteger Integer
n) of
        SomeNat (Proxy n
Proxy :: Proxy n) -> forall a b. (a -> b) -> [a] -> [b]
map (forall a. Integral a => a -> Integer
toInteger forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: Natural). Mod m -> Natural
unMod) forall a b. (a -> b) -> a -> b
$ forall (m :: Natural). KnownNat m => Mod m -> Mod m -> [Mod m]
solveLinear (forall a. Num a => Integer -> a
fromInteger (Integer
bi forall a. Num a => a -> a -> a
- Integer
b2i) :: Mod n) (forall a. Num a => Integer -> a
fromInteger (Integer
ai forall a. Num a => a -> a -> a
- Integer
a2i))
      | Integer
xi forall a. Eq a => a -> a -> Bool
== Integer
x2i                           = []
      | Bool
otherwise                           = (Integer, Integer, Integer)
-> (Integer, Integer, Integer) -> [Integer]
go ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step (Integer, Integer, Integer)
tort) ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step ((Integer, Integer, Integer) -> (Integer, Integer, Integer)
step (Integer, Integer, Integer)
hare))
    runPollard :: (Integer, Integer) -> [Integer]
runPollard        = forall a. (a -> Bool) -> [a] -> [a]
filter Integer -> Bool
check forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer, Integer) -> [Integer]
begin forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer) -> (Integer, Integer, Integer)
initialise