-- |
-- Module:      Math.NumberTheory.Moduli.Sqrt
-- Copyright:   (c) 2011 Daniel Fischer
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Modular square roots and
-- <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol>.
--

{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns  #-}

module Math.NumberTheory.Moduli.Sqrt
  ( -- * Modular square roots
    sqrtsMod
  , sqrtsModFactorisation
  , sqrtsModPrimePower
  , sqrtsModPrime
    -- * Jacobi symbol
  , JacobiSymbol(..)
  , jacobi
  , symbolToNum
  ) where

import Control.Monad (liftM2)
import Data.Bits
import Data.Constraint
import Data.Maybe
import Data.Mod
import Data.Proxy
import GHC.TypeNats (KnownNat, SomeNat(..), natVal, someNatVal)

import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.JacobiSymbol
import Math.NumberTheory.Moduli.Singleton
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils (shiftToOddCount, splitOff)
import Math.NumberTheory.Utils.FromIntegral

-- | List all modular square roots.
--
-- >>> :set -XDataKinds
-- >>> sqrtsMod sfactors (1 :: Mod 60)
-- [(1 `modulo` 60),(49 `modulo` 60),(41 `modulo` 60),(29 `modulo` 60),(31 `modulo` 60),(19 `modulo` 60),(11 `modulo` 60),(59 `modulo` 60)]
sqrtsMod :: SFactors Integer m -> Mod m -> [Mod m]
sqrtsMod :: SFactors Integer m -> Mod m -> [Mod m]
sqrtsMod SFactors Integer m
sm Mod m
a = case SFactors Integer m -> (() :: Constraint) :- KnownNat m
forall a (m :: Nat).
Integral a =>
SFactors a m -> (() :: Constraint) :- KnownNat m
proofFromSFactors SFactors Integer m
sm of
  Sub (() :: Constraint) => Dict (KnownNat m)
Dict -> (Integer -> Mod m) -> [Integer] -> [Mod m]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Mod m
forall a. Num a => Integer -> a
fromInteger ([Integer] -> [Mod m]) -> [Integer] -> [Mod m]
forall a b. (a -> b) -> a -> b
$ Integer -> [(Prime Integer, Word)] -> [Integer]
sqrtsModFactorisation (Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod m -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod m
a)) (SFactors Integer m -> [(Prime Integer, Word)]
forall a (m :: Nat). SFactors a m -> [(Prime a, Word)]
unSFactors SFactors Integer m
sm)

