-- |
-- Module:      Data.Mod
-- Copyright:   (c) 2017-2022 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- <https://en.wikipedia.org/wiki/Modular_arithmetic Modular arithmetic>,
-- promoting moduli to the type level, with an emphasis on performance.
-- Originally part of <https://hackage.haskell.org/package/arithmoi arithmoi> package.
--
-- This module supports moduli of arbitrary size.
-- Use "Data.Mod.Word" to achieve better performance,
-- when your moduli fit into 'Word'.

{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UnboxedTuples         #-}

module Data.Mod
  ( Mod
  , unMod
  , invertMod
  , (^%)
  ) where

import Control.Exception
import Control.DeepSeq
import Data.Ratio
#ifdef MIN_VERSION_semirings
import Data.Euclidean (GcdDomain(..), Euclidean(..), Field)
import Data.Semiring (Semiring(..), Ring(..))
#endif
import GHC.Exts
import GHC.Generics
import GHC.Natural (Natural(..), powModNatural)
import GHC.TypeNats (Nat, KnownNat, natVal)

-- | This data type represents
-- <https://en.wikipedia.org/wiki/Modular_arithmetic#Integers_modulo_n integers modulo m>,
-- equipped with useful instances.
--
-- For example, 3 :: 'Mod' 10 stands for the class of integers
-- congruent to \( 3 \bmod 10 \colon \ldots {}−17, −7, 3, 13, 23 \ldots \)
--
-- >>> :set -XDataKinds
-- >>> 3 + 8 :: Mod 10 -- 3 + 8 = 11 ≡ 1 (mod 10)
-- (1 `modulo` 10)
--
-- __Warning:__ division by residue, which is not
-- <https://en.wikipedia.org/wiki/Coprime_integers coprime>
-- with the modulo, throws 'DivideByZero'.
-- Consider using 'invertMod' for non-prime moduli.
newtype Mod (m :: Nat) = Mod
  { Mod m -> Natural
unMod :: Natural
  -- ^ The canonical representative of the residue class,
  -- always between 0 and \( m - 1 \) inclusively.
  --
  -- >>> :set -XDataKinds
  -- >>> -1 :: Mod 10
  -- (9 `modulo` 10)
  }
  deriving (Mod m -> Mod m -> Bool
(Mod m -> Mod m -> Bool) -> (Mod m -> Mod m -> Bool) -> Eq (Mod m)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (m :: Nat). Mod m -> Mod m -> Bool
/= :: Mod m -> Mod m -> Bool
$c/= :: forall (m :: Nat). Mod m -> Mod m -> Bool
== :: Mod m -> Mod m -> Bool
$c== :: forall (m :: Nat). Mod m -> Mod m -> Bool
Eq, Eq (Mod m)
Eq (Mod m)
-> (Mod m -> Mod m -> Ordering)
-> (Mod m -> Mod m -> Bool)
-> (Mod m -> Mod m -> Bool)
-> (Mod m -> Mod m -> Bool)
-> (Mod m -> Mod m -> Bool)
-> (Mod m -> Mod m -> Mod m)
-> (Mod m -> Mod m -> Mod m)
-> Ord (Mod m)
Mod m -> Mod m -> Bool
Mod m -> Mod m -> Ordering
Mod m -> Mod m -> Mod 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 (Mod m)
forall (m :: Nat). Mod m -> Mod m -> Bool
forall (m :: Nat). Mod m -> Mod m -> Ordering
forall (m :: Nat). Mod m -> Mod m -> Mod m
min :: Mod m -> Mod m -> Mod m
$cmin :: forall (m :: Nat). Mod m -> Mod m -> Mod m
max :: Mod m -> Mod m -> Mod m
$cmax :: forall (m :: Nat). Mod m -> Mod m -> Mod m
>= :: Mod m -> Mod m -> Bool
$c>= :: forall (m :: Nat). Mod m -> Mod m -> Bool
> :: Mod m -> Mod m -> Bool
$c> :: forall (m :: Nat). Mod m -> Mod m -> Bool
<= :: Mod m -> Mod m -> Bool
$c<= :: forall (m :: Nat). Mod m -> Mod m -> Bool
< :: Mod m -> Mod m -> Bool
$c< :: forall (m :: Nat). Mod m -> Mod m -> Bool
compare :: Mod m -> Mod m -> Ordering
$ccompare :: forall (m :: Nat). Mod m -> Mod m -> Ordering
$cp1Ord :: forall (m :: Nat). Eq (Mod m)
Ord, (forall x. Mod m -> Rep (Mod m) x)
-> (forall x. Rep (Mod m) x -> Mod m) -> Generic (Mod m)
forall x. Rep (Mod m) x -> Mod m
forall x. Mod m -> Rep (Mod m) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (m :: Nat) x. Rep (Mod m) x -> Mod m
forall (m :: Nat) x. Mod m -> Rep (Mod m) x
$cto :: forall (m :: Nat) x. Rep (Mod m) x -> Mod m
$cfrom :: forall (m :: Nat) x. Mod m -> Rep (Mod m) x
Generic)

instance NFData (Mod m)

instance KnownNat m => Show (Mod m) where
  show :: Mod m -> String
show Mod m
m = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Natural -> String
forall a. Show a => a -> String
show (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
m) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" `modulo` " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Natural -> String
forall a. Show a => a -> String
show (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
m) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

instance KnownNat m => Enum (Mod m) where
  succ :: Mod m -> Mod m
succ Mod m
x = if Mod m
x Mod m -> Mod m -> Bool
forall a. Eq a => a -> a -> Bool
== Mod m
forall a. Bounded a => a
maxBound then ArithException -> Mod m
forall a e. Exception e => e -> a
throw ArithException
Overflow  else (Natural -> Natural) -> Mod m -> Mod m
coerce (Enum Natural => Natural -> Natural
forall a. Enum a => a -> a
succ @Natural) Mod m
x
  pred :: Mod m -> Mod m
pred Mod m
x = if Mod m
x Mod m -> Mod m -> Bool
forall a. Eq a => a -> a -> Bool
== Mod m
forall a. Bounded a => a
minBound then ArithException -> Mod m
forall a e. Exception e => e -> a
throw ArithException
Underflow else (Natural -> Natural) -> Mod m -> Mod m
coerce (Enum Natural => Natural -> Natural
forall a. Enum a => a -> a
pred @Natural) Mod m
x

  toEnum :: Int -> Mod m
toEnum   = (Int -> Mod m
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int -> Mod m)
  fromEnum :: Mod m -> Int
fromEnum = (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Natural -> Int) (Natural -> Int) -> (Mod m -> Natural) -> Mod m -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod

  enumFrom :: Mod m -> [Mod m]
enumFrom Mod m
x       = Mod m -> Mod m -> [Mod m]
forall a. Enum a => a -> a -> [a]
enumFromTo Mod m
x Mod m
forall a. Bounded a => a
maxBound
  enumFromThen :: Mod m -> Mod m -> [Mod m]
enumFromThen Mod m
x Mod m
y = Mod m -> Mod m -> Mod m -> [Mod m]
forall a. Enum a => a -> a -> a -> [a]
enumFromThenTo Mod m
x Mod m
y (if Mod m
y Mod m -> Mod m -> Bool
forall a. Ord a => a -> a -> Bool
>= Mod m
x then Mod m
forall a. Bounded a => a
maxBound else Mod m
forall a. Bounded a => a
minBound)

  enumFromTo :: Mod m -> Mod m -> [Mod m]
enumFromTo     = (Natural -> Natural -> [Natural]) -> Mod m -> Mod m -> [Mod m]
coerce (Enum Natural => Natural -> Natural -> [Natural]
forall a. Enum a => a -> a -> [a]
enumFromTo     @Natural)
  enumFromThenTo :: Mod m -> Mod m -> Mod m -> [Mod m]
enumFromThenTo = (Natural -> Natural -> Natural -> [Natural])
-> Mod m -> Mod m -> Mod m -> [Mod m]
coerce (Enum Natural => Natural -> Natural -> Natural -> [Natural]
forall a. Enum a => a -> a -> a -> [a]
enumFromThenTo @Natural)

instance KnownNat m => Bounded (Mod m) where
  minBound :: Mod m
minBound = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
0
  maxBound :: Mod m
maxBound = let mx :: Mod m
mx = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) in Mod m
mx

