-- |
-- 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 {
  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 a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
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
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 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
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
min :: MultMod m -> MultMod m -> MultMod m
$cmin :: forall (m :: Nat). MultMod m -> MultMod m -> MultMod m
max :: MultMod m -> MultMod m -> MultMod m
$cmax :: forall (m :: Nat). MultMod m -> MultMod m -> MultMod m
>= :: 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
$c< :: forall (m :: Nat). MultMod m -> MultMod m -> Bool
compare :: MultMod m -> MultMod m -> Ordering
$ccompare :: forall (m :: Nat). MultMod m -> MultMod m -> Ordering
$cp1Ord :: forall (m :: Nat). Eq (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 a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (m :: Nat). KnownNat m => Int -> MultMod m -> ShowS
forall (m :: Nat). KnownNat m => [MultMod m] -> ShowS
forall (m :: Nat). KnownNat m => MultMod m -> String
showList :: [MultMod m] -> ShowS
$cshowList :: forall (m :: Nat). KnownNat m => [MultMod m] -> ShowS
show :: MultMod m -> String
$cshow :: forall (m :: Nat). KnownNat m => MultMod m -> String
showsPrec :: Int -> MultMod m -> ShowS
$cshowsPrec :: forall (m :: Nat). KnownNat m => Int -> 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 :: 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 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 :: Mod m -> Maybe (MultMod m)
isMultElement Mod m
a = if Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
a Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`gcd` Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
a Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
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 :: 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
  { 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 a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
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
$c== :: forall (m :: Nat). 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 a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (m :: Nat). KnownNat m => Int -> PrimitiveRoot m -> ShowS
forall (m :: Nat). KnownNat m => [PrimitiveRoot m] -> ShowS
forall (m :: Nat). KnownNat m => PrimitiveRoot m -> String
showList :: [PrimitiveRoot m] -> ShowS
$cshowList :: forall (m :: Nat). KnownNat m => [PrimitiveRoot m] -> ShowS
show :: PrimitiveRoot m -> String
$cshow :: forall (m :: Nat). KnownNat m => PrimitiveRoot m -> String
showsPrec :: Int -> PrimitiveRoot m -> ShowS
$cshowsPrec :: forall (m :: Nat). KnownNat m => Int -> 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 :: 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 (() :: 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 (Natural -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
r))
    PrimitiveRoot m -> Maybe (PrimitiveRoot m)
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 :: CyclicGroup Integer m -> PrimitiveRoot m -> MultMod m -> Natural
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
    -> Natural
0
    -- the only valid input was a=1, b=1
  CyclicGroup Integer m
CG4
    -> if Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
b Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
1 then Natural
0 else Natural
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 -> Natural
discreteLogarithmPP Integer
p Word
k (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
a)) (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
b))
  CGDoubleOddPrimePower (Prime Integer -> Integer
forall a. Prime a -> a
unPrime -> Integer
p) Word
k
    -> Integer -> Word -> Integer -> Integer -> Natural
discreteLogarithmPP Integer
p Word
k (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
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) (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
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)*