-- | List all square roots modulo a number, the factorisation of which is
-- passed as a second argument.
--
-- >>> sqrtsModFactorisation 1 (factorise 60)
-- [1,49,41,29,31,19,11,59]
sqrtsModFactorisation :: Integer -> [(Prime Integer, Word)] -> [Integer]
sqrtsModFactorisation :: Integer -> [(Prime Integer, Word)] -> [Integer]
sqrtsModFactorisation Integer
_ []  = [Integer
0]
sqrtsModFactorisation Integer
n [(Prime Integer, Word)]
pps = ((Integer, Integer) -> Integer)
-> [(Integer, Integer)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst ([(Integer, Integer)] -> [Integer])
-> [(Integer, Integer)] -> [Integer]
forall a b. (a -> b) -> a -> b
$ ([(Integer, Integer)]
 -> [(Integer, Integer)] -> [(Integer, Integer)])
-> [[(Integer, Integer)]] -> [(Integer, Integer)]
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (((Integer, Integer) -> (Integer, Integer) -> (Integer, Integer))
-> [(Integer, Integer)]
-> [(Integer, Integer)]
-> [(Integer, Integer)]
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (Integer, Integer) -> (Integer, Integer) -> (Integer, Integer)
forall b.
(Ord b, Num b, Ring b, Euclidean b) =>
(b, b) -> (b, b) -> (b, b)
comb) [[(Integer, Integer)]]
cs
  where
    ms :: [Integer]
    ms :: [Integer]
ms = ((Prime Integer, Word) -> Integer)
-> [(Prime Integer, Word)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (\(Prime Integer
p, Word
pow) -> Prime Integer -> Integer
forall a. Prime a -> a
unPrime Prime Integer
p Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
pow) [(Prime Integer, Word)]
pps

    rs :: [[Integer]]
    rs :: [[Integer]]
rs = ((Prime Integer, Word) -> [Integer])
-> [(Prime Integer, Word)] -> [[Integer]]
forall a b. (a -> b) -> [a] -> [b]
map ((Prime Integer -> Word -> [Integer])
-> (Prime Integer, Word) -> [Integer]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Integer -> Prime Integer -> Word -> [Integer]
sqrtsModPrimePower Integer
n)) [(Prime Integer, Word)]
pps

    cs :: [[(Integer, Integer)]]
    cs :: [[(Integer, Integer)]]
cs = ([Integer] -> Integer -> [(Integer, Integer)])
-> [[Integer]] -> [Integer] -> [[(Integer, Integer)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\[Integer]
l Integer
m -> (Integer -> (Integer, Integer))
-> [Integer] -> [(Integer, Integer)]
forall a b. (a -> b) -> [a] -> [b]
map (, Integer
m) [Integer]
l) [[Integer]]
rs [Integer]
ms

    comb :: (b, b) -> (b, b) -> (b, b)
comb (b, b)
t1 (b, b)
t2 = (if b
ch b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
0 then b
ch b -> b -> b
forall a. Num a => a -> a -> a
+ b
m else b
ch, b
m)
      where
        (b
ch, b
m) = Maybe (b, b) -> (b, b)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (b, b) -> (b, b)) -> Maybe (b, b) -> (b, b)
forall a b. (a -> b) -> a -> b
$ (b, b) -> (b, b) -> Maybe (b, b)
forall a.
(Eq a, Ring a, Euclidean a) =>
(a, a) -> (a, a) -> Maybe (a, a)
chinese (b, b)
t1 (b, b)
t2

-- | List all square roots modulo the power of a prime.
--
-- >>> import Data.Maybe
-- >>> import Math.NumberTheory.Primes
-- >>> sqrtsModPrimePower 7 (fromJust (isPrime 3)) 2
-- [4,5]
-- >>> sqrtsModPrimePower 9 (fromJust (isPrime 3)) 3
-- [3,12,21,24,6,15]
sqrtsModPrimePower :: Integer -> Prime Integer -> Word -> [Integer]
sqrtsModPrimePower :: Integer -> Prime Integer -> Word -> [Integer]
sqrtsModPrimePower Integer
nn Prime Integer
p Word
1 = Integer -> Prime Integer -> [Integer]
sqrtsModPrime Integer
nn Prime Integer
p
sqrtsModPrimePower Integer
nn (Prime Integer -> Integer
forall a. Prime a -> a
unPrime -> Integer
prime) Word
expo = let primeExpo :: Integer
primeExpo = Integer
prime Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
expo in
  case Integer -> Integer -> (Word, Integer)
forall a. (Eq a, GcdDomain a) => a -> a -> (Word, a)
splitOff Integer
prime (Integer
nn Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
primeExpo) of
    (Word
_, Integer
0) -> [Integer
0, Integer
prime Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ ((Word
expo Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2) .. Integer
primeExpo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
    (Word
kk, Integer
n)
      | Word -> Bool
forall a. Integral a => a -> Bool
odd Word
kk    -> []
      | Bool
otherwise -> case (if Integer
prime Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 then Integer -> Word -> Maybe Integer
sqM2P Integer
n Word
expo' else Integer -> Integer -> Word -> Maybe Integer
sqrtModPP' Integer
n Integer
prime Word
expo') of
        Maybe Integer
Nothing -> []
        Just Integer
r  -> let rr :: Integer
rr = Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
prime Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
k in
          if Integer
prime Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 Bool -> Bool -> Bool
&& Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1 Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
t
          then Integer -> [Integer] -> [Integer]
go Integer
rr [Integer]
os
          else Integer -> [Integer] -> [Integer]
go Integer
rr [Integer]
os [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ Integer -> [Integer] -> [Integer]
go (Integer
primeExpo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
rr) [Integer]
os
      where
        k :: Word
k = Word
kk Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2
        t :: Word
t = (if Integer
prime Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 then Word
expo Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1 else Word
expo Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
k) Word -> Word -> Word
forall a. Ord a => a -> a -> a
`max` ((Word
expo Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2)
        expo' :: Word
expo' = Word
expo Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
2 Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
k
        os :: [Integer]
os = [Integer
0, Integer
prime Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
t .. Integer
primeExpo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]

        -- equivalent to map ((`mod` primeExpo) . (+ r)) rs,
        -- but avoids division
        go :: Integer -> [Integer] -> [Integer]
go Integer
r [Integer]
rs = (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
r) [Integer]
ps [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
primeExpo)) [Integer]
qs
          where
            ([Integer]
ps, [Integer]
qs) = (Integer -> Bool) -> [Integer] -> ([Integer], [Integer])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
primeExpo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
r) [Integer]
rs

-- | List all square roots by prime modulo.
--
-- >>> import Data.Maybe
-- >>> import Math.NumberTheory.Primes
-- >>> sqrtsModPrime 1 (fromJust (isPrime 5))
-- [1,4]
-- >>> sqrtsModPrime 0 (fromJust (isPrime 5))
-- [0]
-- >>> sqrtsModPrime 2 (fromJust (isPrime 5))
-- []
sqrtsModPrime :: Integer -> Prime Integer -> [Integer]
sqrtsModPrime :: Integer -> Prime Integer -> [Integer]
sqrtsModPrime Integer
n (Prime Integer -> Integer
forall a. Prime a -> a
unPrime -> Integer
2) = [Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
2]
sqrtsModPrime Integer
n (Prime Integer -> Integer
forall a. Prime a -> a
unPrime -> Integer
prime) = case Integer -> Integer -> JacobiSymbol
forall a. (Integral a, Bits a) => a -> a -> JacobiSymbol
jacobi Integer
n Integer
prime of
  JacobiSymbol
MinusOne -> []
  JacobiSymbol
Zero     -> [Integer
0]
  JacobiSymbol
One      -> case Natural -> SomeNat
someNatVal (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
prime) of
    SomeNat (Proxy n
_ :: Proxy p) -> let r :: Integer
r = Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod n -> Natural
forall (m :: Nat). Mod m -> Natural
unMod (Mod n -> Mod n
forall (p :: Nat). KnownNat p => Mod p -> Mod p
sqrtModP' @p (Integer -> Mod n
forall a. Num a => Integer -> a
fromInteger Integer
n))) in [Integer
r, Integer
prime Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
r]

-------------------------------------------------------------------------------
-- Internals

-- | @sqrtModP' square prime@ finds a square root of @square@ modulo
--   prime. @prime@ /must/ be a (positive) prime, and @square@ /must/ be a positive
--   quadratic residue modulo @prime@, i.e. @'jacobi square prime == 1@.
sqrtModP' :: KnownNat p => Mod p -> Mod p
sqrtModP' :: Mod p -> Mod p
sqrtModP' Mod p
square
  | Natural
prime Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
2         = Mod p
square
  | Natural -> Int
forall a. Integral a => a -> Int
rem4 Natural
prime Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3    = Mod p
square Mod p -> Natural -> Mod p
forall a b. (Num a, Integral b) => a -> b -> a
^ ((Natural
prime Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1) Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`quot` Natural
4)
  | Mod p
square Mod p -> Mod p -> Bool
forall a. Eq a => a -> a -> Bool
== Mod p
forall a. Bounded a => a
maxBound = Mod p
forall (p :: Nat). KnownNat p => Mod p
sqrtOfMinusOne
  | Bool
otherwise          = Mod p -> Mod p
forall (p :: Nat). KnownNat p => Mod p -> Mod p
tonelliShanks Mod p
square
  where
    prime :: Natural
prime = Mod p -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod p
square

-- | @p@ must be of form @4k + 1@
sqrtOfMinusOne :: KnownNat p => Mod p
sqrtOfMinusOne :: Mod p
sqrtOfMinusOne = Mod p
res
  where
    p :: Natural
p = Mod p -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod p
res
    k :: Natural
k = (Natural
p Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1) Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`quot` Natural
4
    res :: Mod p
res = [Mod p] -> Mod p
forall a. [a] -> a
head
      ([Mod p] -> Mod p) -> [Mod p] -> Mod p
forall a b. (a -> b) -> a -> b
$ (Mod p -> Bool) -> [Mod p] -> [Mod p]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\Mod p
n -> Mod p
n Mod p -> Mod p -> Bool
forall a. Eq a => a -> a -> Bool
== Mod p
1 Bool -> Bool -> Bool
|| Mod p
n Mod p -> Mod p -> Bool
forall a. Eq a => a -> a -> Bool
== Mod p
forall a. Bounded a => a
maxBound)
      ([Mod p] -> [Mod p]) -> [Mod p] -> [Mod p]
forall a b. (a -> b) -> a -> b
$ (Mod p -> Mod p) -> [Mod p] -> [Mod p]
forall a b. (a -> b) -> [a] -> [b]
map (Mod p -> Natural -> Mod p
forall a b. (Num a, Integral b) => a -> b -> a
^ Natural
k) [Mod p
2 .. Mod p
forall a. Bounded a => a
maxBound Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
- Mod p
1]

-- | @tonelliShanks square prime@ calculates a square root of @square@
--   modulo @prime@, where @prime@ is a prime of the form @4*k + 1@ and
--   @square@ is a positive quadratic residue modulo @prime@, using the
--   Tonelli-Shanks algorithm.
tonelliShanks :: forall p. KnownNat p => Mod p -> Mod p
tonelliShanks :: Mod p -> Mod p
tonelliShanks Mod p
square = Mod p -> Mod p -> Mod p -> Word -> Mod p
loop Mod p
rc Mod p
t1 Mod p
generator Word
log2
  where
    prime :: Natural
prime = Mod p -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod p
square
    (Word
log2, Natural
q) = Natural -> (Word, Natural)
forall a. Integral a => a -> (Word, a)
shiftToOddCount (Natural
prime Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1)
    generator :: Mod p
generator = Mod p
forall (p :: Nat). KnownNat p => Mod p
findNonSquare Mod p -> Natural -> Mod p
forall a b. (Num a, Integral b) => a -> b -> a
^ Natural
q
    rc :: Mod p
rc = Mod p
square Mod p -> Natural -> Mod p
forall a b. (Num a, Integral b) => a -> b -> a
^ ((Natural
q Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1) Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
`quot` Natural
2)
    t1 :: Mod p
t1 = Mod p
square Mod p -> Natural -> Mod p
forall a b. (Num a, Integral b) => a -> b -> a
^ Natural
q

    msquare :: t -> t -> t
msquare t
0 t
x = t
x
    msquare t
k t
x = t -> t -> t
msquare (t
kt -> t -> t
forall a. Num a => a -> a -> a
-t
1) (t
x t -> t -> t
forall a. Num a => a -> a -> a
* t
x)

    findPeriod :: t -> t -> t
findPeriod t
per t
1 = t
per
    findPeriod t
per t
x = t -> t -> t
findPeriod (t
per t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) (t
x t -> t -> t
forall a. Num a => a -> a -> a
* t
x)

    loop :: Mod p -> Mod p -> Mod p -> Word -> Mod p
    loop :: Mod p -> Mod p -> Mod p -> Word -> Mod p
loop !Mod p
r Mod p
t Mod p
c Word
m
        | Mod p
t Mod p -> Mod p -> Bool
forall a. Eq a => a -> a -> Bool
== Mod p
1    = Mod p
r
        | Bool
otherwise = Mod p -> Mod p -> Mod p -> Word -> Mod p
loop Mod p
nextR Mod p
nextT Mod p
nextC Word
nextM
          where
            nextM :: Word
nextM = Word -> Mod p -> Word
forall t t. (Eq t, Num t, Num t) => t -> t -> t
findPeriod Word
0 Mod p
t
            b :: Mod p
b     = Word -> Mod p -> Mod p
forall t t. (Eq t, Num t, Num t) => t -> t -> t
msquare (Word
m Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1 Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
nextM) Mod p
c
            nextR :: Mod p
nextR = Mod p
r Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
* Mod p
b
            nextC :: Mod p
nextC = Mod p
b Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
* Mod p
b
            nextT :: Mod p
nextT = Mod p
t Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
* Mod p
nextC

-- | prime must be odd, n must be coprime with prime
sqrtModPP' :: Integer -> Integer -> Word -> Maybe Integer
sqrtModPP' :: Integer -> Integer -> Word -> Maybe Integer
sqrtModPP' Integer
n Integer
prime Word
expo = case Integer -> Integer -> JacobiSymbol
forall a. (Integral a, Bits a) => a -> a -> JacobiSymbol
jacobi Integer
n Integer
prime of
  JacobiSymbol
MinusOne -> Maybe Integer
forall a. Maybe a
Nothing
  JacobiSymbol
Zero     -> Maybe Integer
forall a. Maybe a
Nothing
  JacobiSymbol
One      -> case Natural -> SomeNat
someNatVal (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
prime) of
    SomeNat (Proxy n
_ :: Proxy p) -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer -> Maybe Integer) -> Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$ Mod n -> Integer
forall (p :: Nat). KnownNat p => Mod p -> Integer
fixup (Mod n -> Integer) -> Mod n -> Integer
forall a b. (a -> b) -> a -> b
$ Mod n -> Mod n
forall (p :: Nat). KnownNat p => Mod p -> Mod p
sqrtModP' @p (Integer -> Mod n
forall a. Num a => Integer -> a
fromInteger Integer
n)
  where
    fixup :: KnownNat p => Mod p -> Integer
    fixup :: Mod p -> Integer
