-- |
-- 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 DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PostfixOperators #-}
{-# 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.List.Infinite (Infinite(..), (...))
import qualified Data.List.Infinite as Inf
import Data.Maybe
import Data.Mod
import Data.Proxy
import GHC.TypeNats (KnownNat, SomeNat(..), natVal, someNatVal, Nat)
import Numeric.Natural (Natural)

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 :: forall (m :: Natural). SFactors Integer m -> Mod m -> [Mod m]
sqrtsMod SFactors Integer m
sm Mod m
a = case forall a (m :: Natural).
Integral a =>
SFactors a m -> (() :: Constraint) :- KnownNat m
proofFromSFactors SFactors Integer m
sm of
  Sub Dict (KnownNat m)
(() :: Constraint) => Dict (KnownNat m)
Dict -> forall a b. (a -> b) -> [a] -> [b]
map forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ Integer -> [(Prime Integer, Word)] -> [Integer]
sqrtsModFactorisation (forall a. Integral a => a -> Integer
toInteger (forall (m :: Natural). Mod m -> Natural
unMod Mod m
a)) (forall a (m :: Natural). 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 = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 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 = forall a b. (a -> b) -> [a] -> [b]
map (\(Prime Integer
p, Word
pow) -> forall a. Prime a -> a
unPrime Prime Integer
p forall a b. (Num a, Integral b) => a -> b -> a
^ Word
pow) [(Prime Integer, Word)]
pps

    rs :: [[Integer]]
    rs :: [[Integer]]
rs = forall a b. (a -> b) -> [a] -> [b]
map (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 = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\[Integer]
l Integer
m -> 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 forall a. Ord a => a -> a -> Bool
< b
0 then b
ch forall a. Num a => a -> a -> a
+ b
m else b
ch, b
m)
      where
        (b
ch, b
m) = forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> 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 (forall a. Prime a -> a
unPrime -> Integer
prime) Word
expo = let primeExpo :: Integer
primeExpo = Integer
prime forall a b. (Num a, Integral b) => a -> b -> a
^ Word
expo in
  case forall a. (Eq a, GcdDomain a) => a -> a -> (Word, a)
splitOff Integer
prime (Integer
nn forall a. Integral a => a -> a -> a
`mod` Integer
primeExpo) of
    (Word
_, Integer
0) -> [Integer
0, Integer
prime forall a b. (Num a, Integral b) => a -> b -> a
^ ((Word
expo forall a. Num a => a -> a -> a
+ Word
1) forall a. Integral a => a -> a -> a
`quot` Word
2) .. Integer
primeExpo forall a. Num a => a -> a -> a
- Integer
1]
    (Word
kk, Integer
n)
      | forall a. Integral a => a -> Bool
odd Word
kk    -> []
      | Bool
otherwise -> case (if Integer
prime 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 forall a. Num a => a -> a -> a
* Integer
prime forall a b. (Num a, Integral b) => a -> b -> a
^ Word
k in
          if Integer
prime forall a. Eq a => a -> a -> Bool
== Integer
2 Bool -> Bool -> Bool
&& Word
k forall a. Num a => a -> a -> a
+ Word
1 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 forall a. [a] -> [a] -> [a]
++ Integer -> [Integer] -> [Integer]
go (Integer
primeExpo forall a. Num a => a -> a -> a
- Integer
rr) [Integer]
os
      where
        k :: Word
k = Word
kk forall a. Integral a => a -> a -> a
`quot` Word
2
        t :: Word
t = (if Integer
prime forall a. Eq a => a -> a -> Bool
== Integer
2 then Word
expo forall a. Num a => a -> a -> a
- Word
k forall a. Num a => a -> a -> a
- Word
1 else Word
expo forall a. Num a => a -> a -> a
- Word
k) forall a. Ord a => a -> a -> a
`max` ((Word
expo forall a. Num a => a -> a -> a
+ Word
1) forall a. Integral a => a -> a -> a
`quot` Word
2)
        expo' :: Word
expo' = Word
expo forall a. Num a => a -> a -> a
- Word
2 forall a. Num a => a -> a -> a
* Word
k
        os :: [Integer]
os = [Integer
0, Integer
prime forall a b. (Num a, Integral b) => a -> b -> a
^ Word
t .. Integer
primeExpo 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 = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ Integer
r) [Integer]
ps forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ (Integer
r forall a. Num a => a -> a -> a
- Integer
primeExpo)) [Integer]
qs
          where
            ([Integer]
ps, [Integer]
qs) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span (forall a. Ord a => a -> a -> Bool
< Integer
primeExpo 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 (forall a. Prime a -> a
unPrime -> Integer
2) = [Integer
n forall a. Integral a => a -> a -> a
`mod` Integer
2]
sqrtsModPrime Integer
n (forall a. Prime a -> a
unPrime -> Integer
prime) = case 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 (forall a. Num a => Integer -> a
fromInteger Integer
prime) of
    SomeNat (Proxy n
_ :: Proxy p) -> let r :: Integer
r = forall a. Integral a => a -> Integer
toInteger (forall (m :: Natural). Mod m -> Natural
unMod (forall (p :: Natural). KnownNat p => Mod p -> Mod p
sqrtModP' @p (forall a. Num a => Integer -> a
fromInteger Integer
n))) in [Integer
r, Integer
prime 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' :: forall (p :: Natural). KnownNat p => Mod p -> Mod p
sqrtModP' Mod p
square
  | Natural
prime forall a. Eq a => a -> a -> Bool
== Natural
2         = Mod p
square
  | forall a. Integral a => a -> Int
rem4 Natural
prime forall a. Eq a => a -> a -> Bool
== Int
3    = Mod p
square forall a b. (Num a, Integral b) => a -> b -> a
^ ((Natural
prime forall a. Num a => a -> a -> a
+ Natural
1) forall a. Integral a => a -> a -> a
`quot` Natural
4)
  | Mod p
square forall a. Eq a => a -> a -> Bool
== forall a. Bounded a => a
maxBound = forall (p :: Natural). KnownNat p => Mod p
sqrtOfMinusOne
  | Bool
otherwise          = forall (p :: Natural). KnownNat p => Mod p -> Mod p
tonelliShanks Mod p
square
  where
    prime :: Natural
prime = forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod p
square

-- | @p@ must be of form @4k + 1@
sqrtOfMinusOne :: forall (p :: Nat). KnownNat p => Mod p
sqrtOfMinusOne :: forall (p :: Natural). KnownNat p => Mod p
sqrtOfMinusOne = case [Mod p]
results of
  [] -> forall a. HasCallStack => [Char] -> a
error [Char]
"sqrtOfMinusOne: internal invariant violated"
  Mod p
hd : [Mod p]
_ -> Mod p
hd
  where
    p :: Natural
    p :: Natural
p = forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal (forall {k} (t :: k). Proxy t
Proxy :: Proxy p)

    k :: Natural
    k :: Natural
k = (Natural
p forall a. Num a => a -> a -> a
- Natural
1) forall a. Integral a => a -> a -> a
`quot` Natural
4

    results :: [Mod p]
    results :: [Mod p]
results = forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\Mod p
n -> Mod p
n forall a. Eq a => a -> a -> Bool
== Mod p
1 Bool -> Bool -> Bool
|| Mod p
n forall a. Eq a => a -> a -> Bool
== forall a. Bounded a => a
maxBound) forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (Num a, Integral b) => a -> b -> a
^ Natural
k) [Mod p
2 .. forall a. Bounded a => a
maxBound 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 :: forall (p :: Natural). KnownNat p => 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 = forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal Mod p
square
    (Word
log2, Natural
q) = forall a. Integral a => a -> (Word, a)
shiftToOddCount (Natural
prime forall a. Num a => a -> a -> a
- Natural
1)
    generator :: Mod p
generator = forall (p :: Natural). KnownNat p => Mod p
findNonSquare forall a b. (Num a, Integral b) => a -> b -> a
^ Natural
q
    rc :: Mod p
rc = Mod p
square forall a b. (Num a, Integral b) => a -> b -> a
^ ((Natural
q forall a. Num a => a -> a -> a
+ Natural
1) forall a. Integral a => a -> a -> a
`quot` Natural
2)
    t1 :: Mod p
t1 = Mod p
square 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
kforall a. Num a => a -> a -> a
-t
1) (t
x 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 forall a. Num a => a -> a -> a
+ t
1) (t
x 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 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 = forall {t} {t}. (Eq t, Num t, Num t) => t -> t -> t
findPeriod Word
0 Mod p
t
            b :: Mod p
b     = forall {t} {t}. (Eq t, Num t, Num t) => t -> t -> t
msquare (Word
m forall a. Num a => a -> a -> a
- Word
1 forall a. Num a => a -> a -> a
- Word
nextM) Mod p
c
            nextR :: Mod p
nextR = Mod p
r forall a. Num a => a -> a -> a
* Mod p
b
            nextC :: Mod p
nextC = Mod p
b forall a. Num a => a -> a -> a
* Mod p
b
            nextT :: Mod p
nextT = Mod p
t 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 forall a. (Integral a, Bits a) => a -> a -> JacobiSymbol
jacobi Integer
n Integer
prime of
  JacobiSymbol
MinusOne -> forall a. Maybe a
Nothing
  JacobiSymbol
Zero     -> forall a. Maybe a
Nothing
  JacobiSymbol
One      -> case Natural -> SomeNat
someNatVal (forall a. Num a => Integer -> a
fromInteger Integer
prime) of
    SomeNat (Proxy n
_ :: Proxy p) -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (p :: Natural). KnownNat p => Mod p -> Integer
fixup forall a b. (a -> b) -> a -> b
$ forall (p :: Natural). KnownNat p => Mod p -> Mod p
sqrtModP' @p (forall a. Num a => Integer -> a
fromInteger Integer
n)
  where
    fixup :: KnownNat p => Mod p -> Integer
    fixup :: forall (p :: Natural). KnownNat p => Mod p -> Integer
fixup Mod p
r
      | Integer
diff' forall a. Eq a => a -> a -> Bool
== Integer
0 = Integer
r'
      | Word
expo forall a. Ord a => a -> a -> Bool
<= Word
e  = Integer
r'
      | Bool
otherwise  = forall (p :: Natural).
KnownNat p =>
Mod p -> Integer -> Mod p -> Integer -> Integer
hoist (forall a. Fractional a => a -> a
recip (Mod p
2 forall a. Num a => a -> a -> a
* Mod p
r)) Integer
r' (forall a. Num a => Integer -> a
fromInteger Integer
q) (Integer
primeforall a b. (Num a, Integral b) => a -> b -> a
^Word
e)
      where
        r' :: Integer
