{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE CPP                  #-}
{-# LANGUAGE ConstraintKinds      #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE ExplicitNamespaces   #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}


-- |
-- Types for working with integers modulo some constant.
module Data.Modular (
  -- $doc

  -- * Preliminaries
  -- $setup

  -- * Modular arithmetic
  Mod, Modulus,
  unMod, toMod, toMod',
  inv, type (/)(), ℤ,
  modVal, SomeMod, someModVal
) where

import           Control.Arrow          (first)

import           Data.Proxy             (Proxy (..))
import           Data.Ratio             (denominator, numerator, (%))

import           Text.Printf            (printf)

#if MIN_VERSION_base(4,11,0)
import           GHC.TypeLits           hiding (Mod)
#else
import           GHC.TypeLits
#endif

#if !MIN_VERSION_base(4,16,0)
import           Data.Type.Equality     ((:~:) (..))

import           GHC.TypeLits.Compare   ((%<=?), (:<=?) (LE, NLE))
import           GHC.TypeLits.Witnesses (SNat (SNat))
#endif

-- $setup
--
-- To use type level numeric literals you need to enable
-- the @DataKinds@ extension:
--
-- >>> :set -XDataKinds
--
-- To use infix syntax for @'Mod'@ or the @/@ synonym,
-- enable @TypeOperators@:
--
-- >>> :set -XTypeOperators
--
-- To use type applications with @'toMod'@ and friends:
--
-- >>> :set -XTypeApplications
--

-- $doc
--
-- @'Mod'@ and its synonym @/@ let you wrap arbitrary numeric types
-- in a modulus. To work with integers (mod 7) backed by @'Integer'@,
-- you could use one of the following equivalent types:
--
-- > Mod Integer 7
-- > Integer `Mod` 7
-- > Integer/7
-- > ℤ/7
--
-- (@'ℤ'@ is a synonym for @'Integer'@ provided by this library. In
-- Emacs, you can use the TeX input mode to type it with @\\Bbb{Z}@.)
--
-- The usual numeric typeclasses are defined for these types. You can
-- always extract the underlying value with @'unMod'@.
--
-- Here is a quick example:
--
-- >>> 10 * 11 :: ℤ/7
-- 5
--
-- It also works correctly with negative numeric literals:
--
-- >>> (-10) * 11 :: ℤ/7
-- 2
--
-- Modular division is an inverse of modular multiplication.
-- It is defined when divisor is coprime to modulus:
--
-- >>> 7 / 3 :: ℤ/16
-- 13
-- >>> 3 * 13 :: ℤ/16
-- 7
--
-- Note that it raises an exception if the divisor is *not* coprime to
-- the modulus:
--
-- >>> 7 / 4 :: ℤ/16
-- *** Exception: Cannot invert 4 (mod 16): not coprime to modulus.
-- ...
--
-- To use type level numeric literals you need to enable the
-- @DataKinds@ extension and to use infix syntax for @Mod@ or the @/@
-- synonym, you need @TypeOperators@.

-- | Wraps an underlying @Integeral@ type @i@ in a newtype annotated
-- with the bound @n@.
newtype i `Mod` (n :: Nat) = Mod i deriving (Mod i n -> Mod i n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall i (n :: Natural). Eq i => Mod i n -> Mod i n -> Bool
/= :: Mod i n -> Mod i n -> Bool
$c/= :: forall i (n :: Natural). Eq i => Mod i n -> Mod i n -> Bool
== :: Mod i n -> Mod i n -> Bool
$c== :: forall i (n :: Natural). Eq i => Mod i n -> Mod i n -> Bool
Eq, Mod i n -> Mod i n -> Bool
Mod i n -> Mod i n -> Ordering
Mod i n -> Mod i n -> Mod i n
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 {i} {n :: Natural}. Ord i => Eq (Mod i n)
forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Bool
forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Ordering
forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Mod i n
min :: Mod i n -> Mod i n -> Mod i n
$cmin :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Mod i n
max :: Mod i n -> Mod i n -> Mod i n
$cmax :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Mod i n
>= :: Mod i n -> Mod i n -> Bool
$c>= :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Bool
> :: Mod i n -> Mod i n -> Bool
$c> :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Bool
<= :: Mod i n -> Mod i n -> Bool
$c<= :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Bool
< :: Mod i n -> Mod i n -> Bool
$c< :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Bool
compare :: Mod i n -> Mod i n -> Ordering
$ccompare :: forall i (n :: Natural). Ord i => Mod i n -> Mod i n -> Ordering
Ord)

-- | Extract the underlying integral value from a modular type.
--
-- >>> unMod (10 :: ℤ/4)
-- 2
unMod :: i `Mod` n -> i
unMod :: forall i (n :: Natural). Mod i n -> i
unMod (Mod i
i) = i
i

-- | A synonym for @Mod@, inspired by the ℤ/n syntax from mathematics.
type (/) = Mod

-- | A synonym for Integer, also inspired by the ℤ/n syntax.
type ℤ   = Integer

-- | The modulus has to be a non-zero type-level natural number.
type Modulus n = (KnownNat n, 1 <= n)

-- | Helper function to get the modulus of a @ℤ/n@ as a value. Used
-- with type applications:
--
-- >>> modulus @5
-- 5
--
modulus :: forall n i. (Integral i, Modulus n) => i
modulus :: forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus = forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal (forall {k} (t :: k). Proxy t
Proxy :: Proxy n)

-- | Injects a value of the underlying type into the modulus type,
-- wrapping as appropriate.
--
-- If @n@ is ambiguous, you can specify it with @TypeApplications@:
--
-- >>> toMod @6 10
-- 4
--
-- Note that @n@ cannot be 0.
toMod :: forall n i. (Integral i, Modulus n) => i -> i `Mod` n
toMod :: forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod i
i = forall i (n :: Natural). i -> Mod i n
Mod forall a b. (a -> b) -> a -> b
$ i
i forall a. Integral a => a -> a -> a
`mod` (forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n)

-- | Convert an integral number @i@ into a @'Mod'@ value with the
-- type-level modulus @n@ specified with a proxy argument.
--
-- This lets you use 'toMod' without @TypeApplications@ in contexts
-- where @n@ is ambiguous.
modVal :: forall i proxy n. (Integral i, Modulus n)
       => i
       -> proxy n
       -> Mod i n
modVal :: forall i (proxy :: Natural -> *) (n :: Natural).
(Integral i, Modulus n) =>
i -> proxy n -> Mod i n
modVal i
i proxy n
_ = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod i
i

-- | Wraps an integral number, converting between integral types.
toMod' :: forall n i j. (Integral i, Integral j, Modulus n)
       => i
       -> j `Mod` n
toMod' :: forall (n :: Natural) i j.
(Integral i, Integral j, Modulus n) =>
i -> Mod j n
toMod' i
i = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ i
i forall a. Integral a => a -> a -> a
`mod` (forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n)

instance Show i => Show (i `Mod` n) where show :: Mod i n -> String
show (Mod i
i) = forall a. Show a => a -> String
show i
i
instance (Read i, Integral i, Modulus n) => Read (i `Mod` n)
  where readsPrec :: Int -> ReadS (Mod i n)
readsPrec Int
prec = forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Read a => Int -> ReadS a
readsPrec Int
prec

instance (Integral i, Modulus n) => Num (i `Mod` n) where
  fromInteger :: Integer -> Mod i n
fromInteger = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => Integer -> a
fromInteger

  Mod i
i₁ + :: Mod i n -> Mod i n -> Mod i n
+ Mod i
i₂ = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall a b. (a -> b) -> a -> b
$ i
i₁ forall a. Num a => a -> a -> a
+ i
i₂
  Mod i
i₁ * :: Mod i n -> Mod i n -> Mod i n
* Mod i
i₂ = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall a b. (a -> b) -> a -> b
$ i
i₁ forall a. Num a => a -> a -> a
* i
i₂

  abs :: Mod i n -> Mod i n
abs    (Mod i
i) = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
abs i
i
  signum :: Mod i n -> Mod i n
signum (Mod i
i) = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
signum i
i
  negate :: Mod i n -> Mod i n
negate (Mod i
i) = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
negate i
i

instance (Integral i, Modulus n) => Enum (i `Mod` n) where
  toEnum :: Int -> Mod i n
toEnum = forall a b. (Integral a, Num b) => a -> b
fromIntegral
  fromEnum :: Mod i n -> Int
fromEnum = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i (n :: Natural). Mod i n -> i
unMod

  -- implementation straight from the report
  enumFrom :: Mod i n -> [Mod i n]
enumFrom     Mod i n
x   = forall a. Enum a => a -> a -> [a]
enumFromTo     Mod i n
x forall a. Bounded a => a
maxBound
  enumFromThen :: Mod i n -> Mod i n -> [Mod i n]
enumFromThen Mod i n
x Mod i n
y = forall a. Enum a => a -> a -> a -> [a]
enumFromThenTo Mod i n
x Mod i n
y Mod i n
bound
    where
      bound :: Mod i n
bound | forall a. Enum a => a -> Int
fromEnum Mod i n
y forall a. Ord a => a -> a -> Bool
>= forall a. Enum a => a -> Int
fromEnum Mod i n
x = forall a. Bounded a => a
maxBound
            | Bool
otherwise                = forall a. Bounded a => a
minBound

instance (Integral i, Modulus n) => Bounded (i `Mod` n) where
  maxBound :: Mod i n
maxBound = forall i (n :: Natural). i -> Mod i n
Mod forall a b. (a -> b) -> a -> b
$ forall a. Enum a => a -> a
pred (forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n)
  minBound :: Mod i n
minBound = Mod i n
0

instance (Integral i, Modulus n) => Real (i `Mod` n) where
  toRational :: Mod i n -> Rational
toRational (Mod i
i) = forall a. Integral a => a -> Integer
toInteger i
i forall a. Integral a => a -> a -> Ratio a
% Integer
1

-- | Division uses modular inverse 'inv' so it is only possible to
-- divide by numbers coprime to @n@.
--
-- >>> 1 / 3 :: ℤ/7
-- 5
-- >>> 3 * 5 :: ℤ/7
-- 1
--
-- >>> 2 / 5 :: ℤ/7
-- 6
-- >>> 5 * 6 :: ℤ/7
-- 2
--
-- Dividing by a number that is not coprime to @n@ will raise an
-- error. Use 'inv' directly if you want to avoid this.
--
-- >>> 2 / 7 :: ℤ/7
-- *** Exception: Cannot invert 0 (mod 7): not coprime to modulus.
-- ...
--
instance (Integral i, Modulus n) => Fractional (i `Mod` n) where
  fromRational :: Rational -> Mod i n
fromRational Rational
r =
    forall a. Num a => Integer -> a
fromInteger (forall a. Ratio a -> a
numerator Rational
r) forall a. Fractional a => a -> a -> a
/ forall a. Num a => Integer -> a
fromInteger (forall a. Ratio a -> a
denominator Rational
r)
  recip :: Mod i n -> Mod i n
recip Mod i n
i = Maybe (Mod i n) -> Mod i n
unwrap forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) i.
(Modulus n, Integral i) =>
(i / n) -> Maybe (i / n)
inv Mod i n
i
    where
      unwrap :: Maybe (Mod i n) -> Mod i n
unwrap (Just Mod i n
x) = Mod i n
x
      unwrap Maybe (Mod i n)
Nothing  =
        let i' :: Integer
i'     = forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall i (n :: Natural). Mod i n -> i
unMod Mod i n
i
            bound' :: Integer
bound' = forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n @Integer
        in forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$
             forall r. PrintfType r => String -> r
printf String
"Cannot invert %d (mod %d): not coprime to modulus." Integer
i' Integer
bound'

-- | The modular inverse.
--
-- >>> inv 3 :: Maybe (ℤ/7)
-- Just 5
-- >>> 3 * 5 :: ℤ/7
-- 1
--
-- Note that only numbers coprime to @n@ have an inverse modulo @n@:
--
-- >>> inv 6 :: Maybe (ℤ/15)
-- Nothing
--
inv :: forall n i. (Modulus n, Integral i) => (i/n) -> Maybe (i/n)
inv :: forall (n :: Natural) i.
(Modulus n, Integral i) =>
(i / n) -> Maybe (i / n)
inv (Mod i
k) = forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {b}. Integral b => b -> b -> Maybe (b, b)
inv' (forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n) i
k
  where
    -- backwards Euclidean algorithm
    inv' :: b -> b -> Maybe (b, b)
inv' b
_ b
0 = forall a. Maybe a
Nothing
    inv' b
_ b
1 = forall a. a -> Maybe a
Just (b
0, b
1)
    inv' b
n b
x = do
      let (b
q,  b
r)  = b
n forall a. Integral a => a -> a -> (a, a)
`quotRem` b
x
      (b
q', b
r') <- b -> b -> Maybe (b, b)
inv' b
x b
r
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
r', b
q' forall a. Num a => a -> a -> a
- b
r' forall a. Num a => a -> a -> a
* b
q)

-- | A modular number with an unknown modulus.
--
-- Conceptually @SomeMod i = ∃n. i/n@.
data SomeMod i where
  SomeMod :: forall i (n :: Nat). Modulus n => Mod i n -> SomeMod i

-- | Shows both the number *and* its modulus:
--
-- >>> show (someModVal 10 4)
-- "Just (someModVal 2 4)"
--
-- This doesn't *quite* follow the rule that the show instance should
-- be a Haskell expression that evaluates to the given
-- value—'someModVal' returns a 'Maybe'—but this seems like the
-- closest we can reasonably get.
instance Show i => Show (SomeMod i) where
  showsPrec :: Int -> SomeMod i -> ShowS
showsPrec Int
p (SomeMod (i / n
x :: i/n)) = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
10) forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString forall a b. (a -> b) -> a -> b
$ forall r. PrintfType r => String -> r
printf String
"someModVal %s %d" (forall a. Show a => a -> String
show i / n
x) (forall (n :: Natural) i. (Integral i, Modulus n) => i
modulus @n @Integer)

