-- |
-- 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 ScopedTypeVariables #-}

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.Integer.GMP.Internals
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' :: CyclicGroup a m -> a -> Bool
isPrimitiveRoot' CyclicGroup a m
cg a
r =
  case CyclicGroup a m
cg of
    CyclicGroup a m
CG2                       -> a
r a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1
    CyclicGroup a m
CG4                       -> a
r a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
3
    CGOddPrimePower Prime a
p Word
k       -> a -> Word -> a -> Bool
forall a a.
(UniqueFactorisation a, Num a, Eq a, Integral a) =>
a -> a -> a -> Bool
oddPrimePowerTest (Prime a -> a
forall a. Prime a -> a
unPrime Prime a
p) Word
k a
r
    CGDoubleOddPrimePower Prime a
p Word
k -> a -> Word -> a -> Bool
forall a a.
(Integral a, UniqueFactorisation a, Num a, Eq a) =>
a -> a -> a -> Bool
doubleOddPrimePowerTest (Prime a -> a
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       = a -> a -> Bool
forall b. (UniqueFactorisation b, Integral b) => b -> b -> Bool
oddPrimeTest a
p (a
g a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
p)
    oddPrimePowerTest a
p a
_ a
g       = a -> a -> Bool
forall b. (UniqueFactorisation b, Integral b) => b -> b -> Bool
oddPrimeTest a
p (a
g a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
p) Bool -> Bool -> Bool
&& case Natural -> SomeNat
someNatVal (a -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral' (a
p a -> a -> a
forall a. Num a => a -> a -> a
* a
p)) of
      SomeNat (Proxy n
_ :: Proxy pp) -> a -> Mod n
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
g Mod n -> a -> Mod n
forall a b. (Num a, Integral b) => a -> b -> a
^ (a
p a -> a -> a
forall a. Num a => a -> a -> a
- a
1) Mod n -> Mod n -> Bool
forall a. Eq a => a -> a -> Bool
/= (Mod n
1 :: Mod pp)

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

    oddPrimeTest :: b -> b -> Bool
oddPrimeTest b
p b
g = b
g b -> b -> Bool
forall a. Eq a => a -> a -> Bool
/= b
0 Bool -> Bool -> Bool
&& b -> b -> b
forall a. Integral a => a -> a -> a
gcd b
g b
p b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
1 Bool -> Bool -> Bool
&& case Natural -> SomeNat
someNatVal (b -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral' b
p) of
      SomeNat (Proxy n
_ :: Proxy p) -> (b -> Bool) -> [b] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\b
x -> b -> Mod n
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
g Mod n -> b -> Mod n
forall a b. (Num a, Integral b) => a -> b -> a
^ b
x Mod n -> Mod n -> Bool
forall a. Eq a => a -> a -> Bool
/= (Mod n
1 :: Mod p)) [b]
pows
      where
        pows :: [b]