r' = forall a. Integral a => a -> Integer
toInteger (forall (m :: Natural). Mod m -> Natural
unMod Mod p
r)
        diff' :: Integer
diff' = Integer
r' forall a. Num a => a -> a -> a
* Integer
r' forall a. Num a => a -> a -> a
- Integer
n
        (Word
e, Integer
q) = 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 :: forall (p :: Natural).
KnownNat p =>
Mod p -> Integer -> Mod p -> Integer -> Integer
hoist Mod p
inv Integer
root Mod p
elim Integer
pp
      | Integer
diff' forall a. Eq a => a -> a -> Bool
== Integer
0    = Integer
root'
      | Word
expo forall a. Ord a => a -> a -> Bool
<= Word
ex    = Integer
root'
      | Bool
otherwise     = forall (p :: Natural).
KnownNat p =>
Mod p -> Integer -> Mod p -> Integer -> Integer
hoist Mod p
inv Integer
root' (forall a. Num a => Integer -> a
fromInteger Integer
nelim) (Integer
prime forall a b. (Num a, Integral b) => a -> b -> a
^ Word
ex)
        where
          root' :: Integer
root' = Integer
root forall a. Num a => a -> a -> a
+ forall a. Integral a => a -> Integer
toInteger (forall (m :: Natural). Mod m -> Natural
unMod (Mod p
inv forall a. Num a => a -> a -> a
* forall a. Num a => a -> a
negate Mod p
elim)) forall a. Num a => a -> a -> a
* Integer
pp
          diff' :: Integer
