    -- |
    -- Module:      Math.NumberTheory.Moduli.Multiplicative
    -- Copyright:   (c) 2017 Andrew Lelechenko
    -- Licence:     MIT
    -- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
    --
    -- Multiplicative groups of integers modulo m.
    --

    {-# LANGUAGE ScopedTypeVariables #-}
    {-# LANGUAGE ViewPatterns #-}

    module Math.NumberTheory.Moduli.Multiplicative
      ( -- * Multiplicative group
        MultMod
      , multElement
      , isMultElement
      , invertGroup
      -- * Primitive roots
      , PrimitiveRoot
      , unPrimitiveRoot
      , isPrimitiveRoot
      , discreteLogarithm
      ) where

    import Control.Monad
    import Data.Constraint
    import Data.Mod
    import Data.Semigroup
    import GHC.TypeNats (KnownNat, natVal)
    import Numeric.Natural

    import Math.NumberTheory.Moduli.Internal
    import Math.NumberTheory.Moduli.Singleton
    import Math.NumberTheory.Primes

    -- | This type represents elements of the multiplicative group mod m, i.e.
    -- those elements which are coprime to m. Use @isMultElement@ to construct.
    newtype MultMod m = MultMod {
      forall (m :: Nat). MultMod m -> Mod m
    multElement :: Mod m -- ^ Unwrap a residue.
      } deriving (MultMod m -> MultMod m -> Bool
    (MultMod m -> MultMod m -> Bool)
    -> (MultMod m -> MultMod m -> Bool) -> Eq (MultMod m)
    forall (m :: Nat). MultMod m -> MultMod m -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    == :: MultMod m -> MultMod m -> Bool
    $c/= :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    /= :: MultMod m -> MultMod m -> Bool
    Eq, Eq (MultMod m)
    Eq (MultMod m) =>
    (MultMod m -> MultMod m -> Ordering)
    -> (MultMod m -> MultMod m -> Bool)
    -> (MultMod m -> MultMod m -> Bool)
    -> (MultMod m -> MultMod m -> Bool)
    -> (MultMod m -> MultMod m -> Bool)
    -> (MultMod m -> MultMod m -> MultMod m)
    -> (MultMod m -> MultMod m -> MultMod m)
    -> Ord (MultMod m)
    MultMod m -> MultMod m -> Bool
    MultMod m -> MultMod m -> Ordering
    MultMod m -> MultMod m -> MultMod m
    forall (m :: Nat). Eq (MultMod m)
    forall (m :: Nat). MultMod m -> MultMod m -> Bool
    forall (m :: Nat). MultMod m -> MultMod m -> Ordering
    forall (m :: Nat). MultMod m -> MultMod m -> MultMod m
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: forall (m :: Nat). MultMod m -> MultMod m -> Ordering
    compare :: MultMod m -> MultMod m -> Ordering
    $c< :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    < :: MultMod m -> MultMod m -> Bool
    $c<= :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    <= :: MultMod m -> MultMod m -> Bool
    $c> :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    > :: MultMod m -> MultMod m -> Bool
    $c>= :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
    >= :: MultMod m -> MultMod m -> Bool
    $cmax :: forall (m :: Nat). MultMod m -> MultMod m -> MultMod m
    max :: MultMod m -> MultMod m -> MultMod m
    $cmin :: forall (m :: Nat). MultMod m -> MultMod m -> MultMod m
    min :: MultMod m -> MultMod m -> MultMod m
    Ord, Int -> MultMod m -> ShowS
    [MultMod m] -> ShowS
    MultMod m -> String
    (Int -> MultMod m -> ShowS)
    -> (MultMod m -> String)
    -> ([MultMod m] -> ShowS)
    -> Show (MultMod m)
    forall (m :: Nat). Int -> MultMod m -> ShowS
    forall (m :: Nat). [MultMod m] -> ShowS
    forall (m :: Nat). MultMod m -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall (m :: Nat). Int -> MultMod m -> ShowS
    showsPrec :: Int -> MultMod m -> ShowS
    $cshow :: forall (m :: Nat). MultMod m -> String
    show :: MultMod m -> String
    $cshowList :: forall (m :: Nat). [MultMod m] -> ShowS
    showList :: [MultMod m] -> ShowS
    Show)

    instance KnownNat m => Semigroup (MultMod m) where
      MultMod Mod m
    a <> :: MultMod m -> MultMod m -> MultMod m
    <> MultMod Mod m
    b = Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod (Mod m
    a Mod m -> Mod m -> Mod m
    forall a. Num a => a -> a -> a
    * Mod m
    b)
      stimes :: forall b. Integral b => b -> MultMod m -> MultMod m
    stimes b
    k a :: MultMod m
    a@(MultMod Mod m
    a')
        | b
    k b -> b -> Bool
    forall a. Ord a => a -> a -> Bool
    >= b
    0 = Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod (Mod m
    a' Mod m -> b -> Mod m
    forall (m :: Nat) a.
    (KnownNat m, Integral a) =>
    Mod m -> a -> Mod m
    ^% b
    k)
        | Bool
    otherwise = MultMod m -> MultMod m
    forall (m :: Nat). KnownNat m => MultMod m -> MultMod m
    invertGroup (MultMod m -> MultMod m) -> MultMod m -> MultMod m
    forall a b. (a -> b) -> a -> b
    $ b -> MultMod m -> MultMod m
    forall b. Integral b => b -> MultMod m -> MultMod m
    forall a b. (Semigroup a, Integral b) => b -> a -> a
    stimes (-b
    k) MultMod m
    a
      -- ^ This Semigroup is in fact a group, so @stimes@ can be called with a negative first argument.

    instance KnownNat m => Monoid (MultMod m) where
      mempty :: MultMod m
    mempty = Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod Mod m
    1
      mappend :: MultMod m -> MultMod m -> MultMod m
    mappend = MultMod m -> MultMod m -> MultMod m
    forall a. Semigroup a => a -> a -> a
    (<>)

    instance KnownNat m => Bounded (MultMod m) where
      minBound :: MultMod m
    minBound = Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod Mod m
    1
      maxBound :: MultMod m
    maxBound = Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod (-Mod m
    1)

    -- | Attempt to construct a multiplicative group element.
    isMultElement :: KnownNat m => Mod m -> Maybe (MultMod m)
    isMultElement :: forall (m :: Nat). KnownNat m => Mod m -> Maybe (MultMod m)
    isMultElement Mod m
    a = if Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    a Nat -> Nat -> Nat
    forall a. Integral a => a -> a -> a
    `gcd` Mod m -> Nat
    forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
    natVal Mod m
    a Nat -> Nat -> Bool
    forall a. Eq a => a -> a -> Bool
    == Nat
    1
                         then MultMod m -> Maybe (MultMod m)
    forall a. a -> Maybe a
    Just (MultMod m -> Maybe (MultMod m)) -> MultMod m -> Maybe (MultMod m)
    forall a b. (a -> b) -> a -> b
    $ Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod Mod m
    a
                         else Maybe (MultMod m)
    forall a. Maybe a
    Nothing

    -- | For elements of the multiplicative group, we can safely perform the inverse
    -- without needing to worry about failure.
    invertGroup :: KnownNat m => MultMod m -> MultMod m
    invertGroup :: forall (m :: Nat). KnownNat m => MultMod m -> MultMod m
    invertGroup (MultMod Mod m
    a) = case Mod m -> Maybe (Mod m)
    forall (m :: Nat). KnownNat m => Mod m -> Maybe (Mod m)
    invertMod Mod m
    a of
                                Just Mod m
    b -> Mod m -> MultMod m
    forall (m :: Nat). Mod m -> MultMod m
    MultMod Mod m
    b
                                Maybe (Mod m)
    Nothing -> String -> MultMod m
    forall a. HasCallStack => String -> a
    error String
    "Math.NumberTheory.Moduli.invertGroup: failed to invert element"

    -- | 'PrimitiveRoot' m is a type which is only inhabited
    -- by <https://en.wikipedia.org/wiki/Primitive_root_modulo_n primitive roots> of m.
    newtype PrimitiveRoot m = PrimitiveRoot
      { forall (m :: Nat). PrimitiveRoot m -> MultMod m
    unPrimitiveRoot :: MultMod m -- ^ Extract primitive root value.
      }
      deriving (PrimitiveRoot m -> PrimitiveRoot m -> Bool
    (PrimitiveRoot m -> PrimitiveRoot m -> Bool)
    -> (PrimitiveRoot m -> PrimitiveRoot m -> Bool)
    -> Eq (PrimitiveRoot m)
    forall (m :: Nat). PrimitiveRoot m -> PrimitiveRoot m -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall (m :: Nat). PrimitiveRoot m -> PrimitiveRoot m -> Bool
    == :: PrimitiveRoot m -> PrimitiveRoot m -> Bool
    $c/= :: forall (m :: Nat). PrimitiveRoot m -> PrimitiveRoot m -> Bool
    /= :: PrimitiveRoot m -> PrimitiveRoot m -> Bool
    Eq, Int -> PrimitiveRoot m -> ShowS
    [PrimitiveRoot m] -> ShowS
    PrimitiveRoot m -> String
    (Int -> PrimitiveRoot m -> ShowS)
    -> (PrimitiveRoot m -> String)
    -> ([PrimitiveRoot m] -> ShowS)
    -> Show (PrimitiveRoot m)
    forall (m :: Nat). Int -> PrimitiveRoot m -> ShowS
    forall (m :: Nat). [PrimitiveRoot m] -> ShowS
    forall (m :: Nat). PrimitiveRoot m -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall (m :: Nat). Int -> PrimitiveRoot m -> ShowS
    showsPrec :: Int -> PrimitiveRoot m -> ShowS
    $cshow :: forall (m :: Nat). PrimitiveRoot m -> String
    show :: PrimitiveRoot m -> String
    $cshowList :: forall (m :: Nat). [PrimitiveRoot m] -> ShowS
    showList :: [PrimitiveRoot m] -> ShowS
    Show)

    -- | Check whether a given modular residue is
    -- a <https://en.wikipedia.org/wiki/Primitive_root_modulo_n primitive root>.
    --
    -- >>> :set -XDataKinds
    -- >>> import Data.Maybe
    -- >>> isPrimitiveRoot (fromJust cyclicGroup) (1 :: Mod 13)
    -- Nothing
    -- >>> isPrimitiveRoot (fromJust cyclicGroup) (2 :: Mod 13)
    -- Just (PrimitiveRoot {unPrimitiveRoot = MultMod {multElement = (2 `modulo` 13)}})
    isPrimitiveRoot
      :: (Integral a, UniqueFactorisation a)
      => CyclicGroup a m
      -> Mod m
      -> Maybe (PrimitiveRoot m)
    isPrimitiveRoot :: forall a (m :: Nat).
    (Integral a, UniqueFactorisation a) =>
    CyclicGroup a m -> Mod m -> Maybe (PrimitiveRoot m)
    isPrimitiveRoot CyclicGroup a m
    cg Mod m
    r = case CyclicGroup a m -> (() :: Constraint) :- KnownNat m
    forall a (m :: Nat).
    Integral a =>
    CyclicGroup a m -> (() :: Constraint) :- KnownNat m
    proofFromCyclicGroup CyclicGroup a m
    cg of
      Sub Dict (KnownNat m)
    (() :: Constraint) => Dict (KnownNat m)
    Dict -> do
        MultMod m
    r' <- Mod m -> Maybe (MultMod m)
    forall (m :: Nat). KnownNat m => Mod m -> Maybe (MultMod m)
    isMultElement Mod m
    r
        Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ CyclicGroup a m -> a -> Bool
    forall a (m :: Nat).
    (Integral a, UniqueFactorisation a) =>
    CyclicGroup a m -> a -> Bool
    isPrimitiveRoot' CyclicGroup a m
    cg (Nat -> a
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    r))
        PrimitiveRoot m -> Maybe (PrimitiveRoot m)
    forall a. a -> Maybe a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (PrimitiveRoot m -> Maybe (PrimitiveRoot m))
    -> PrimitiveRoot m -> Maybe (PrimitiveRoot m)
    forall a b. (a -> b) -> a -> b
    $ MultMod m -> PrimitiveRoot m
    forall (m :: Nat). MultMod m -> PrimitiveRoot m
    PrimitiveRoot MultMod m
    r'

    -- | Computes the discrete logarithm. Currently uses a combination of the baby-step
    -- giant-step method and Pollard's rho algorithm, with Bach reduction.
    --
    -- >>> :set -XDataKinds
    -- >>> import Data.Maybe
    -- >>> let cg = fromJust cyclicGroup :: CyclicGroup Integer 13
    -- >>> let rt = fromJust (isPrimitiveRoot cg 2)
    -- >>> let x  = fromJust (isMultElement 11)
    -- >>> discreteLogarithm cg rt x
    -- 7
    discreteLogarithm :: CyclicGroup Integer m -> PrimitiveRoot m -> MultMod m -> Natural
    discreteLogarithm :: forall (m :: Nat).
    CyclicGroup Integer m -> PrimitiveRoot m -> MultMod m -> Nat
    discreteLogarithm CyclicGroup Integer m
    cg (MultMod m -> Mod m
    forall (m :: Nat). MultMod m -> Mod m
    multElement (MultMod m -> Mod m)
    -> (PrimitiveRoot m -> MultMod m) -> PrimitiveRoot m -> Mod m
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PrimitiveRoot m -> MultMod m
    forall (m :: Nat). PrimitiveRoot m -> MultMod m
    unPrimitiveRoot -> Mod m
    a) (MultMod m -> Mod m
    forall (m :: Nat). MultMod m -> Mod m
    multElement -> Mod m
    b) = case CyclicGroup Integer m
    cg of
      CyclicGroup Integer m
    CG2
        -> Nat
    0
        -- the only valid input was a=1, b=1
      CyclicGroup Integer m
    CG4
        -> if Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    b Nat -> Nat -> Bool
    forall a. Eq a => a -> a -> Bool
    == Nat
    1 then Nat
    0 else Nat
    1
        -- the only possible input here is a=3 with b = 1 or 3
      CGOddPrimePower (Prime Integer -> Integer
    forall a. Prime a -> a
    unPrime -> Integer
    p) Word
    k
        -> Integer -> Word -> Integer -> Integer -> Nat
    discreteLogarithmPP Integer
    p Word
    k (Nat -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    a)) (Nat -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    b))
      CGDoubleOddPrimePower (Prime Integer -> Integer
    forall a. Prime a -> a
    unPrime -> Integer
    p) Word
    k
        -> Integer -> Word -> Integer -> Integer -> Nat
    discreteLogarithmPP Integer
    p Word
    k (Nat -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    a) Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    pInteger -> Word -> Integer
    forall a b. (Num a, Integral b) => a -> b -> a
    ^Word
    k) (Nat -> Integer
    forall a. Integral a => a -> Integer
    toInteger (Mod m -> Nat
    forall (m :: Nat). Mod m -> Nat
    unMod Mod m
    b) Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    pInteger -> Word -> Integer
    forall a b. (Num a, Integral b) => a -> b -> a
    ^Word
    k)
        -- we have the isomorphism t -> t `rem` p^k from (Z/2p^kZ)* -> (Z/p^kZ)*