fixup Mod p
r
      | Integer
diff' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Integer
r'
      | Word
expo Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
e  = Integer
r'
      | Bool
otherwise  = Mod p -> Integer -> Mod p -> Integer -> Integer
forall (p :: Nat).
KnownNat p =>
Mod p -> Integer -> Mod p -> Integer -> Integer
hoist (Mod p -> Mod p
forall a. Fractional a => a -> a
recip (Mod p
2 Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
* Mod p
r)) Integer
r' (Integer -> Mod p
forall a. Num a => Integer -> a
fromInteger Integer
q) (Integer
primeInteger -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Word
e)
      where
        r' :: Integer
r' = Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod p -> Natural
forall (m :: Nat). Mod m -> Natural
unMod Mod p
r)
        diff' :: Integer
diff' = Integer
r' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
r' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
n
        (Word
e, Integer
q) = Integer -> Integer -> (Word, Integer)
forall a. (Eq a, GcdDomain a) => a -> a -> (Word, a)
splitOff Integer
prime Integer
diff'

    hoist :: KnownNat p => Mod p -> Integer -> Mod p -> Integer -> Integer
    hoist :: Mod p -> Integer -> Mod p -> Integer -> Integer
hoist Mod p
inv Integer
root Mod p
elim Integer
pp
      | Integer