diff' = Integer
root' forall a. Num a => a -> a -> a
* Integer
root' forall a. Num a => a -> a -> a
- Integer
n
          (Word
ex, Integer
nelim) = 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 forall a. Ord a => a -> a -> Bool
< Word
2     = forall a. a -> Maybe a
Just (Integer
n forall a. Integral a => a -> a -> a
`mod` Integer
2)
    | Integer
n' forall a. Eq a => a -> a -> Bool
== Integer
0   = forall a. a -> Maybe a
Just Integer
0
    | forall a. Integral a => a -> Bool
odd Word
k     = forall a. Maybe a
Nothing
    | Bool
otherwise = (forall a. Integral a => a -> a -> a
`mod` Integer
mdl) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Bits a => a -> Int -> a
`shiftL` Word -> Int
wordToInt Word
k2) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {a} {b}.
(Num a, Eq a, Bits b, Integral b) =>
b -> a -> Maybe b
solve Integer
s Word
e2
      where
        mdl :: Integer
mdl = Integer
1 forall a. Bits a => a -> Int -> a
`shiftL` Word -> Int
wordToInt Word
e
        n' :: Integer
n' = Integer
n forall a. Integral a => a -> a -> a
`mod` Integer
mdl
        (Word
k, Integer
s) = forall a. Integral a => a -> (Word, a)
shiftToOddCount Integer
n'
        k2 :: Word
