-- |
-- Module:      Math.NumberTheory.Moduli.SomeMod
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Safe modular arithmetic with modulo on type level.
--

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.Moduli.SomeMod
  ( SomeMod(..)
  , modulo
  , invertSomeMod
  , powSomeMod
  ) where

import Data.Euclidean (GcdDomain(..), Euclidean(..), Field)
import Data.Mod
import Data.Proxy
#if __GLASGOW_HASKELL__ < 803
import Data.Semigroup
#endif
import Data.Semiring (Semiring(..), Ring(..))
import Data.Type.Equality
import GHC.TypeNats (KnownNat, SomeNat(..), sameNat, natVal, someNatVal)
import Numeric.Natural

-- | This type represents residues with unknown modulo and rational numbers.
-- One can freely combine them in arithmetic expressions, but each operation
-- will spend time on modulo's recalculation:
--
-- >>> 2 `modulo` 10 + 4 `modulo` 15
-- (1 `modulo` 5)
-- >>> (2 `modulo` 10) * (4 `modulo` 15)
-- (3 `modulo` 5)
-- >>> import Data.Ratio
-- >>> 2 `modulo` 10 + fromRational (3 % 7)
-- (1 `modulo` 10)
-- >>> 2 `modulo` 10 * fromRational (3 % 7)
-- (8 `modulo` 10)
--
-- If performance is crucial, it is recommended to extract @Mod m@ for further processing
-- by pattern matching. E. g.,
--
-- > case modulo n m of
-- >   SomeMod k -> process k -- Here k has type Mod m
-- >   InfMod{}  -> error "impossible"
data SomeMod where
  SomeMod :: KnownNat m => Mod m -> SomeMod
  InfMod  :: Rational -> SomeMod

instance Eq SomeMod where
  SomeMod Mod m
mx == :: SomeMod -> SomeMod -> Bool
== SomeMod Mod m
my =
    Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
my Bool -> Bool -> Bool
&& Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
mx Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
my
  InfMod Rational
rx  == InfMod Rational
ry  = Rational
rx Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
ry
  SomeMod
_          == SomeMod
_          = Bool
False

instance Ord SomeMod where
  SomeMod Mod m
mx compare :: SomeMod -> SomeMod -> Ordering
`compare` SomeMod Mod m
my =
    Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx Natural -> Natural -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
my Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
mx Natural -> Natural -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
my
  SomeMod{} `compare` InfMod{} = Ordering
LT
  InfMod{} `compare` SomeMod{} = Ordering
GT
  InfMod Rational
rx `compare` InfMod Rational
ry = Rational
rx Rational -> Rational -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Rational
ry

instance Show SomeMod where
  show :: SomeMod -> String
show = \case
    SomeMod Mod m
m -> Mod m -> String
forall a. Show a => a -> String
show Mod m
m
    InfMod  Rational
r -> Rational -> String
forall a. Show a => a -> String
show Rational
r

-- | Create modular value by representative of residue class and modulo.
-- One can use the result either directly (via functions from 'Num' and 'Fractional'),
-- or deconstruct it by pattern matching. Note that 'modulo' never returns 'InfMod'.
modulo :: Integer -> Natural -> SomeMod
modulo :: Integer -> Natural -> SomeMod
modulo Integer
n Natural
m = case Natural -> SomeNat
someNatVal Natural
m of
  SomeNat (Proxy n
_ :: Proxy t) -> Mod n -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Integer -> Mod n
forall a. Num a => Integer -> a
fromInteger Integer
n :: Mod t)
{-# INLINABLE modulo #-}
infixl 7 `modulo`

liftUnOp
  :: (forall k. KnownNat k => Mod k -> Mod k)
  -> (Rational -> Rational)
  -> SomeMod
  -> SomeMod
liftUnOp :: (forall (k :: Nat). KnownNat k => Mod k -> Mod k)
-> (Rational -> Rational) -> SomeMod -> SomeMod
liftUnOp forall (k :: Nat). KnownNat k => Mod k -> Mod k
fm Rational -> Rational
fr = \case
  SomeMod Mod m
m -> Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m -> Mod m
forall (k :: Nat). KnownNat k => Mod k -> Mod k
fm Mod m
m)
  InfMod  Rational
