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

-- |
-- 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           Data.Type.Equality     ((:~:)(..))

import           Text.Printf   (printf)

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

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

-- $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
(Mod i n -> Mod i n -> Bool)
-> (Mod i n -> Mod i n -> Bool) -> Eq (Mod i n)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall i (n :: Nat). Eq i => Mod i n -> Mod i n -> Bool
/= :: Mod i n -> Mod i n -> Bool
$c/= :: forall i (n :: Nat). Eq i => Mod i n -> Mod i n -> Bool
== :: Mod i n -> Mod i n -> Bool
$c== :: forall i (n :: Nat). Eq i => Mod i n -> Mod i n -> Bool
Eq, Eq (Mod i n)
Eq (Mod i n) =>
(Mod i n -> Mod i n -> Ordering)
-> (Mod i n -> Mod i n -> Bool)
-> (Mod i n -> Mod i n -> Bool)
-> (Mod i n -> Mod i n -> Bool)
-> (Mod i n -> Mod i n -> Bool)
-> (Mod i n -> Mod i n -> Mod i n)
-> (Mod i n -> Mod i n -> Mod i n)
-> Ord (Mod i n)
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 :: Nat). Ord i => Eq (Mod i n)
forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Bool
forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Ordering
forall i (n :: Nat). 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 :: Nat). 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 :: Nat). Ord i => Mod i n -> Mod i n -> Mod i n
>= :: Mod i n -> Mod i n -> Bool
$c>= :: forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Bool
> :: Mod i n -> Mod i n -> Bool
$c> :: forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Bool
<= :: Mod i n -> Mod i n -> Bool
$c<= :: forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Bool
< :: Mod i n -> Mod i n -> Bool
$c< :: forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Bool
compare :: Mod i n -> Mod i n -> Ordering
$ccompare :: forall i (n :: Nat). Ord i => Mod i n -> Mod i n -> Ordering
$cp1Ord :: forall i (n :: Nat). Ord i => Eq (Mod i n)
Ord)