k2 = Word
k forall a. Integral a => a -> a -> a
`quot` Word
2
        e2 :: Word
e2 = Word
e forall a. Num a => a -> a -> a
- Word
k
        solve :: b -> a -> Maybe b
solve b
_ a
1 = forall a. a -> Maybe a
Just b
1
        solve b
1 a
_ = forall a. a -> Maybe a
Just b
1
        solve b
r a
_
            | forall a. Integral a => a -> Int
rem4 b
r forall a. Eq a => a -> a -> Bool
== Int
3   = forall a. Maybe a
Nothing  -- otherwise r ≡ 1 (mod 4)
            | forall a. Integral a => a -> Int
rem8 b
r forall a. Eq a => a -> a -> Bool
== Int
5   = forall a. Maybe a
Nothing  -- otherwise r ≡ 1 (mod 8)
            | Bool
otherwise     = b -> Word -> Maybe b
fixup b
r (forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> (Word, a)
shiftToOddCount (b
rforall a. Num a => a -> a -> a
-b
1))
              where
                fixup :: b -> Word -> Maybe b
fixup b
x Word
pw
                    | Word
pw forall a. Ord a => a -> a -> Bool
>= Word
e2  = forall a. a -> Maybe a
Just b
x
                    | Bool
otherwise = b -> Word -> Maybe b
fixup b
x' Word
pw'
                      where
                        x' :: b
x' = b
x forall a. Num a => a -> a -> a
+ (b
1 forall a. Bits a => a -> Int -> a
`shiftL` (Word -> Int
wordToInt Word
pw forall a. Num a => a -> a -> a
- Int
1))
                        d :: b
d = b
x'forall a. Num a => a -> a -> a
*b
x' forall a. Num a => a -> a -> a
- b
r
                        pw' :: Word
pw' = if b
d forall a. Eq a => a -> a -> Bool
== b
0 then Word
e2 else forall a b. (a, b) -> a
fst (forall a. Integral a => a -> (Word, a)
shiftToOddCount b
d)

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

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

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

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

    -- 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 :: Infinite Natural
    candidates :: Infinite Natural
candidates = Natural
3 forall a. a -> Infinite a -> Infinite a
:< Natural
5 forall a. a -> Infinite a -> Infinite a
:< Natural
7 forall a. a -> Infinite a -> Infinite a
:< Natural
11 forall a. a -> Infinite a -> Infinite a
:< Natural
13 forall a. a -> Infinite a -> Infinite a
:< Natural
17 forall a. a -> Infinite a -> Infinite a
:< Natural
19 forall a. a -> Infinite a -> Infinite a
:< Natural
23 forall a. a -> Infinite a -> Infinite a
:< Natural
29 forall a. a -> Infinite a -> Infinite a
:< Natural
31 forall a. a -> Infinite a -> Infinite a
:<
      Natural
37 forall a. a -> Infinite a -> Infinite a
:< Natural
41 forall a. a -> Infinite a -> Infinite a
:< Natural
43 forall a. a -> Infinite a -> Infinite a
:< Natural
47 forall a. a -> Infinite a -> Infinite a
:< Natural
53 forall a. a -> Infinite a -> Infinite a
:< Natural
59 forall a. a -> Infinite a -> Infinite a
:< Natural
61 forall a. a -> Infinite a -> Infinite a
:< Natural
67 forall a. a -> Infinite a -> Infinite a
:< (Natural
71...)