r -> Rational -> SomeMod
InfMod  (Rational -> Rational
fr Rational
r)
{-# INLINEABLE liftUnOp #-}

liftBinOpMod
  :: (KnownNat m, KnownNat n)
  => (forall k. KnownNat k => Mod k -> Mod k -> Mod k)
  -> Mod m
  -> Mod n
  -> SomeMod
liftBinOpMod :: (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> Mod m -> Mod n -> SomeMod
liftBinOpMod forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
f Mod m
mx Mod n
my = case Natural -> SomeNat
someNatVal Natural
m of
  SomeNat (Proxy n
_ :: Proxy t) ->
    Mod n -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Natural -> Mod n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural
x Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`mod` Natural
m) Mod n -> Mod n -> Mod n
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
`f` Natural -> Mod n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural
y Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`mod` Natural
m) :: Mod t)
  where
    x :: Natural
x = Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
mx
    y :: Natural
y = Mod n -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod n
my
    m :: Natural
m = Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`Prelude.gcd` Mod n -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod n
my

liftBinOp
  :: (forall k. KnownNat k => Mod k -> Mod k -> Mod k)
  -> (Rational -> Rational -> Rational)
  -> SomeMod
  -> SomeMod
  -> SomeMod
liftBinOp :: (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> (Rational -> Rational -> Rational)
-> SomeMod
-> SomeMod
-> SomeMod
liftBinOp forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
_ Rational -> Rational -> Rational
fr (InfMod Rational
rx)  (InfMod Rational
ry)  = Rational -> SomeMod
InfMod  (Rational
rx Rational -> Rational -> Rational
`fr` Rational
ry)
liftBinOp forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
fm Rational -> Rational -> Rational
_ (InfMod Rational
rx)  (SomeMod Mod m
my) = Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Rational -> Mod m
forall a. Fractional a => Rational -> a
fromRational Rational
rx Mod m -> Mod m -> Mod m
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
`fm` Mod m
my)
liftBinOp forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
fm Rational -> Rational -> Rational
_ (SomeMod Mod m
mx) (InfMod Rational
ry)  = Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
mx Mod m -> Mod m -> Mod m
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
`fm` Rational -> Mod m
forall a. Fractional a => Rational -> a
fromRational Rational
ry)
liftBinOp forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
fm Rational -> Rational -> Rational
_ (SomeMod (Mod m
mx :: Mod m)) (SomeMod (Mod m
my :: Mod n))
  = case (Proxy m
forall k (t :: k). Proxy t
Proxy :: Proxy m) Proxy m -> Proxy m -> Maybe (m :~: m)
forall (a :: Nat) (b :: Nat).
(KnownNat a, KnownNat b) =>
Proxy a -> Proxy b -> Maybe (a :~: b)
`sameNat` (Proxy m
forall k (t :: k). Proxy t
Proxy :: Proxy n) of
    Maybe (m :~: m)