addMod :: Natural -> Natural -> Natural -> Natural
addMod :: Natural -> Natural -> Natural -> Natural
addMod Natural
m Natural
x Natural
y = let z :: Natural
z = Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
y in if Natural
z Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
>= Natural
m then Natural
z Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
m else Natural
z

subMod :: Natural -> Natural -> Natural -> Natural
subMod :: Natural -> Natural -> Natural -> Natural
subMod Natural
m Natural
x Natural
y = if Natural
x Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
>= Natural
y then Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
y else Natural
m Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
y

negateMod :: Natural -> Natural -> Natural
negateMod :: Natural -> Natural -> Natural
negateMod !Natural
_ Natural
0 = Natural
0
negateMod Natural
m Natural
x = Natural
m Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
x

mulMod :: Natural -> Natural -> Natural -> Natural
mulMod :: Natural -> Natural -> Natural -> Natural
mulMod Natural
m Natural
x Natural
y = (Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
* Natural
y) Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`Prelude.rem` Natural
m

instance KnownNat m => Num (Mod m) where
  mx :: Mod m
mx@(Mod !Natural
x) + :: Mod m -> Mod m -> Mod m
+ (Mod !Natural
y) = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural -> Natural
addMod (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx) Natural
x Natural
y
  {-# INLINE (+) #-}
  mx :: Mod m
mx@(Mod !Natural
x) - :: Mod m -> Mod m -> Mod m
- (Mod !Natural
y) = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural -> Natural
subMod (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx) Natural
x Natural
y
  {-# INLINE (-) #-}
  negate :: Mod m -> Mod m
negate mx :: Mod m
mx@(Mod !Natural
x) = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural
negateMod (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx) Natural
x
  {-# INLINE negate #-}
  mx :: Mod m
mx@(Mod !Natural
x) * :: Mod m -> Mod m -> Mod m
* (Mod !Natural
y) = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural -> Natural
mulMod (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx) Natural
x Natural
y
  {-# INLINE (*) #-}
  abs :: Mod m -> Mod m
abs = Mod m -> Mod m
forall a. a -> a
id
  {-# INLINE abs #-}
  signum :: Mod m -> Mod m
signum = Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const Mod m
x
    where
      x :: Mod m
x = if Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
x Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
1 then Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
1 else Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
0
  {-# INLINE signum #-}
  fromInteger :: Integer -> Mod m
fromInteger Integer
x = Mod m
mx
    where
      mx :: Mod m
mx = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx)
  {-# INLINE fromInteger #-}

#ifdef MIN_VERSION_semirings

instance KnownNat m => Semiring (Mod m) where
  plus :: Mod m -> Mod m -> Mod m
plus  = Mod m -> Mod m -> Mod m
forall a. Num a => a -> a -> a
(+)
  {-# INLINE plus #-}
  times :: Mod m -> Mod m -> Mod m
times = Mod m -> Mod m -> Mod m
forall a. Num a => a -> a -> a
(*)
  {-# INLINE times #-}
  zero :: Mod m
zero  = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
0
  {-# INLINE zero #-}
  one :: Mod m
one   = Mod m
mx
    where
      mx :: Mod m
mx = if Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
1 then Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
1 else Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod Natural
0
  {-# INLINE one #-}
  fromNatural :: Natural -> Mod m
fromNatural Natural
x = Mod m
mx
    where
      mx :: Mod m
mx = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural
x Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`mod` Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx
  {-# INLINE fromNatural #-}

instance KnownNat m => Ring (Mod m) where
  negate :: Mod m -> Mod m
negate = Mod m -> Mod m
forall a. Num a => a -> a
Prelude.negate
  {-# INLINE negate #-}

-- | See the warning about division above.
instance KnownNat m => GcdDomain (Mod m) where
  divide :: Mod m -> Mod m -> Maybe (Mod m)
divide Mod m
x Mod m
y = Mod m -> Maybe (Mod m)
forall a. a -> Maybe a
Just (Mod m
x Mod m -> Mod m -> Mod m
forall a. Fractional a => a -> a -> a
/ Mod m
y)
  gcd :: Mod m -> Mod m -> Mod m
gcd        = (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const ((Mod m -> Mod m) -> Mod m -> Mod m -> Mod m)
-> (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. (a -> b) -> a -> b
$ Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const Mod m
1
  lcm :: Mod m -> Mod m -> Mod m
lcm        = (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const ((Mod m -> Mod m) -> Mod m -> Mod m -> Mod m)
-> (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. (a -> b) -> a -> b
$ Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const Mod m
1
  coprime :: Mod m -> Mod m -> Bool
coprime    = (Mod m -> Bool) -> Mod m -> Mod m -> Bool
forall a b. a -> b -> a
const ((Mod m -> Bool) -> Mod m -> Mod m -> Bool)
-> (Mod m -> Bool) -> Mod m -> Mod m -> Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Mod m -> Bool
forall a b. a -> b -> a
const Bool
True

-- | See the warning about division above.
instance KnownNat m => Euclidean (Mod m) where
  degree :: Mod m -> Natural
degree      = Natural -> Mod m -> Natural
forall a b. a -> b -> a
const Natural
0
  quotRem :: Mod m -> Mod m -> (Mod m, Mod m)
quotRem Mod m
x Mod m
y = (Mod m
x Mod m -> Mod m -> Mod m
forall a. Fractional a => a -> a -> a
/ Mod m
y, Mod m
0)
  quot :: Mod m -> Mod m -> Mod m
quot        = Mod m -> Mod m -> Mod m
forall a. Fractional a => a -> a -> a
(/)
  rem :: Mod m -> Mod m -> Mod m
rem         = (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const ((Mod m -> Mod m) -> Mod m -> Mod m -> Mod m)
-> (Mod m -> Mod m) -> Mod m -> Mod m -> Mod m
forall a b. (a -> b) -> a -> b
$ Mod m -> Mod m -> Mod m
forall a b. a -> b -> a
const Mod m
0

-- | See the warning about division above.
instance KnownNat m => Field (Mod m)

#endif

-- | See the warning about division above.
instance KnownNat m => Fractional (Mod m) where
  fromRational :: Rational -> Mod m
fromRational Rational
r = case Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r of
    Integer
1   -> Mod m
num
    Integer
den -> Mod m
num Mod m -> Mod m -> Mod m
forall a. Fractional a => a -> a -> a
/ Integer -> Mod m
forall a. Num a => Integer -> a
fromInteger Integer
den
    where
      num :: Mod m
num = Integer -> Mod m
forall a. Num a => Integer -> a
fromInteger (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r)
  {-# INLINE fromRational #-}
  recip :: Mod m -> Mod m
recip Mod m
mx = case Mod m -> Maybe (Mod m)
forall (m :: Nat). KnownNat m => Mod m -> Maybe (Mod m)
invertMod Mod m
mx of
    Maybe (Mod m)
Nothing -> ArithException -> Mod m
forall a e. Exception e => e -> a
throw ArithException
DivideByZero
    Just Mod m
y  -> Mod m
y
  {-# INLINE recip #-}

-- | If an argument is
-- <https://en.wikipedia.org/wiki/Coprime_integers coprime>
-- with the modulo, return its modular inverse.
-- Otherwise return 'Nothing'.
--
-- >>> :set -XDataKinds
-- >>> invertMod 3 :: Mod 10 -- 3 * 7 = 21 ≡ 1 (mod 10)
-- Just (7 `modulo` 10)
-- >>> invertMod 4 :: Mod 10 -- 4 and 10 are not coprime
-- Nothing
invertMod :: KnownNat m => Mod m -> Maybe (Mod m)
invertMod :: Mod m -> Maybe (Mod m)
invertMod Mod m
mx
  = if Integer
y Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0
    then Maybe (Mod m)
forall a. Maybe a
Nothing
    else Mod m -> Maybe (Mod m)
forall a. a -> Maybe a
Just (Mod m -> Maybe (Mod m)) -> Mod m -> Maybe (Mod m)
forall a b. (a -> b) -> a -> b
$ Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
y
  where
    y :: Integer
y = Integer -> Integer -> Integer
recipModInteger (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
mx)) (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx))
{-# INLINABLE invertMod #-}

recipModInteger :: Integer -> Integer -> Integer
recipModInteger :: Integer -> Integer -> Integer
recipModInteger Integer
x Integer
m = case Integer -> Integer -> (Integer, Integer)
gcdExt Integer
x Integer
m of
  (Integer
1, Integer
s) -> Integer
s Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
  (Integer, Integer)
_ -> -Integer
1

gcdExt :: Integer -> Integer -> (Integer, Integer)
gcdExt :: Integer -> Integer -> (Integer, Integer)
gcdExt = Integer -> Integer -> Integer -> Integer -> (Integer, Integer)
forall t. Integral t => t -> t -> t -> t -> (t, t)
go Integer
1 Integer
0
  where
    go :: t -> t -> t -> t -> (t, t)
go t
s !t
_ t
r t
0 = (t
r, t
s)
    go t
s t
s' t
r t
r' = case t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
Prelude.quotRem t
r t
r' of
      (t
q, t
r'') -> t -> t -> t -> t -> (t, t)
go t
s' (t
s t -> t -> t
forall a. Num a => a -> a -> a
- t
q t -> t -> t
forall a. Num a => a -> a -> a
* t
s') t
r' t
r''

-- | Drop-in replacement for 'Prelude.^' with much better performance.
-- Negative powers are allowed, but may throw 'DivideByZero', if an argument
-- is not <https://en.wikipedia.org/wiki/Coprime_integers coprime> with the modulo.
--
-- Building with @-O@ triggers a rewrite rule 'Prelude.^' = '^%'.
--
-- >>> :set -XDataKinds
-- >>> 3 ^% 4 :: Mod 10    -- 3 ^ 4 = 81 ≡ 1 (mod 10)
-- (1 `modulo` 10)
-- >>> 3 ^% (-1) :: Mod 10 -- 3 * 7 = 21 ≡ 1 (mod 10)
-- (7 `modulo` 10)
-- >>> 4 ^% (-1) :: Mod 10 -- 4 and 10 are not coprime
-- (*** Exception: divide by zero
(^%) :: (KnownNat m, Integral a) => Mod m -> a -> Mod m
Mod m
mx ^% :: Mod m -> a -> Mod m
^% a
a
  | a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = case Mod m -> Maybe (Mod m)
forall (m :: Nat). KnownNat m => Mod m -> Maybe (Mod m)
invertMod Mod m
mx of
    Maybe (Mod m)
Nothing ->  ArithException -> Mod m
forall a e. Exception e => e -> a
throw ArithException
DivideByZero
    Just Mod m
my ->  Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural -> Natural
powModNatural (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
my) (a -> Natural
fromIntegral' (-a
a)) (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx)
  | Bool
otherwise = Natural -> Mod m
forall (m :: Nat). Natural -> Mod m
Mod (Natural -> Mod m) -> Natural -> Mod m
forall a b. (a -> b) -> a -> b
$ Natural -> Natural -> Natural -> Natural
powModNatural (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
mx) (a -> Natural
fromIntegral' a
a)    (Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx)
  where
#if __GLASGOW_HASKELL__ == 900 && __GLASGOW_HASKELL_PATCHLEVEL1__ == 1
    -- Cannot use fromIntegral because of https://gitlab.haskell.org/ghc/ghc/-/issues/19411
    fromIntegral' = fromInteger . toInteger
#else
    fromIntegral' :: a -> Natural
fromIntegral' = a -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral
#endif
{-# INLINABLE [1] (^%) #-}

{-# SPECIALISE [1] (^%) ::
  KnownNat m => Mod m -> Integer -> Mod m,
  KnownNat m => Mod m -> Natural -> Mod m,
  KnownNat m => Mod m -> Int     -> Mod m,
  KnownNat m => Mod m -> Word    -> Mod m #-}

{-# RULES
"powMod"               forall (x :: KnownNat m => Mod m) p. x ^ p = x ^% p

"powMod/2/Integer"     forall x. x ^% (2 :: Integer) = let u = x in u*u
"powMod/3/Integer"     forall x. x ^% (3 :: Integer) = let u = x in u*u*u
"powMod/2/Int"         forall x. x ^% (2 :: Int)     = let u = x in u*u
"powMod/3/Int"         forall x. x ^% (3 :: Int)     = let u = x in u*u*u
"powMod/2/Word"        forall x. x ^% (2 :: Word)    = let u = x in u*u
"powMod/3/Word"        forall x. x ^% (3 :: Word)    = let u = x in u*u*u #-}

infixr 8 ^%