    -- |
    -- 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' :: 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 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, Integral a, Num a, Eq 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 {a}. (UniqueFactorisation a, Integral a) => a -> a -> 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 {a}. (UniqueFactorisation a, Integral a) => a -> a -> 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, Integral a, Num a, Eq a) =>
    a -> a -> a -> Bool
    oddPrimePowerTest a
    p a
    k a
    g

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