diff' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0    = Integer
root'
      | Word
expo Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
ex    = Integer
root'
      | Bool
otherwise     = Mod p -> Integer -> Mod p -> Integer -> Integer
forall (p :: Nat).
KnownNat p =>
Mod p -> Integer -> Mod p -> Integer -> Integer
hoist Mod p
inv Integer
root' (Integer -> Mod p
forall a. Num a => Integer -> a
fromInteger Integer
nelim) (Integer
prime Integer -> Word -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
ex)
        where
          root' :: Integer
root' = Integer
root Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Mod p -> Natural
forall (m :: Nat). Mod m -> Natural
unMod (Mod p
inv Mod p -> Mod p -> Mod p
forall a. Num a => a -> a -> a
* Mod p -> Mod p
forall a. Num a => a -> a
negate Mod p
elim)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
pp
          diff' :: Integer
diff' = Integer
root' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
root' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
n
          (Word
ex, Integer
nelim) = Integer -> Integer -> (Word, Integer)
forall a. (Eq a, GcdDomain a) => a -> a -> (Word, a)
splitOff Integer
prime Integer
diff'

-- dirty, dirty
sqM2P :: Integer -> Word -> Maybe Integer
sqM2P :: Integer -> Word -> Maybe Integer
sqM2P Integer
n Word
e
    | Word
e Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
2     = Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
2)
    | Integer
n' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0   = Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
0
    | Word -> Bool
forall a. Integral a => a -> Bool
odd Word
k     = Maybe Integer
forall a. Maybe a
Nothing
    | Bool
otherwise = (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
mdl) (Integer -> Integer) -> (Integer -> Integer) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Word -> Int
wordToInt Word
k2) (Integer -> Integer) -> Maybe Integer -> Maybe Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Integer -> Word -> Maybe Integer
forall a a. (Num a, Eq a, Bits a, Integral a) => a -> a -> Maybe a
solve Integer
s Word
e2
      where
        mdl :: Integer
mdl = Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Word -> Int
wordToInt Word
e
        n' :: Integer
n' = Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
mdl
        (Word
k, Integer
s) = Integer -> (Word, Integer)
forall a. Integral a => a -> (Word, a)
shiftToOddCount Integer
n'
        k2 :: Word
k2 = Word
k Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2
        e2 :: Word
e2 = Word
e Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
k
        solve :: a -> a -> Maybe a
solve a
_ a
1 = a -> Maybe a
forall a. a -> Maybe a
Just a
1
        solve a
1 a
_ = a -> Maybe a
forall a. a -> Maybe a
Just a
1
        solve a
r a
_
            | a -> Int
forall a. Integral a => a -> Int
rem4 a
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3   = Maybe a
forall a. Maybe a
Nothing  -- otherwise r ≡ 1 (mod 4)
            | a -> Int