-- | Convert an integral number @i@ into @'SomeMod'@ with the modulus
-- given at runtime.
--
-- That is, given @i :: ℤ@, @someModVal i modulus@ is equivalent to @i ::
-- ℤ/modulus@ except we don't know @modulus@ statically.
--
-- >>> someModVal 10 4
-- Just (someModVal 2 4)
--
-- Will return 'Nothing' if the modulus is 0 or negative:
--
-- >>> someModVal 10 (-10)
-- Nothing
--
-- >>> someModVal 10 0
-- Nothing
--
someModVal :: Integral i
           => i
           -- ^ Underlying integer @i@
           -> Integer
           -- ^ Modulus @n@
           -> Maybe (SomeMod i)
someModVal :: forall i. Integral i => i -> Integer -> Maybe (SomeMod i)
someModVal i
i Integer
n = do
  SomeNat (Proxy n
_ :: Proxy n) <- Integer -> Maybe SomeNat
someNatVal Integer
n
#if MIN_VERSION_base(4,16,0)
  case forall {k} (t :: k). Proxy t
Proxy @1 forall (a :: Natural) (b :: Natural) (proxy1 :: Natural -> *)
       (proxy2 :: Natural -> *).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI a b
`cmpNat` forall {k} (t :: k). Proxy t
Proxy @n of
    OrderingI 1 n
LTI -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall i (n :: Natural). Modulus n => Mod i n -> SomeMod i
SomeMod forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod @n i
i
    OrderingI 1 n
EQI -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall i (n :: Natural). Modulus n => Mod i n -> SomeMod i
SomeMod forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) i. (Integral i, Modulus n) => i -> Mod i n
toMod @n i
i
    OrderingI 1 n
GTI -> forall a. Maybe a
Nothing
#else
  case SNat @1 %<=? SNat @n of
    LE Refl -> pure $ SomeMod $ toMod @n i
    NLE _ _ -> Nothing
#endif