-- |
-- 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 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
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 =
    forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx forall a. Eq a => a -> a -> Bool
== forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
my Bool -> Bool -> Bool
&& forall (m :: Natural). Mod m -> Natural
unMod Mod m
mx forall a. Eq a => a -> a -> Bool
== forall (m :: Natural). Mod m -> Natural
unMod Mod m
my
  InfMod Rational
rx  == InfMod Rational
ry  = Rational
rx 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 =
    forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx forall a. Ord a => a -> a -> Ordering
`compare` forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
my forall a. Semigroup a => a -> a -> a
<> forall (m :: Natural). Mod m -> Natural
unMod Mod m
mx forall a. Ord a => a -> a -> Ordering
`compare` forall (m :: Natural). 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 forall a. Ord a => a -> a -> Ordering
`compare` Rational
ry

instance Show SomeMod where
  show :: SomeMod -> String
show = \case
    SomeMod Mod m
m -> forall a. Show a => a -> String
show Mod m
m
    InfMod  Rational
r -> 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) -> forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (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 :: Natural). KnownNat k => Mod k -> Mod k)
-> (Rational -> Rational) -> SomeMod -> SomeMod
liftUnOp forall (k :: Natural). KnownNat k => Mod k -> Mod k
fm Rational -> Rational
fr = \case
  SomeMod Mod m
m -> forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (forall (k :: Natural). 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 (m :: Natural) (n :: Natural).
(KnownNat m, KnownNat n) =>
(forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k)
-> Mod m -> Mod n -> SomeMod
liftBinOpMod forall (k :: Natural). 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) ->
    forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural
x forall a. Integral a => a -> a -> a
`mod` Natural
m) forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
`f` forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural
y forall a. Integral a => a -> a -> a
`mod` Natural
m) :: Mod t)
  where
    x :: Natural
x = forall (m :: Natural). Mod m -> Natural
unMod Mod m
mx
    y :: Natural
y = forall (m :: Natural). Mod m -> Natural
unMod Mod n
my
    m :: Natural
m = forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
mx forall a. Integral a => a -> a -> a
`Prelude.gcd` forall (n :: Natural) (proxy :: Natural -> *).
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 :: Natural). KnownNat k => Mod k -> Mod k -> Mod k)
-> (Rational -> Rational -> Rational)
-> SomeMod
-> SomeMod
-> SomeMod
liftBinOp forall (k :: Natural). 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 :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
fm Rational -> Rational -> Rational
_ (InfMod Rational
rx)  (SomeMod Mod m
my) = forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (forall a. Fractional a => Rational -> a
fromRational Rational
rx forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
`fm` Mod m
my)
liftBinOp forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
fm Rational -> Rational -> Rational
_ (SomeMod Mod m
mx) (InfMod Rational
ry)  = forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
mx forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
`fm` forall a. Fractional a => Rational -> a
fromRational Rational
ry)
liftBinOp forall (k :: Natural). 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 (forall {k} (t :: k). Proxy t
Proxy :: Proxy m) forall (a :: Natural) (b :: Natural) (proxy1 :: Natural -> *)
       (proxy2 :: Natural -> *).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe (a :~: b)
`sameNat` (forall {k} (t :: k). Proxy t
Proxy :: Proxy n) of
    Maybe (m :~: m)
Nothing   -> forall (m :: Natural) (n :: Natural).
(KnownNat m, KnownNat n) =>
(forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k)
-> Mod m -> Mod n -> SomeMod
liftBinOpMod forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
fm Mod m
mx Mod m
my
    Just m :~: m
Refl -> forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
mx forall (k :: Natural). KnownNat k => Mod k -> Mod k -> Mod k
`fm` Mod m
my)

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

instance Semiring SomeMod where
  plus :: SomeMod -> SomeMod -> SomeMod
plus  = forall a. Num a => a -> a -> a
(+)
  times :: SomeMod -> SomeMod -> SomeMod
times = 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 = forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance Ring SomeMod where
  negate :: SomeMod -> SomeMod
negate = 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 -> 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 = forall a. a -> Maybe a
Just (SomeMod
x forall a. Fractional a => a -> a -> a
/ SomeMod
y)
  gcd :: SomeMod -> SomeMod -> SomeMod
gcd        = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const SomeMod
1
  lcm :: SomeMod -> SomeMod -> SomeMod
lcm        = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const SomeMod
1
  coprime :: SomeMod -> SomeMod -> Bool
coprime    = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const Bool
True

-- | See the warning about division above.
instance Euclidean SomeMod where
  degree :: SomeMod -> Natural
degree      = forall a b. a -> b -> a
const Natural
0
  quotRem :: SomeMod -> SomeMod -> (SomeMod, SomeMod)
quotRem SomeMod
x SomeMod
y = (SomeMod
x forall a. Fractional a => a -> a -> a
/ SomeMod
y, SomeMod
0)
  quot :: SomeMod -> SomeMod -> SomeMod
quot        = forall a. Fractional a => a -> a -> a
(/)
  rem :: SomeMod -> SomeMod -> SomeMod
rem         = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ 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 -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (forall (m :: Natural). KnownNat m => Mod m -> Maybe (Mod m)
invertMod Mod m
m)
  InfMod  Rational
r -> forall a. a -> Maybe a
Just (Rational -> SomeMod
InfMod (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 :: forall a. Integral a => SomeMod -> a -> SomeMod
powSomeMod (SomeMod Mod m
m) a
a = forall (m :: Natural). KnownNat m => Mod m -> SomeMod
SomeMod (Mod m
m forall (m :: Natural) a.
(KnownNat m, Integral a) =>
Mod m -> a -> Mod m
^% a
a)
powSomeMod (InfMod  Rational
r) a
a = Rational -> SomeMod
InfMod  (Rational
r 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 #-}