-- | Extract the underlying integral value from a modular type.
--
-- >>> unMod (10 :: ℤ/4)
-- 2
unMod :: i `Mod` n -> i
unMod :: Mod i n -> i
unMod (Mod i :: 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 :: i
modulus = Integer -> i
forall a. Num a => Integer -> a
fromInteger (Integer -> i) -> Integer -> i
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
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 :: i -> Mod i n
toMod i :: i
i = i -> Mod i n
forall i (n :: Nat). i -> Mod i n
Mod (i -> Mod i n) -> i -> Mod i n
forall a b. (a -> b) -> a -> b
$ i
i i -> i -> i
forall a. Integral a => a -> a -> a
`mod` (forall i. (Integral i, Modulus n) => i
forall (n :: Nat) 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 :: i -> proxy n -> Mod i n
modVal i :: i
i _ = i -> Mod i n
forall (n :: Nat) 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' :: i -> Mod j n
toMod' i :: i
i = j -> Mod j n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (j -> Mod j n) -> (i -> j) -> i -> Mod j n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> j
forall a b. (Integral a, Num b) => a -> b
fromIntegral (i -> Mod j n) -> i -> Mod j n
forall a b. (a -> b) -> a -> b
$ i
i i -> i -> i
forall a. Integral a => a -> a -> a
`mod` (forall i. (Integral i, Modulus n) => i
forall (n :: Nat) 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
i) = i -> String
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 prec :: Int
prec = ((i, String) -> (Mod i n, String))
-> [(i, String)] -> [(Mod i n, String)]
forall a b. (a -> b) -> [a] -> [b]
map ((i -> Mod i n) -> (i, String) -> (Mod i n, String)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod) ([(i, String)] -> [(Mod i n, String)])
-> (String -> [(i, String)]) -> ReadS (Mod i n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String -> [(i, String)]
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 = i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (i -> Mod i n) -> (Integer -> i) -> Integer -> Mod i n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> i
forall a. Num a => Integer -> a
fromInteger

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

  abs :: Mod i n -> Mod i n
abs    (Mod i :: i
i) = i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (i -> Mod i n) -> i -> Mod i n
forall a b. (a -> b) -> a -> b
$ i -> i
forall a. Num a => a -> a
abs i
i
  signum :: Mod i n -> Mod i n
signum (Mod i :: i
i) = i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (i -> Mod i n) -> i -> Mod i n
forall a b. (a -> b) -> a -> b
$ i -> i
forall a. Num a => a -> a
signum i
i
  negate :: Mod i n -> Mod i n
negate (Mod i :: i
i) = i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (i -> Mod i n) -> i -> Mod i n
forall a b. (a -> b) -> a -> b
$ i -> i
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 = Int -> Mod i n
forall a b. (Integral a, Num b) => a -> b
fromIntegral
  fromEnum :: Mod i n -> Int
fromEnum = i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (i -> Int) -> (Mod i n -> i) -> Mod i n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mod i n -> i
forall i (n :: Nat). Mod i n -> i
unMod

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

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

instance (Integral i, Modulus n) => Real (i `Mod` n) where
  toRational :: Mod i n -> Rational
toRational (Mod i :: i
i) = i -> Integer
forall a. Integral a => a -> Integer
toInteger i
i Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% 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 r :: Rational
r =
    Integer -> Mod i n
forall a. Num a => Integer -> a
fromInteger (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r) Mod i n -> Mod i n -> Mod i n
forall a. Fractional a => a -> a -> a
/ Integer -> Mod i n
forall a. Num a => Integer -> a
fromInteger (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r)
  recip :: Mod i n -> Mod i n
recip i :: Mod i n
i = Maybe (Mod i n) -> Mod i n
unwrap (Maybe (Mod i n) -> Mod i n) -> Maybe (Mod i n) -> Mod i n
forall a b. (a -> b) -> a -> b
$ Mod i n -> Maybe (Mod i n)
forall (n :: Nat) 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 x :: Mod i n
x) = Mod i n
x
      unwrap Nothing  =
        let i' :: Integer
i'     = i -> Integer
forall a. Integral a => a -> Integer
toInteger (i -> Integer) -> i -> Integer
forall a b. (a -> b) -> a -> b
$ Mod i n -> i
forall i (n :: Nat). Mod i n -> i
unMod Mod i n
i
            bound' :: Integer
bound' = (Integral Integer, Modulus n) => Integer
forall (n :: Nat) i. (Integral i, Modulus n) => i
modulus @n @Integer
        in String -> Mod i n
forall a. HasCallStack => String -> a
error (String -> Mod i n) -> String -> Mod i n
forall a b. (a -> b) -> a -> b
$
             String -> Integer -> Integer -> String
forall r. PrintfType r => String -> r
printf "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 :: (i / n) -> Maybe (i / n)
inv (Mod k :: i
k) = i -> i / n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod (i -> i / n) -> ((i, i) -> i) -> (i, i) -> i / n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (i, i) -> i
forall a b. (a, b) -> b
snd ((i, i) -> i / n) -> Maybe (i, i) -> Maybe (i / n)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> i -> i -> Maybe (i, i)
forall b. Integral b => b -> b -> Maybe (b, b)
inv' (forall i. (Integral i, Modulus n) => i
forall (n :: Nat) i. (Integral i, Modulus n) => i
modulus @n) i
k
  where
    -- backwards Euclidean algorithm
    inv' :: b -> b -> Maybe (b, b)
inv' _ 0 = Maybe (b, b)
forall a. Maybe a
Nothing
    inv' _ 1 = (b, b) -> Maybe (b, b)
forall a. a -> Maybe a
Just (0, 1)
    inv' n :: b
n x :: b
x = do
      let (q :: b
q,  r :: b
r)  = b
n b -> b -> (b, b)
forall a. Integral a => a -> a -> (a, a)
`quotRem` b
x
      (q' :: b
q', r' :: b
r') <- b -> b -> Maybe (b, b)
inv' b
x b
r
      (b, b) -> Maybe (b, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
r', b
q' b -> b -> b
forall a. Num a => a -> a -> a
- b
r' b -> b -> b
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 p :: Int
p (SomeMod (i / n
x :: i/n)) = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString (String -> ShowS) -> String -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> String -> Integer -> String
forall r. PrintfType r => String -> r
printf "someModVal %s %d" ((i / n) -> String
forall a. Show a => a -> String
show i / n
x) ((Integral Integer, Modulus n) => Integer
forall (n :: Nat) 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 :: i -> Integer -> Maybe (SomeMod i)
someModVal i :: i
i n :: Integer
n = do
  SomeNat (Proxy n
_ :: p n) <- Integer -> Maybe SomeNat
someNatVal Integer
n
  case KnownNat 1 => SNat 1
forall (n :: Nat). KnownNat n => SNat n
SNat @1 SNat 1 -> SNat n -> 1 :<=? n
forall (m :: Nat) (n :: Nat) (p :: Nat -> *) (q :: Nat -> *).
(KnownNat m, KnownNat n) =>
p m -> q n -> m :<=? n
%<=? KnownNat n => SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat @n of
    LE Refl -> SomeMod i -> Maybe (SomeMod i)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SomeMod i -> Maybe (SomeMod i)) -> SomeMod i -> Maybe (SomeMod i)
forall a b. (a -> b) -> a -> b
$ Mod i n -> SomeMod i
forall i (n :: Nat). Modulus n => Mod i n -> SomeMod i
SomeMod (Mod i n -> SomeMod i) -> Mod i n -> SomeMod i
forall a b. (a -> b) -> a -> b
$ i -> Mod i n
forall (n :: Nat) i. (Integral i, Modulus n) => i -> Mod i n
toMod @n i
i
    NLE _ _ -> Maybe (SomeMod i)
forall a. Maybe a
Nothing