forall a. Integral a => a -> Int
rem8 a
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5   = Maybe a
forall a. Maybe a
Nothing  -- otherwise r ≡ 1 (mod 8)
            | Bool
otherwise     = a -> Word -> Maybe a
fixup a
r ((Word, a) -> Word
forall a b. (a, b) -> a
fst ((Word, a) -> Word) -> (Word, a) -> Word
forall a b. (a -> b) -> a -> b
$ a -> (Word, a)
forall a. Integral a => a -> (Word, a)
shiftToOddCount (a
ra -> a -> a
forall a. Num a => a -> a -> a
-a
1))
              where
                fixup :: a -> Word -> Maybe a
fixup a
x Word
pw
                    | Word
pw Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
e2  = a -> Maybe a
forall a. a -> Maybe a
Just a
x
                    | Bool
otherwise = a -> Word -> Maybe a
fixup a
x' Word
pw'
                      where
                        x' :: a
x' = a
x a -> a -> a
forall a. Num a => a -> a -> a
+ (a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` (Word -> Int
wordToInt Word
pw Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
                        d :: a
d = a
x'a -> a -> a
forall a. Num a => a -> a -> a
*a
x' a -> a -> a
forall a. Num a => a -> a -> a
- a
r
                        pw' :: Word
pw' = if a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then Word
e2 else (Word, a) -> Word
forall a b. (a, b) -> a
fst (a -> (Word, a)
forall a. Integral a => a -> (Word, a)
shiftToOddCount a
d)

-------------------------------------------------------------------------------
-- Utilities

rem4 :: Integral a => a -> Int
rem4 :: a -> Int
rem4 a
n = a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
3

rem8 :: Integral a => a -> Int
rem8 :: a -> Int
rem8 a
n = a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7

findNonSquare :: KnownNat n => Mod n
findNonSquare :: Mod n
findNonSquare = Mod n
res
  where
    n :: Natural
n = Mod n -> Natural
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod n
res
    res :: Mod n
res
      | Natural -> Int
forall a. Integral a => a -> Int
rem8 Natural
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3 Bool -> Bool -> Bool
|| Natural -> Int
forall a. Integral a => a -> Int
rem8 Natural
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5 = Mod n
2
      | Bool
otherwise = Natural -> Mod n
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural -> Mod n) -> Natural -> Mod n
forall a b. (a -> b) -> a -> b
$ [Natural] -> Natural
forall a. [a] -> a
head ([Natural] -> Natural) -> [Natural] -> Natural
forall a b. (a -> b) -> a -> b
$
        (Natural -> Bool) -> [Natural] -> [Natural]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\Natural
p -> Natural -> Natural -> JacobiSymbol
forall a. (Integral a, Bits a) => a -> a -> JacobiSymbol
jacobi Natural
p Natural
n JacobiSymbol -> JacobiSymbol -> Bool
forall a. Eq a => a -> a -> Bool
/= JacobiSymbol
MinusOne) [Natural]
candidates

    -- It is enough to consider only prime candidates, but
    -- the probability that the smallest non-residue is > 67
    -- is small and 'jacobi' test is fast,
    -- so we use [71..n] instead of filter isPrime [71..n].
    candidates :: [Natural]
candidates = Natural
3Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
5Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
7Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
11Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
13Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
17Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
19Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
23Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
29Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
31Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
37Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
41Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
43Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
47Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
53Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
59Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
61Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:Natural
67Natural -> [Natural] -> [Natural]
forall a. a -> [a] -> [a]
:[Natural
71..Natural
n]