-- |
-- Module:      Math.NumberTheory.Moduli.Chinese
-- Copyright:   (c) 2011 Daniel Fischer, 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Chinese remainder theorem
--

{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.Moduli.Chinese
  ( -- * Safe interface
    chinese
  , chineseSomeMod
  ) where

import Prelude hiding ((^), (+), (-), (*), rem, mod, quot, gcd, lcm)

import Data.Euclidean
import Data.Mod
import Data.Ratio
import Data.Semiring (Semiring(..), (+), (-), (*), Ring)
import GHC.TypeNats (KnownNat, natVal)

import Math.NumberTheory.Moduli.SomeMod

-- | 'chinese' @(n1, m1)@ @(n2, m2)@ returns @(n, lcm m1 m2)@ such that
-- @n \`mod\` m1 == n1@ and @n \`mod\` m2 == n2@, if exists.
-- Moduli @m1@ and @m2@ are allowed to have common factors.
--
-- >>> chinese (1, 2) (2, 3)
-- Just (-1, 6)
-- >>> chinese (3, 4) (5, 6)
-- Just (-1, 12)
-- >>> chinese (3, 4) (2, 6)
-- Nothing
chinese :: forall a. (Eq a, Ring a, Euclidean a) => (a, a) -> (a, a) -> Maybe (a, a)
chinese :: (a, a) -> (a, a) -> Maybe (a, a)
chinese (a
n1, a
m1) (a
n2, a
m2)
  | a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Semiring a => a
one
  = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just ((a
v a -> a -> a
forall a. Semiring a => a -> a -> a
* a
m2 a -> a -> a
forall a. Semiring a => a -> a -> a
* a
n1 a -> a -> a
forall a. Semiring a => a -> a -> a
+ a
u a -> a -> a
forall a. Semiring a => a -> a -> a
* a
m1 a -> a -> a
forall a. Semiring a => a -> a -> a
* a
n2) a -> a -> a
forall a. Euclidean a => a -> a -> a
`rem` a
m, a
m)
  | (a
n1 a -> a -> a
forall a. Ring a => a -> a -> a
- a
n2) a -> a -> a
forall a. Euclidean a => a -> a -> a
`rem` a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Semiring a => a
zero
  = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just ((a
v a -> a -> a
forall a. Semiring a => a -> a -> a
* (a
m2 a -> a -> a
forall a. Euclidean a => a -> a -> a
`quot` a
d) a -> a -> a
forall a. Semiring a => a -> a -> a
* a
n1 a -> a -> a
forall a. Semiring a => a -> a -> a
+ a
u a -> a -> a
forall a. Semiring a => a -> a -> a
* (a
m1 a -> a -> a
forall a. Euclidean a => a -> a -> a
`quot` a
d) a -> a -> a
forall a. Semiring a => a -> a -> a
* a
n2) a -> a -> a
forall a. Euclidean a => a -> a -> a
`rem` a
m, a
m)
  | Bool
otherwise
  = Maybe (a, a)
forall a. Maybe a
Nothing
  where
    (a
d, a
u, a
v) = a -> a -> (a, a, a)
forall a. (Eq a, Ring a, Euclidean a) => a -> a -> (a, a, a)
extendedGCD a
m1 a
m2
    m :: a
m = if a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. Semiring a => a
one then a
m1 a -> a -> a
forall a. Semiring a => a -> a -> a
* a
m2 else (a
m1 a -> a -> a
forall a. Euclidean a => a -> a -> a
`quot` a
d) a -> a -> a
forall a. Semiring a => a -> a -> a
* a
m2

{-# SPECIALISE chinese :: (Int, Int) -> (Int, Int) -> Maybe (Int, Int) #-}
{-# SPECIALISE chinese :: (Word, Word) -> (Word, Word) -> Maybe (Word, Word) #-}
{-# SPECIALISE chinese :: (Integer, Integer) -> (Integer, Integer) -> Maybe (Integer, Integer) #-}

isCompatible :: KnownNat m => Mod m -> Rational -> Bool
isCompatible :: Mod m -> Rational -> Bool
isCompatible Mod m
n Rational
r = case Mod m -> Maybe (Mod m)
forall (m :: Nat). KnownNat m => Mod m -> Maybe (Mod m)
invertMod (Integer -> Mod m
forall a. Num a => Integer -> a
fromInteger (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r)) of
  Maybe (Mod m)
Nothing -> Bool
False
  Just Mod m
r' -> Mod m
r' Mod m -> Mod m -> Mod m
forall a. Semiring a => a -> a -> a
* Integer -> Mod m
forall a. Num a => Integer -> a
fromInteger (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r) Mod m -> Mod m -> Bool
forall a. Eq a => a -> a -> Bool
== Mod m
n

-- | Same as 'chinese', but operates on residues.
--
-- >>> :set -XDataKinds
-- >>> import Data.Mod
-- >>> (1 `modulo` 2) `chineseSomeMod` (2 `modulo` 3)
-- Just (5 `modulo` 6)
-- >>> (3 `modulo` 4) `chineseSomeMod` (5 `modulo` 6)
-- Just (11 `modulo` 12)
-- >>> (3 `modulo` 4) `chineseSomeMod` (2 `modulo` 6)
-- Nothing
chineseSomeMod :: SomeMod -> SomeMod -> Maybe SomeMod
chineseSomeMod :: SomeMod -> SomeMod -> Maybe SomeMod
chineseSomeMod (SomeMod Mod m
n1) (SomeMod Mod m
n2)
  = (\(Integer
n, Integer
m) -> Integer
n Integer -> Natural -> SomeMod
`modulo` Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
m) ((Integer, Integer) -> SomeMod)
-> Maybe (Integer, Integer) -> Maybe SomeMod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Integer, Integer)
-> (Integer, Integer) -> Maybe (Integer, Integer)
forall a.
(Eq a, Ring a, Euclidean a) =>
(a, a) -> (a, a) -> Maybe (a, a)
chinese
    (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
n1, Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
n1)
    (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
n2, Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Mod m -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod m
n2)
chineseSomeMod (SomeMod Mod m
n) (InfMod Rational
r)
  | Mod m -> Rational -> Bool
forall (m :: Nat). KnownNat m => Mod m -> Rational -> Bool
isCompatible Mod m
n Rational
r = SomeMod -> Maybe SomeMod
forall a. a -> Maybe a
Just (SomeMod -> Maybe SomeMod) -> SomeMod -> Maybe SomeMod
forall a b. (a -> b) -> a -> b
$ Rational -> SomeMod
InfMod Rational
r
  | Bool
otherwise        = Maybe SomeMod
forall a. Maybe a
Nothing
chineseSomeMod (InfMod Rational
r) (SomeMod Mod m
n)
  | Mod m -> Rational -> Bool
forall (m :: Nat). KnownNat m => Mod m -> Rational -> Bool
isCompatible Mod m
n Rational
r = SomeMod -> Maybe SomeMod
forall a. a -> Maybe a
Just (SomeMod -> Maybe SomeMod) -> SomeMod -> Maybe SomeMod
forall a b. (a -> b) -> a -> b
$ Rational -> SomeMod
InfMod Rational
r
  | Bool
otherwise        = Maybe SomeMod
forall a. Maybe a
Nothing
chineseSomeMod (InfMod Rational
r1) (InfMod Rational
r2)
  | Rational
r1 Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
r2  = SomeMod -> Maybe SomeMod
forall a. a -> Maybe a
Just (SomeMod -> Maybe SomeMod) -> SomeMod -> Maybe SomeMod
forall a b. (a -> b) -> a -> b
$ Rational -> SomeMod
InfMod Rational
r1
  | Bool
otherwise = Maybe SomeMod
forall a. Maybe a
Nothing

-------------------------------------------------------------------------------
-- Utils

extendedGCD :: (Eq a, Ring a, Euclidean a) => a -> a -> (a, a, a)
extendedGCD :: a -> a -> (a, a, a)
extendedGCD a
a a
b = (a
g, a
s, a
t)
  where
    (a
g, a
s) = a -> a -> (a, a)
forall a. (Eq a, Euclidean a, Ring a) => a -> a -> (a, a)
gcdExt a
a a
b
    t :: a
t = (a
g a -> a -> a
forall a. Ring a => a -> a -> a
- a
a a -> a -> a
forall a. Semiring a => a -> a -> a
* a
s) a -> a -> a
forall a. Euclidean a => a -> a -> a
`quot` a
b