Nothing   -> (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> Mod m -> Mod m -> SomeMod
forall (m :: Nat) (n :: Nat).
(KnownNat m, KnownNat n) =>
(forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> Mod m -> Mod n -> SomeMod
liftBinOpMod forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
fm Mod m
mx Mod m
my
    Just m :~: m
Refl -> Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
mx Mod m -> Mod m -> Mod m
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
`fm` Mod m
Mod m
my)

instance Num SomeMod where
  + :: SomeMod -> SomeMod -> SomeMod
(+)    = (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> (Rational -> Rational -> Rational)
-> SomeMod
-> SomeMod
-> SomeMod
liftBinOp forall a. Num a => a -> a -> a
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
(+) Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
(+)
  (-)    = (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> (Rational -> Rational -> Rational)
-> SomeMod
-> SomeMod
-> SomeMod
liftBinOp (-) (-)
  negate :: SomeMod -> SomeMod
negate = (forall (k :: Nat). KnownNat k => Mod k -> Mod k)
-> (Rational -> Rational) -> SomeMod -> SomeMod
liftUnOp forall a. Num a => a -> a
forall (k :: Nat). KnownNat k => Mod k -> Mod k
Prelude.negate Rational -> Rational
forall a. Num a => a -> a
Prelude.negate
  {-# INLINE negate #-}
  * :: SomeMod -> SomeMod -> SomeMod
(*)    = (forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k)
-> (Rational -> Rational -> Rational)
-> SomeMod
-> SomeMod
-> SomeMod
liftBinOp forall a. Num a => a -> a -> a
forall (k :: Nat). KnownNat k => Mod k -> Mod k -> Mod k
(*) Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
(*)
  abs :: SomeMod -> SomeMod
abs    = SomeMod -> SomeMod
forall a. a -> a
id
  {-# INLINE abs #-}
  signum :: SomeMod -> SomeMod
signum = SomeMod -> SomeMod -> SomeMod
forall a b. a -> b -> a
const SomeMod
1
  {-# INLINE signum #-}
  fromInteger :: Integer -> SomeMod
fromInteger = Rational -> SomeMod
InfMod (Rational -> SomeMod)
-> (Integer -> Rational) -> Integer -> SomeMod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Rational
forall a. Num a => Integer -> a
fromInteger
  {-# INLINE fromInteger #-}

instance Semiring SomeMod where
  plus :: SomeMod -> SomeMod -> SomeMod
plus  = SomeMod -> SomeMod -> SomeMod
forall a. Num a => a -> a -> a
(+)
  times :: SomeMod -> SomeMod -> SomeMod
times = SomeMod -> SomeMod -> SomeMod
forall a. Num a => a -> a -> a
(*)
  zero :: SomeMod
zero  = Rational -> SomeMod
InfMod Rational
0
  one :: SomeMod
one   = Rational -> SomeMod
InfMod Rational
1
  fromNatural :: Natural -> SomeMod
fromNatural = Natural -> SomeMod
forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance Ring SomeMod where
  negate :: SomeMod -> SomeMod
negate = SomeMod -> SomeMod
forall a. Num a => a -> a
Prelude.negate

-- | Beware that division by residue, which is not coprime with the modulo,
-- will result in runtime error. Consider using 'invertSomeMod' instead.
instance Fractional SomeMod where
  fromRational :: Rational -> SomeMod
fromRational = Rational -> SomeMod
InfMod
  {-# INLINE fromRational #-}
  recip :: SomeMod -> SomeMod
recip SomeMod
x = case SomeMod -> Maybe SomeMod
invertSomeMod SomeMod
x of
    Maybe SomeMod
Nothing -> String -> SomeMod
forall a. HasCallStack => String -> a
error String
"recip{SomeMod}: residue is not coprime with modulo"
    Just SomeMod
y  -> SomeMod
y

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

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

-- | See the warning about division above.
instance Field SomeMod

-- | Computes the inverse value, if it exists.
--
-- >>> invertSomeMod (3 `modulo` 10) -- because 3 * 7 = 1 :: Mod 10
-- Just (7 `modulo` 10)
-- >>> invertSomeMod (4 `modulo` 10)
-- Nothing
-- >>> import Data.Ratio
-- >>> invertSomeMod (fromRational (2 % 5))
-- Just 5 % 2
invertSomeMod :: SomeMod -> Maybe SomeMod
invertSomeMod :: SomeMod -> Maybe SomeMod
invertSomeMod = \case
  SomeMod Mod m
m -> (Mod m -> SomeMod) -> Maybe (Mod m) -> Maybe SomeMod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m -> Maybe (Mod m)
forall (m :: Nat). KnownNat m => Mod m -> Maybe (Mod m)
invertMod Mod m
m)
  InfMod  Rational
r -> SomeMod -> Maybe SomeMod
forall a. a -> Maybe a
Just (Rational -> SomeMod
InfMod (Rational -> Rational
forall a. Fractional a => a -> a
recip Rational
r))
{-# INLINABLE [1] invertSomeMod #-}

{-# SPECIALISE [1] powSomeMod ::
  SomeMod -> Integer -> SomeMod,
  SomeMod -> Natural -> SomeMod,
  SomeMod -> Int     -> SomeMod,
  SomeMod -> Word    -> SomeMod #-}

-- | Drop-in replacement for 'Prelude.^', with much better performance.
-- When -O is enabled, there is a rewrite rule, which specialises 'Prelude.^' to 'powSomeMod'.
--
-- >>> powSomeMod (3 `modulo` 10) 4
-- (1 `modulo` 10)
powSomeMod :: Integral a => SomeMod -> a -> SomeMod
powSomeMod :: SomeMod -> a -> SomeMod
powSomeMod (SomeMod Mod m
m) a
a = Mod m -> SomeMod
forall (m :: Nat). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
m Mod m -> a -> Mod m
forall (m :: Nat) a.
(KnownNat m, Integral a) =>
Mod m -> a -> Mod m
^% a
a)
powSomeMod (InfMod  Rational
r) a
a = Rational -> SomeMod
InfMod  (Rational
r Rational -> a -> Rational
forall a b. (Num a, Integral b) => a -> b -> a
^  a
a)
{-# INLINABLE [1] powSomeMod #-}

{-# RULES "^%SomeMod" forall x p. x ^ p = powSomeMod x p #-}