pows = ((Prime b, Word) -> b) -> [(Prime b, Word)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (\(Prime b
q, Word
_) -> (b
p b -> b -> b
forall a. Num a => a -> a -> a
- b
1) b -> b -> b
forall a. Integral a => a -> a -> a
`quot` Prime b -> b
forall a. Prime a -> a
unPrime Prime b
q) (b -> [(Prime b, Word)]
forall a. UniqueFactorisation a => a -> [(Prime a, Word)]
factorise (b
p b -> b -> b
forall a. Num a => a -> a -> a
- b
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 = Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ if Integer
result Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then Integer
result Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
pkMinusPk1 else Integer
result
  where
    baseSol :: Integer
baseSol    = Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> Natural
discreteLogarithmPrime Integer
p (Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
p) (Integer
b Integer -> Integer -> Integer
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
pInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^(Word
kWord -> Word -> Word
forall a. Num a => a -> a -> a
-Word
1)
    c :: Integer
c          = (Integer -> Integer -> Integer
recipModInteger Integer
thetaA Integer
pkMinusOne Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
thetaB) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
pkMinusOne
    (Integer
result, Integer
pkMinusPk1) = Maybe (Integer, Integer) -> (Integer, Integer)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Integer, Integer) -> (Integer, Integer))
-> Maybe (Integer, Integer) -> (Integer, Integer)
forall a b. (a -> b) -> a -> b
$ (Integer, Integer)
-> (Integer, Integer) -> Maybe (Integer, Integer)
forall a.
(Eq a, Ring a, Euclidean a) =>
(a, a) -> (a, a) -> Maybe (a, a)
chinese (Integer
baseSol, Integer
pInteger -> Integer -> Integer
forall 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 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
pk) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
pkMinusOne
  where
    pk :: Integer
pk           = Integer
pkMinusOne Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p
    p2kMinusOne :: Integer
p2kMinusOne  = Integer
pkMinusOne Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
pk
    numerator :: Integer
numerator    = (Integer -> Integer -> Integer -> Integer
powModInteger Integer
a (Integer
pk Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
pkMinusOne) Integer
p2kMinusOne Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` 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 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
100000000 = Int -> Natural
intToNatural (Int -> Natural) -> Int -> Natural
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Int
discreteLogarithmPrimeBSGS (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
p) (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
a) (Integer -> Int
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 = [Int] -> Int
forall a. [a] -> a
head [Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j | (Int
v,Int
i) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
giants [Int
0..Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1], Int
j <- Maybe Int -> [Int]
forall a. Maybe a -> [a]
maybeToList (Int -> Map Int Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
v Map Int Int
table)]
  where
    m :: Int
m        = Int -> Int
forall a. Integral a => a -> a
integerSquareRoot (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 -- simple way of ceiling (sqrt (p-1))
    babies :: [Int]
babies   = (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
.* Int
a) Int
1
    table :: Map Int Int
table    = [(Int, Int)] -> Map Int Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
babies [Int
0..Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1])
    aInv :: Integer
aInv     = Integer -> Integer -> Integer
recipModInteger (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
a) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
p)
    bigGiant :: Int
bigGiant = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> Integer
powModInteger Integer
aInv (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
m) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
p)
    giants :: [Int]
giants   = (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
.* Int
bigGiant) Int
b
    Int
x .* :: Int -> Int -> Int
.* Int
y   = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
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 ((Integer, Integer) -> [Integer])
-> [(Integer, Integer)] -> [Integer]
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]
_)  -> Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
t
    []     -> [Char] -> Natural
forall a. HasCallStack => [Char] -> a
error ([Char]
"discreteLogarithm: pollard's rho failed, please report this as a bug. inputs " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Integer] -> [Char]
forall a. Show a => a -> [Char]
show [Integer
p,Integer
a,Integer
b])
  where
    n :: Integer
n                 = Integer
pInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1 -- order of the cyclic group
    halfN :: Integer
halfN             = Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` Integer
2
    mul2 :: Integer -> Integer
mul2 Integer
m            = if Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
halfN then Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
2 else Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
n
    sqrtN :: Integer
sqrtN             = Integer -> Integer
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 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
3 of
                          Integer
0 -> (Integer
xiInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
xi Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
p, Integer -> Integer
mul2 Integer
ai, Integer -> Integer
mul2 Integer
bi)
                          Integer
1 -> ( Integer
aInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
xi Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
p,    Integer
aiInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1,      Integer
bi)
                          Integer
_ -> ( Integer
bInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
xi Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
p,      Integer
ai,    Integer
biInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)
    initialise :: (Integer, Integer) -> (Integer, Integer, Integer)
initialise (Integer
x,Integer
y)  = (Integer -> Integer -> Integer -> Integer
powModInteger Integer
a Integer
x Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Integer -> Integer -> Integer
powModInteger Integer
b Integer
y Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
n, Integer
x, Integer
y)
    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           = Integer -> Integer -> Integer -> Integer
powModInteger Integer
a Integer
t Integer
p Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
b
    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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
x2i, Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd (Integer
bi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b2i) Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
sqrtN = case Natural -> SomeNat
someNatVal (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
n) of
        SomeNat (Proxy n
Proxy :: Proxy n) -> (Mod n -> Integer) -> [Mod n] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> (Mod n -> Natural) -> Mod n -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mod n -> Natural
forall (m :: Nat). Mod m -> Natural
unMod) ([Mod n] -> [Integer]) -> [Mod n] -> [Integer]
forall a b. (a -> b) -> a -> b
$ Mod n -> Mod n -> [Mod n]
forall (m :: Nat). KnownNat m => Mod m -> Mod m -> [Mod m]
solveLinear (Integer -> Mod n
forall a. Num a => Integer -> a
fromInteger (Integer
bi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b2i) :: Mod n) (Integer -> Mod n
forall a. Num a => Integer -> a
fromInteger (Integer
ai Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
a2i))
      | Integer
xi Integer -> Integer -> Bool
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        = (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter Integer -> Bool
check ([Integer] -> [Integer])
-> ((Integer, Integer) -> [Integer])
-> (Integer, Integer)
-> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer, Integer) -> [Integer]
begin ((Integer, Integer, Integer) -> [Integer])
-> ((Integer, Integer) -> (Integer, Integer, Integer))
-> (Integer, Integer)
-> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer) -> (Integer, Integer, Integer)
initialise