    -- |
    -- 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 :: forall (m :: Natural). 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 :: Natural).
    Integral a =>
    SFactors a m -> (() :: Constraint) :- KnownNat m
    proofFromSFactors SFactors Integer m
    sm of
      Sub Dict (KnownNat m)
    (() :: 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 :: Natural). Mod m -> Natural
    unMod Mod m
    a)) (SFactors Integer m -> [(Prime Integer, Word)]
    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 = ((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 a. (a -> a -> a) -> [a] -> a
    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 :: Natural). Mod m -> Natural
    unMod (forall (p :: Natural). 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' :: forall (p :: Natural). KnownNat p => 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 :: Natural). KnownNat p => Mod p
    sqrtOfMinusOne
      | Bool
    otherwise          = Mod p -> Mod p
    forall (p :: Natural). KnownNat p => Mod p -> Mod p
    tonelliShanks Mod p
    square
      where
        prime :: Natural
    prime = Mod p -> Natural
    forall (n :: Natural) (proxy :: Natural -> *).
    KnownNat n =>
    proxy n -> Natural
    natVal Mod p
    square

    -- | @p@ must be of form @4k + 1@
    sqrtOfMinusOne :: KnownNat p => Mod p
    sqrtOfMinusOne :: forall (p :: Natural). KnownNat p => Mod p
    sqrtOfMinusOne = Mod p
    res
      where
        p :: Natural
    p = Mod p -> Natural
    forall (n :: Natural) (proxy :: Natural -> *).
    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. HasCallStack => [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 :: 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 = Mod p -> Natural
    forall (n :: Natural) (proxy :: Natural -> *).
    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 :: Natural). 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 :: Natural). KnownNat p => Mod p -> Integer
    fixup (Mod n -> Integer) -> Mod n -> Integer
    forall a b. (a -> b) -> a -> b
    $ forall (p :: Natural). 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 :: forall (p :: Natural). KnownNat p => 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 :: Natural).
    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 :: Natural). 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 :: 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' 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 :: Natural).
    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 :: Natural). 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 {b} {a}.
    (Bits b, Eq a, Num a, Integral b) =>
    b -> a -> Maybe b
    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 :: b -> a -> Maybe b
    solve b
    _ a
    1 = b -> Maybe b
    forall a. a -> Maybe a
    Just b
    1
            solve b
    1 a
    _ = b -> Maybe b
    forall a. a -> Maybe a
    Just b
    1
            solve b
    r a
    _
                | b -> Int
    forall a. Integral a => a -> Int
    rem4 b
    r Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    3   = Maybe b
    forall a. Maybe a
    Nothing  -- otherwise r ≡ 1 (mod 4)
                | b -> Int
    forall a. Integral a => a -> Int
    rem8 b
    r Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    5   = Maybe b
    forall a. Maybe a
    Nothing  -- otherwise r ≡ 1 (mod 8)
                | Bool
    otherwise     = b -> Word -> Maybe b
    fixup b
    r ((Word, b) -> Word
    forall a b. (a, b) -> a
    fst ((Word, b) -> Word) -> (Word, b) -> Word
    forall a b. (a -> b) -> a -> b
    $ b -> (Word, b)
    forall a. Integral a => a -> (Word, a)
    shiftToOddCount (b
    rb -> b -> b
    forall a. Num a => a -> a -> a
    -b
    1))
                  where
                    fixup :: b -> Word -> Maybe b
    fixup b
    x Word
    pw
                        | Word
    pw Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Word
    e2  = b -> Maybe b
    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 b -> b -> b
    forall a. Num a => a -> a -> a
    + (b
    1 b -> Int -> b
    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 :: b
    d = b
    x'b -> b -> b
    forall a. Num a => a -> a -> a
    *b
    x' b -> b -> b
    forall a. Num a => a -> a -> a
    - b
    r
                            pw' :: Word
    pw' = if b
    d b -> b -> Bool
    forall a. Eq a => a -> a -> Bool
    == b
    0 then Word
    e2 else (Word, b) -> Word
    forall a b. (a, b) -> a
    fst (b -> (Word, b)
    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 = 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 :: forall a. Integral a => 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 :: forall (p :: Natural). KnownNat p => Mod p
    findNonSquare = Mod n
    res
      where
        n :: Natural
    n = Mod n -> Natural
    forall (n :: Natural) (proxy :: Natural -> *).
    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. HasCallStack => [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]