-- SPDX-License-Identifier: MPL-2.0

{- | Exact integer algorithms. Expensive searches take an explicit nonnegative
iteration budget and report 'NonConvergence'. Primality by trial division is
proven; Miller--Rabin and Baillie--PSW are explicitly probable-prime tests.
-}
module HMath.NumberTheory (
    gcdInteger,
    lcmInteger,
    extendedGCD,
    bezout,
    modPow,
    modInverse,
    chineseRemainder,
    linearDiophantine,
    integerSquareRoot,
    perfectPower,
    isPrimeExact,
    nextPrime,
    millerRabin,
    bailliePSW,
    sieveEratosthenes,
    sieveAtkin,
    segmentedSieve,
    primeFactorization,
    pollardRho,
    pollardPMinusOne,
    fermatFactor,
    divisors,
    divisorSum,
    divisorCount,
    eulerTotient,
    carmichael,
    mobius,
    liouville,
    multiplicativeOrder,
    primitiveRoot,
    discreteLog,
    legendreSymbol,
    jacobiSymbol,
    kroneckerSymbol,
    tonelliShanks,
    continuedFraction,
    fromContinuedFraction,
    convergents,
    farey,
    pellFundamental,
    isPerfectNumber,
    areAmicable,
    partitions,
    partitionCount,
    fibonacciFast,
) where

import Data.List (group, sort)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import Data.Ratio (denominator, numerator, (%))
import qualified Data.Set as Set
import HMath.Error (MathError (..))

-- | Nonnegative greatest common divisor; @gcdInteger 0 0 = 0@.
gcdInteger :: Integer -> Integer -> Integer
gcdInteger :: Integer -> Integer -> Integer
gcdInteger = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd

-- | Nonnegative least common multiple; zero if either argument is zero.
lcmInteger :: Integer -> Integer -> Integer
lcmInteger :: Integer -> Integer -> Integer
lcmInteger = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
lcm

-- | @(g,x,y)@ with @a*x+b*y=g>=0@, also for negative inputs and zero.
extendedGCD :: Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD :: Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD Integer
a Integer
b =
    let (Integer
g, Integer
x, Integer
y) = Integer -> Integer -> (Integer, Integer, Integer)
forall {c}. Integral c => c -> c -> (c, c, c)
go (Integer -> Integer
forall a. Num a => a -> a
abs Integer
a) (Integer -> Integer
forall a. Num a => a -> a
abs Integer
b)
     in (Integer
g, Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Integer
forall a. Num a => a -> a
signum Integer
a, Integer
y Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer -> Integer
forall a. Num a => a -> a
signum Integer
b)
  where
    go :: c -> c -> (c, c, c)
go c
u c
0 = (c
u, c
1, c
0)
    go c
u c
v =
        let (c
q, c
r) = c
u c -> c -> (c, c)
forall a. Integral a => a -> a -> (a, a)
`divMod` c
v; (c
g, c
x, c
y) = c -> c -> (c, c, c)
go c
v c
r
         in (c
g, c
y, c
x c -> c -> c
forall a. Num a => a -> a -> a
- c
q c -> c -> c
forall a. Num a => a -> a -> a
* c
y)

-- | Bezout coefficients only; see 'extendedGCD'.
bezout :: Integer -> Integer -> (Integer, Integer)
bezout :: Integer -> Integer -> (Integer, Integer)
bezout Integer
a Integer
b = let (Integer
_, Integer
x, Integer
y) = Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD Integer
a Integer
b in (Integer
x, Integer
y)

{- | Binary modular exponentiation. Modulus must be positive. Negative powers
use a modular inverse and fail when the base is not a unit.
-}
modPow :: Integer -> Integer -> Integer -> Either MathError Integer
modPow :: Integer -> Integer -> Integer -> Either MathError Integer
modPow Integer
a Integer
e Integer
m
    | Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"modulus must be positive")
    | Integer
e Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 = do
        v <- Integer -> Integer -> Either MathError Integer
modInverse Integer
a Integer
m; pure (powMod v (-e) m)
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right (Integer -> Integer -> Integer -> Integer
powMod Integer
a Integer
e Integer
m)

powMod :: Integer -> Integer -> Integer -> Integer
powMod :: Integer -> Integer -> Integer -> Integer
powMod Integer
a Integer
e Integer
m = Integer -> Integer -> Integer -> Integer
forall {a}. Integral a => Integer -> a -> Integer -> Integer
go (Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m) Integer
e (Integer
1 Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m)
  where
    go :: Integer -> a -> Integer -> Integer
go Integer
_ a
0 Integer
r = Integer
r
    go Integer
x a
k Integer
r =
        Integer -> a -> Integer -> Integer
go
            ((Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m)
            (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2)
            (if a -> Bool
forall a. Integral a => a -> Bool
odd a
k then (Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m else Integer
r)

-- | Canonical inverse residue in @[0,m)@. Modulus must exceed one.
modInverse :: Integer -> Integer -> Either MathError Integer
modInverse :: Integer -> Integer -> Either MathError Integer
modInverse Integer
a Integer
m
    | Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"inverse modulus must exceed one")
    | Integer
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"residue has no modular inverse")
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right (Integer
x Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m)
  where
    (Integer
g, Integer
x, Integer
_) = Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD Integer
a Integer
m

{- | Generalized CRT for @(residue,positive modulus)@ pairs. Compatible
noncoprime moduli are supported. Returns canonical residue and lcm modulus.
-}
chineseRemainder :: [(Integer, Integer)] -> Either MathError (Integer, Integer)
chineseRemainder :: [(Integer, Integer)] -> Either MathError (Integer, Integer)
chineseRemainder = (Either MathError (Integer, Integer)
 -> (Integer, Integer) -> Either MathError (Integer, Integer))
-> Either MathError (Integer, Integer)
-> [(Integer, Integer)]
-> Either MathError (Integer, Integer)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' Either MathError (Integer, Integer)
-> (Integer, Integer) -> Either MathError (Integer, Integer)
combine ((Integer, Integer) -> Either MathError (Integer, Integer)
forall a b. b -> Either a b
Right (Integer
0, Integer
1))
  where
    combine :: Either MathError (Integer, Integer)
-> (Integer, Integer) -> Either MathError (Integer, Integer)
combine Either MathError (Integer, Integer)
previous (Integer
b, Integer
n) = do
        (a, m) <- Either MathError (Integer, Integer)
previous
        if n <= 0
            then Left (DomainError "CRT moduli must be positive")
            else do
                let (g, u, _) = extendedGCD m n
                if (b - a) `mod` g /= 0
                    then Left (DomainError "inconsistent congruences")
                    else
                        let modulus = Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g)
                            k = (((Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
a) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
u) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` (Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g)
                         in Right ((a + m * k) `mod` modulus, modulus)

{- | One solution and a direction: @(x,y,dx,dy)@ represents every solution
@(x+k*dx,y+k*dy)@ of @a*x+b*y=c@. The all-zero equation has a two-dimensional
solution set and is rejected instead of returning a misleading direction.
-}
linearDiophantine ::
    Integer ->
    Integer ->
    Integer ->
    Either MathError (Integer, Integer, Integer, Integer)
linearDiophantine :: Integer
-> Integer
-> Integer
-> Either MathError (Integer, Integer, Integer, Integer)
linearDiophantine Integer
a Integer
b Integer
c
    | Integer
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = MathError -> Either MathError (Integer, Integer, Integer, Integer)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"degenerate Diophantine equation")
    | Integer
c Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
g Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 = MathError -> Either MathError (Integer, Integer, Integer, Integer)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"no integer solution")
    | Bool
otherwise = (Integer, Integer, Integer, Integer)
-> Either MathError (Integer, Integer, Integer, Integer)
forall a b. b -> Either a b
Right (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
c Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g), Integer
y Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
c Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g), Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g, -Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
g)
  where
    (Integer
g, Integer
x, Integer
y) = Integer -> Integer -> (Integer, Integer, Integer)
extendedGCD Integer
a Integer
b

-- | Floor of the square root, computed by exact integer Newton iteration.
integerSquareRoot :: Integer -> Either MathError Integer
integerSquareRoot :: Integer -> Either MathError Integer
integerSquareRoot Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"negative integer square root")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
n
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right (Integer -> Integer
go Integer
n)
  where
    go :: Integer -> Integer
go Integer
x = let y :: Integer
y = (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
x) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2 in if Integer
y Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
x then Integer
x else Integer -> Integer
go Integer
y

rootFloor :: Int -> Integer -> Integer
rootFloor :: Int -> Integer -> Integer
rootFloor Int
k Integer
n = Integer -> Integer -> Integer
search Integer
0 (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
  where
    search :: Integer -> Integer -> Integer
search Integer
a Integer
b
        | Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 = Integer
a
        | Bool
otherwise =
            let c :: Integer
c = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2
             in if Integer
c Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
n then Integer -> Integer -> Integer
search Integer
c Integer
b else Integer -> Integer -> Integer
search Integer
a Integer
c

{- | Find a representation @n = base^exponent@ with exponent greater than one.
Zero, one, and minus one use exponent 2, 2, and 3 respectively.
-}
perfectPower :: Integer -> Maybe (Integer, Int)
perfectPower :: Integer -> Maybe (Integer, Int)
perfectPower Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = (Integer, Int) -> Maybe (Integer, Int)
forall a. a -> Maybe a
Just (Integer
0, Int
2)
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = (Integer, Int) -> Maybe (Integer, Int)
forall a. a -> Maybe a
Just (Integer
1, Int
2)
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== -Integer
1 = (Integer, Int) -> Maybe (Integer, Int)
forall a. a -> Maybe a
Just (-Integer
1, Int
3)
    | Bool
otherwise = Int -> Maybe (Integer, Int)
seek Int
2
  where
    cap :: Int
cap = Integer -> Int
bitLength (Integer -> Integer
forall a. Num a => a -> a
abs Integer
n)
    seek :: Int -> Maybe (Integer, Int)
seek Int
k
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cap = Maybe (Integer, Int)
forall a. Maybe a
Nothing
        | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
&& Int -> Bool
forall a. Integral a => a -> Bool
even Int
k = Int -> Maybe (Integer, Int)
seek (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise =
            let r :: Integer
r = Int -> Integer -> Integer
rootFloor Int
k (Integer -> Integer
forall a. Num a => a -> a
abs Integer
n)
             in if Integer
r Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer -> Integer
forall a. Num a => a -> a
abs Integer
n
                    then (Integer, Int) -> Maybe (Integer, Int)
forall a. a -> Maybe a
Just (if Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then -Integer
r else Integer
r, Int
k)
                    else Int -> Maybe (Integer, Int)
seek (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

bitLength :: Integer -> Int
bitLength :: Integer -> Int
bitLength = Int -> Integer -> Int
forall {t} {t}. (Num t, Integral t) => t -> t -> t
go Int
0
  where
    go :: t -> t -> t
go t
k t
0 = t
k; go t
k t
n = t -> t -> t
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) (t
n t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
2)

-- | Proven primality using trial division. The budget counts trial divisors.
isPrimeExact :: Int -> Integer -> Either MathError Bool
isPrimeExact :: Int -> Integer -> Either MathError Bool
isPrimeExact Int
budget Integer
n
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Bool
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative primality budget")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
True
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Bool
otherwise = Int -> Integer -> Either MathError Bool
forall {t}. (Eq t, Num t) => t -> Integer -> Either MathError Bool
go Int
budget Integer
3
  where
    go :: t -> Integer -> Either MathError Bool
go t
left Integer
d
        | Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
n = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
True
        | t
left t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = MathError -> Either MathError Bool
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"trial primality" Int
budget)
        | Integer
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
d Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
        | Bool
otherwise = t -> Integer -> Either MathError Bool
go (t
left t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
2)

{- | First proven prime at least the requested integer. The budget bounds both
candidate count and trial divisions per candidate.
-}
nextPrime :: Int -> Integer -> Either MathError Integer
nextPrime :: Int -> Integer -> Either MathError Integer
nextPrime Int
budget Integer
n
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative prime-search budget")
    | Bool
otherwise = Int -> Integer -> Either MathError Integer
forall {t}.
(Eq t, Num t) =>
t -> Integer -> Either MathError Integer
go Int
budget (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
2 Integer
n)
  where
    go :: t -> Integer -> Either MathError Integer
go t
0 Integer
_ = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"prime search" Int
budget)
    go t
left Integer
k = do
        p <- Int -> Integer -> Either MathError Bool
isPrimeExact Int
budget Integer
k
        if p then Right k else go (left - 1) (k + 1)

{- | Strong Miller--Rabin for explicitly supplied bases (each >=2). Passing is
probable primality, not a proof. No random-base probability bound is claimed
for caller-selected deterministic bases. An empty base list is rejected.
-}
millerRabin :: [Integer] -> Integer -> Either MathError Bool
millerRabin :: [Integer] -> Integer -> Either MathError Bool
millerRabin [Integer]
bases Integer
n
    | [Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
bases Bool -> Bool -> Bool
|| (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2) [Integer]
bases = MathError -> Either MathError Bool
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"Miller-Rabin needs bases >= 2")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 Bool -> Bool -> Bool
|| Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
True
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Bool
otherwise = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right ((Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Integer -> Bool
strong [Integer]
bases)
  where
    (Int
s, Integer
d) = Integer -> (Int, Integer)
splitTwos (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)
    strong :: Integer -> Bool
strong Integer
a
        | Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = Bool
True
        | Bool
otherwise =
            let x :: Integer
x = Integer -> Integer -> Integer -> Integer
powMod Integer
a Integer
d Integer
n
             in Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1
                    Bool -> Bool -> Bool
|| (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
                        (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)
                        (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
s ((Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (\Integer
z -> Integer
z Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
z Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) Integer
x))

splitTwos :: Integer -> (Int, Integer)
splitTwos :: Integer -> (Int, Integer)
splitTwos = Int -> Integer -> (Int, Integer)
forall {b} {a}. (Integral b, Num a) => a -> b -> (a, b)
go Int
0
  where
    go :: a -> b -> (a, b)
go a
s b
d
        | b -> Bool
forall a. Integral a => a -> Bool
even b
d Bool -> Bool -> Bool
&& b
d b -> b -> Bool
forall a. Eq a => a -> a -> Bool
/= b
0 = a -> b -> (a, b)
go (a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (b
d b -> b -> b
forall a. Integral a => a -> a -> a
`div` b
2)
        | Bool
otherwise = (a
s, b
d)

{- | Baillie--PSW: base-2 strong Miller--Rabin plus strong Lucas--Selfridge.
Perfect squares are rejected. Passing is probable primality for arbitrary
integers; no unproved claim of deterministic primality is made. Selfridge
parameter search is limited to 10000 discriminants.
-}
bailliePSW :: Integer -> Either MathError Bool
bailliePSW :: Integer -> Either MathError Bool
bailliePSW Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
True
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
    | Bool
otherwise = do
        r <- Integer -> Either MathError Integer
integerSquareRoot Integer
n
        mr <- millerRabin [2] n
        if r * r == n || not mr then Right False else choose 0 5
  where
    choose :: Int -> Integer -> Either MathError Bool
    choose :: Int -> Integer -> Either MathError Bool
choose Int
k Integer
d
        | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
10000 = MathError -> Either MathError Bool
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Lucas-Selfridge parameters" Int
k)
        | Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
n (Integer -> Integer
forall a. Num a => a -> a
abs Integer
d) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
n (Integer -> Integer
forall a. Num a => a -> a
abs Integer
d) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n = Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
False
        | Bool
otherwise = do
            j <- Integer -> Integer -> Either MathError Integer
jacobiSymbol Integer
d Integer
n
            if j == -1
                then Right (strongLucas d)
                else choose (k + 1) (if d > 0 then -d - 2 else -d + 2)
    strongLucas :: Integer -> Bool
strongLucas Integer
d =
        let q :: Integer
q = (Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
d) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
4
            (Int
s, Integer
k) = Integer -> (Int, Integer)
splitTwos (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
            (Integer
u, Integer
v) = Integer -> Integer -> Integer -> Integer -> (Integer, Integer)
lucasUV Integer
1 Integer
q Integer
k Integer
n
            vs :: [Integer]
vs =
                Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take
                    Int
s
                    ( ((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) -> [(Integer, Integer)]
forall a. (a -> a) -> a -> [a]
iterate
                            ( \(Integer
x, Integer
z) ->
                                ((Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
z) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n, Integer
z Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
z Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n)
                            )
                            (Integer
v, Integer -> Integer -> Integer -> Integer
powMod Integer
q Integer
k Integer
n)
                        )
                    )
         in Integer
u Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 Bool -> Bool -> Bool
|| (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0) [Integer]
vs

lucasUV :: Integer -> Integer -> Integer -> Integer -> (Integer, Integer)
lucasUV :: Integer -> Integer -> Integer -> Integer -> (Integer, Integer)
lucasUV Integer
p Integer
q Integer
k Integer
m =
    let (Integer
a, Integer
_, Integer
c, Integer
_) = (Integer, Integer, Integer, Integer)
-> Integer -> (Integer, Integer, Integer, Integer)
forall {a}.
Integral a =>
(Integer, Integer, Integer, Integer)
-> a -> (Integer, Integer, Integer, Integer)
matrixPower (Integer
p, -Integer
q, Integer
1, Integer
0) Integer
k
     in (Integer
c, (Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
c) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m)
  where
    mul :: (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
mul (Integer
a, Integer
b, Integer
c, Integer
d) (Integer
e, Integer
f, Integer
g, Integer
h) =
        ( (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
e Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
g) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
        , (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
f Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
h) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
        , (Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
e Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
g) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
        , (Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
f Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
h) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m
        )
    matrixPower :: (Integer, Integer, Integer, Integer)
-> a -> (Integer, Integer, Integer, Integer)
matrixPower (Integer, Integer, Integer, Integer)
a a
e = (Integer, Integer, Integer, Integer)
-> a
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
forall {a}.
Integral a =>
(Integer, Integer, Integer, Integer)
-> a
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
go (Integer, Integer, Integer, Integer)
a a
e (Integer
1, Integer
0, Integer
0, Integer
1)
    go :: (Integer, Integer, Integer, Integer)
-> a
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
go (Integer, Integer, Integer, Integer)
_ a
0 (Integer, Integer, Integer, Integer)
r = (Integer, Integer, Integer, Integer)
r
    go (Integer, Integer, Integer, Integer)
a a
e (Integer, Integer, Integer, Integer)
r = (Integer, Integer, Integer, Integer)
-> a
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
go ((Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
mul (Integer, Integer, Integer, Integer)
a (Integer, Integer, Integer, Integer)
a) (a
e a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2) (if a -> Bool
forall a. Integral a => a -> Bool
odd a
e then (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
-> (Integer, Integer, Integer, Integer)
mul (Integer, Integer, Integer, Integer)
r (Integer, Integer, Integer, Integer)
a else (Integer, Integer, Integer, Integer)
r)

-- | Eratosthenes sieve through an inclusive nonnegative upper bound.
sieveEratosthenes :: Int -> [Integer]
sieveEratosthenes :: Int -> [Integer]
sieveEratosthenes Int
bound = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int] -> [Int]
go [Int
2 .. Int
bound])
  where
    go :: [Int] -> [Int]
go [] = []
    go (Int
p : [Int]
xs)
        | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bound Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
p = Int
p Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
xs
        | Bool
otherwise = Int
p Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int] -> [Int]
go ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Int
x -> Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) [Int]
xs)

-- | Atkin quadratic-residue sieve, followed by removal of prime-square multiples.
sieveAtkin :: Int -> [Integer]
sieveAtkin :: Int -> [Integer]
sieveAtkin Int
bound
    | Int
bound Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = []
    | Bool
otherwise = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bound) [Int
2, Int
3] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Set Int -> [Int]
forall a. Set a -> [a]
Set.toAscList Set Int
cleaned)
  where
    limit :: Int
limit = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Double
forall a. Floating a => a -> a
sqrt (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bound :: Double))
    flipMember :: a -> Set a -> Set a
flipMember a
x Set a
s = if a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member a
x Set a
s then a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
x Set a
s else a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x Set a
s
    toggles :: [Int]
toggles =
        [ Int
n
        | Int
x <- [Int
1 .. Int
limit]
        , Int
y <- [Int
1 .. Int
limit]
        , Int
n <- [Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y]
        , Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bound
        , Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
12 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int
1, Int
5]
        ]
            [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [ Int
n
               | Int
x <- [Int
1 .. Int
limit]
               , Int
y <- [Int
1 .. Int
limit]
               , Int
n <- [Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y]
               , Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bound
               , Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
12 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
7
               ]
            [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [ Int
n
               | Int
x <- [Int
1 .. Int
limit]
               , Int
y <- [Int
1 .. Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
               , Int
n <- [Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y]
               , Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bound
               , Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
12 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
11
               ]
    candidates :: Set Int
candidates = (Set Int -> Int -> Set Int) -> Set Int -> [Int] -> Set Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' ((Int -> Set Int -> Set Int) -> Set Int -> Int -> Set Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
flipMember) Set Int
forall a. Set a
Set.empty [Int]
toggles
    cleaned :: Set Int
cleaned =
        (Set Int -> Int -> Set Int) -> Set Int -> [Int] -> Set Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl'
            ( \Set Int
s Int
p ->
                if Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
p Set Int
s
                    then (Set Int -> Int -> Set Int) -> Set Int -> [Int] -> Set Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' ((Int -> Set Int -> Set Int) -> Set Int -> Int -> Set Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.delete) Set Int
s [Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p, Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
p .. Int
bound]
                    else Set Int
s
            )
            Set Int
candidates
            [Int
5 .. Int
limit]

{- | Segmented sieve on the inclusive interval @[lo,hi]@. The segment width and
the square-root base sieve must fit an Int; reversed intervals are rejected.
-}
segmentedSieve :: Integer -> Integer -> Either MathError [Integer]
segmentedSieve :: Integer -> Integer -> Either MathError [Integer]
segmentedSieve Integer
lo Integer
hi
    | Integer
hi Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
lo = MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"reversed sieve interval")
    | Integer
hi Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2 = [Integer] -> Either MathError [Integer]
forall a b. b -> Either a b
Right []
    | Integer
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) = MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sieve segment too large")
    | Bool
otherwise = do
        r <- Integer -> Either MathError Integer
integerSquareRoot Integer
hi
        if r > toInteger (maxBound :: Int)
            then Left (InvalidArgument "base sieve too large")
            else
                let low = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
2 Integer
lo
                    remove Set Integer
s Integer
p =
                        (Set Integer -> Integer -> Set Integer)
-> Set Integer -> [Integer] -> Set Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl'
                            ((Integer -> Set Integer -> Set Integer)
-> Set Integer -> Integer -> Set Integer
forall a b c. (a -> b -> c) -> b -> a -> c
flip Integer -> Set Integer -> Set Integer
forall a. Ord a => a -> Set a -> Set a
Set.delete)
                            Set Integer
s
                            [ Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max (Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p) (((Integer
low Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
p) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p)
                            , Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max (Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p) (((Integer
low Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
p) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
p
                            .. Integer
hi
                            ]
                 in Right
                        ( Set.toAscList
                            ( List.foldl'
                                remove
                                (Set.fromList [low .. hi])
                                (sieveEratosthenes (fromInteger r))
                            )
                        )

{- | Exact trial factorization of a positive integer. The budget counts attempted
divisors (including repeated factors). One has an empty factorization.
-}
primeFactorization :: Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization :: Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError [(Integer, Int)]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative factorization budget")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1 = MathError -> Either MathError [(Integer, Int)]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"factorization requires a positive integer")
    | Bool
otherwise =
        ([Integer] -> [(Integer, Int)])
-> Either MathError [Integer] -> Either MathError [(Integer, Int)]
forall a b. (a -> b) -> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Integer] -> (Integer, Int)) -> [[Integer]] -> [(Integer, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\[Integer]
xs -> (case [Integer]
xs of Integer
x : [Integer]
_ -> Integer
x; [] -> Integer
1, [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs)) ([[Integer]] -> [(Integer, Int)])
-> ([Integer] -> [[Integer]]) -> [Integer] -> [(Integer, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [[Integer]]
forall a. Eq a => [a] -> [[a]]
group) (Int -> Integer -> Integer -> Either MathError [Integer]
forall {a} {t}.
(Eq t, Num t, Integral a) =>
t -> a -> a -> Either MathError [a]
go Int
budget Integer
n Integer
2)
  where
    go :: t -> a -> a -> Either MathError [a]
go t
_ a
1 a
_ = [a] -> Either MathError [a]
forall a b. b -> Either a b
Right []
    go t
left a
k a
d
        | a
d a -> a -> a
forall a. Num a => a -> a -> a
* a
d a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
k = [a] -> Either MathError [a]
forall a b. b -> Either a b
Right [a
k]
        | t
left t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = MathError -> Either MathError [a]
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"trial factorization" Int
budget)
        | a
k a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = (a
d a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> Either MathError [a] -> Either MathError [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> a -> a -> Either MathError [a]
go (t
left t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
d) a
d
        | Bool
otherwise = t -> a -> a -> Either MathError [a]
go (t
left t -> t -> t
forall a. Num a => a -> a -> a
- t
1) a
k (if a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
2 then a
3 else a
d a -> a -> a
forall a. Num a => a -> a -> a
+ a
2)

{- | Pollard rho with Floyd cycle detection, explicit seed and iteration budget.
Returns a proper divisor or failure; a failed seed does not prove primality.
-}
pollardRho :: Int -> Integer -> Integer -> Either MathError Integer
pollardRho :: Int -> Integer -> Integer -> Either MathError Integer
pollardRho Int
budget Integer
seed Integer
n
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative rho budget")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
4 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"rho requires a composite candidate >=4")
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
2
    | Bool
otherwise = Int -> Integer -> Integer -> Either MathError Integer
forall {t}.
(Eq t, Num t) =>
t -> Integer -> Integer -> Either MathError Integer
go Int
budget (Integer
seed Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) (Integer
seed Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n)
  where
    f :: Integer -> Integer
f Integer
x = (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n
    go :: t -> Integer -> Integer -> Either MathError Integer
go t
0 Integer
_ Integer
_ = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Pollard rho" Int
budget)
    go t
left Integer
x Integer
y =
        let a :: Integer
a = Integer -> Integer
f Integer
x; b :: Integer
b = Integer -> Integer
f (Integer -> Integer
f Integer
y); d :: Integer
d = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd (Integer -> Integer
forall a. Num a => a -> a
abs (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b)) Integer
n
         in if Integer
d Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1
                then t -> Integer -> Integer -> Either MathError Integer
go (t
left t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Integer
a Integer
b
                else
                    if Integer
d Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n
                        then MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Pollard rho cycle" Int
budget)
                        else Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
d

{- | Pollard p-1 stage one with a smoothness bound >=2 and base. The exponent
includes the maximal powers of every prime <= the bound.
-}
pollardPMinusOne :: Int -> Integer -> Integer -> Either MathError Integer
pollardPMinusOne :: Int -> Integer -> Integer -> Either MathError Integer
pollardPMinusOne Int
bound Integer
base Integer
n
    | Int
bound Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 Bool -> Bool -> Bool
|| Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
4 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"invalid p-1 bound or candidate")
    | Integer
g Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
g Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
g
    | Bool
otherwise =
        let power :: Integer
power = [Integer] -> Integer
lcmList [Integer
1 .. Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
bound]
            d :: Integer
d = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd (Integer -> Integer -> Integer -> Integer
powMod Integer
base Integer
power Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer
n
         in if Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
1 Bool -> Bool -> Bool
&& Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
n
                then Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
d
                else MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Pollard p-1" Int
bound)
  where
    g :: Integer
g = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
base Integer
n
    lcmList :: [Integer] -> Integer
lcmList = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
lcm Integer
1

-- | Fermat difference-of-squares factorization, with an explicit step budget.
fermatFactor :: Int -> Integer -> Either MathError Integer
fermatFactor :: Int -> Integer -> Either MathError Integer
fermatFactor Int
budget Integer
n
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative Fermat budget")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
4 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Fermat requires candidate >=4")
    | Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
2
    | Bool
otherwise = do
        r <- Integer -> Either MathError Integer
integerSquareRoot Integer
n
        go budget (if r * r == n then r else r + 1)
  where
    go :: t -> Integer -> Either MathError Integer
go t
0 Integer
_ = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Fermat factorization" Int
budget)
    go t
left Integer
a = do
        b <- Integer -> Either MathError Integer
integerSquareRoot (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
n)
        if b * b == a * a - n && a - b > 1
            then Right (a - b)
            else go (left - 1) (a + 1)

-- | All positive divisors in increasing order, under a factorization budget.
divisors :: Int -> Integer -> Either MathError [Integer]
divisors :: Int -> Integer -> Either MathError [Integer]
divisors Int
budget Integer
n = do
    fs <- Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
    pure (sort (List.foldl' (\[Integer]
xs (Integer
p, Int
k) -> [Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
p Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
j | Integer
x <- [Integer]
xs, Int
j <- [Int
0 .. Int
k]]) [1] fs))

-- | Sum of positive divisors.
divisorSum :: Int -> Integer -> Either MathError Integer
divisorSum :: Int -> Integer -> Either MathError Integer
divisorSum Int
budget Integer
n = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer)
-> Either MathError [Integer] -> Either MathError Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Integer -> Either MathError [Integer]
divisors Int
budget Integer
n

-- | Number of positive divisors.
divisorCount :: Int -> Integer -> Either MathError Integer
divisorCount :: Int -> Integer -> Either MathError Integer
divisorCount Int
budget Integer
n = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Integer] -> Integer)
-> ([(Integer, Int)] -> [Integer]) -> [(Integer, Int)] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer)
-> ((Integer, Int) -> Int) -> (Integer, Int) -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int) -> ((Integer, Int) -> Int) -> (Integer, Int) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Integer, Int)] -> Integer)
-> Either MathError [(Integer, Int)] -> Either MathError Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n

-- | Euler totient, including @phi(1)=1@.
eulerTotient :: Int -> Integer -> Either MathError Integer
eulerTotient :: Int -> Integer -> Either MathError Integer
eulerTotient Int
budget Integer
n = do
    fs <- Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
    pure (List.foldl' (\Integer
a (Integer
p, Int
_) -> Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)) n fs)

-- | Carmichael's exponent of the multiplicative unit group.
carmichael :: Int -> Integer -> Either MathError Integer
carmichael :: Int -> Integer -> Either MathError Integer
carmichael Int
budget Integer
n = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
lcm Integer
1 ([Integer] -> Integer)
-> ([(Integer, Int)] -> [Integer]) -> [(Integer, Int)] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, Int) -> Integer) -> [(Integer, Int)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Int) -> Integer
forall {b} {a}. (Num a, Integral b, Eq a) => (a, b) -> a
term ([(Integer, Int)] -> Integer)
-> Either MathError [(Integer, Int)] -> Either MathError Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
  where
    term :: (a, b) -> a
term (a
2, b
k) | b
k b -> b -> Bool
forall a. Ord a => a -> a -> Bool
>= b
3 = a
2 a -> b -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (b
k b -> b -> b
forall a. Num a => a -> a -> a
- b
2)
    term (a
p, b
k) = (a
p a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Num a => a -> a -> a
* a
p a -> b -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (b
k b -> b -> b
forall a. Num a => a -> a -> a
- b
1)

-- | Mobius function on positive integers.
mobius :: Int -> Integer -> Either MathError Integer
mobius :: Int -> Integer -> Either MathError Integer
mobius Int
budget Integer
n = do
    fs <- Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
    pure (if any ((> 1) . snd) fs then 0 else if even (length fs) then 1 else -1)

-- | Liouville function, counting prime factors with multiplicity.
liouville :: Int -> Integer -> Either MathError Integer
liouville :: Int -> Integer -> Either MathError Integer
liouville Int
budget Integer
n = do
    fs <- Int -> Integer -> Either MathError [(Integer, Int)]
primeFactorization Int
budget Integer
n
    pure (if even (sum (map snd fs)) then 1 else -1)

-- | Multiplicative order of a unit modulo a modulus >1.
multiplicativeOrder :: Int -> Integer -> Integer -> Either MathError Integer
multiplicativeOrder :: Int -> Integer -> Integer -> Either MathError Integer
multiplicativeOrder Int
budget Integer
a Integer
m
    | Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 Bool -> Bool -> Bool
|| Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
a Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"order requires a unit modulo m>1")
    | Bool
otherwise = do
        phi <- Int -> Integer -> Either MathError Integer
eulerTotient Int
budget Integer
m
        ds <- divisors budget phi
        case filter (\Integer
d -> Integer -> Integer -> Integer -> Integer
powMod Integer
a Integer
d Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1) ds of
            Integer
d : [Integer]
_ -> Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
d
            [] -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"multiplicative order" Int
budget)

{- | Smallest generator of the units modulo m, if that group is cyclic. Search
and factorization are both bounded; noncyclic groups report a domain error.
-}
primitiveRoot :: Int -> Integer -> Either MathError Integer
primitiveRoot :: Int -> Integer -> Either MathError Integer
primitiveRoot Int
budget Integer
m = do
    phi <- Int -> Integer -> Either MathError Integer
eulerTotient Int
budget Integer
m
    if m <= 1
        then Left (DomainError "primitive root modulus must exceed one")
        else do
            fs <- primeFactorization budget phi
            let candidates = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
budget [Integer
1 .. Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
                valid Integer
a = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
a Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 Bool -> Bool -> Bool
&& ((Integer, Int) -> Bool) -> [(Integer, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Integer
p, Int
_) -> Integer -> Integer -> Integer -> Integer
powMod Integer
a (Integer
phi Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
p) Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1) [(Integer, Int)]
fs
            case filter valid candidates of
                Integer
a : [Integer]
_ -> Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
a
                []
                    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
budget Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"unit group has no primitive root")
                    | Bool
otherwise -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"primitive root search" Int
budget)

{- | Baby-step giant-step discrete logarithm of a unit. Budget bounds the table
dimension, so storage is O(budget) and work O(budget*log budget), apart
from integer arithmetic. The smallest nonnegative exponent is returned.
-}
discreteLog :: Int -> Integer -> Integer -> Integer -> Either MathError Integer
discreteLog :: Int -> Integer -> Integer -> Integer -> Either MathError Integer
discreteLog Int
budget Integer
a Integer
b Integer
m
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative discrete-log budget")
    | Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
1 Bool -> Bool -> Bool
|| Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
a Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1 Bool -> Bool -> Bool
|| Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
gcd Integer
b Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
1 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"discrete log needs units modulo m>1")
    | Bool
otherwise = do
        r <- Integer -> Either MathError Integer
integerSquareRoot Integer
m
        let width = Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
        if width > toInteger budget
            then Left (NonConvergence "discrete logarithm" budget)
            else do
                inv <- modInverse (powMod a width m) m
                let table =
                        (Integer -> Integer -> Integer)
-> [(Integer, Integer)] -> Map Integer Integer
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith
                            Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min
                            ( [Integer] -> [Integer] -> [(Integer, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip
                                ( Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take
                                    (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
width)
                                    ((Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (\Integer
x -> Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m) Integer
1)
                                )
                                [Integer
0 ..]
                            )
                    matches =
                        [ Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
width Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
j
                        | (Integer
i, Integer
v) <- [Integer] -> [Integer] -> [(Integer, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0 .. Integer
width] ((Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (\Integer
x -> Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
inv Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m) (Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
m))
                        , Just Integer
j <- [Integer -> Map Integer Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Integer
v Map Integer Integer
table]
                        ]
                case matches of
                    [] -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"discrete logarithm does not exist")
                    [Integer]
xs -> Integer -> Either MathError Integer
forall a b. b -> Either a b
Right ([Integer] -> Integer
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Integer]
xs)

-- | Legendre symbol for a proven odd prime modulus, with proof budget.
legendreSymbol :: Int -> Integer -> Integer -> Either MathError Integer
legendreSymbol :: Int -> Integer -> Integer -> Either MathError Integer
legendreSymbol Int
budget Integer
a Integer
p = do
    prime <- Int -> Integer -> Either MathError Bool
isPrimeExact Int
budget Integer
p
    if not prime || p == 2
        then Left (DomainError "Legendre modulus must be an odd prime")
        else let x = Integer -> Integer -> Integer -> Integer
powMod Integer
a ((Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2) Integer
p in Right (if x == p - 1 then -1 else x)

-- | Jacobi symbol for a positive odd denominator.
jacobiSymbol :: Integer -> Integer -> Either MathError Integer
jacobiSymbol :: Integer -> Integer -> Either MathError Integer
jacobiSymbol Integer
a Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 Bool -> Bool -> Bool
|| Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Jacobi denominator must be positive odd")
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right (Integer -> Integer -> Integer -> Integer
forall {t}. Num t => Integer -> Integer -> t -> t
go (Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) Integer
n Integer
1)
  where
    go :: Integer -> Integer -> t -> t
go Integer
0 Integer
m t
r = if Integer
m Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then t
r else t
0
    go Integer
x Integer
m t
r =
        let (Int
s, Integer
u) = Integer -> (Int, Integer)
splitTwos Integer
x
            r1 :: t
r1 = if Int -> Bool
forall a. Integral a => a -> Bool
odd Int
s Bool -> Bool -> Bool
&& Integer
m Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
8 Integer -> [Integer] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Integer
3, Integer
5] then -t
r else t
r
            r2 :: t
r2 = if Integer
u Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3 Bool -> Bool -> Bool
&& Integer
m Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
3 then -t
r1 else t
r1
         in Integer -> Integer -> t -> t
go (Integer
m Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
u) Integer
u t
r2

-- | Kronecker symbol, extending Jacobi to every integer denominator.
kroneckerSymbol :: Integer -> Integer -> Integer
kroneckerSymbol :: Integer -> Integer -> Integer
kroneckerSymbol Integer
a Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 = if Integer -> Integer
forall a. Num a => a -> a
abs Integer
a Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then Integer
1 else Integer
0
    | Bool
otherwise =
        let sign :: Integer
sign = if Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
&& Integer
a Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then -Integer
1 else Integer
1
            (Int
s, Integer
u) = Integer -> (Int, Integer)
splitTwos (Integer -> Integer
forall a. Num a => a -> a
abs Integer
n)
            two :: Integer
two = if Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
a then Integer
0 else if Integer
a Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
8 Integer -> [Integer] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Integer
1, Integer
7] then Integer
1 else -Integer
1
            oddPart :: Integer
oddPart = (MathError -> Integer)
-> (Integer -> Integer) -> Either MathError Integer -> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Integer -> MathError -> Integer
forall a b. a -> b -> a
const Integer
0) Integer -> Integer
forall a. a -> a
id (Integer -> Integer -> Either MathError Integer
jacobiSymbol Integer
a Integer
u)
         in Integer
sign Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (if Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Integer
1 else Integer
two Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
s) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
oddPart

{- | Tonelli--Shanks modular square root for a proven prime modulus. Returns the
smaller of the two roots; a nonresidue is a domain error. Proof and nonresidue
search each respect the supplied budget.
-}
tonelliShanks :: Int -> Integer -> Integer -> Either MathError Integer
tonelliShanks :: Int -> Integer -> Integer -> Either MathError Integer
tonelliShanks Int
budget Integer
a Integer
p = do
    prime <- Int -> Integer -> Either MathError Bool
isPrimeExact Int
budget Integer
p
    if not prime
        then Left (DomainError "square-root modulus must be prime")
        else
            if p == 2
                then Right (a `mod` 2)
                else
                    if a `mod` p == 0
                        then Right 0
                        else do
                            symbol <- legendreSymbol budget a p
                            if symbol /= 1
                                then Left (DomainError "quadratic nonresidue")
                                else do
                                    let (s, q) = splitTwos (p - 1)
                                        nonres = [Integer
z | Integer
z <- Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
budget [Integer
2 .. Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1], Integer -> Integer -> Integer -> Integer
powMod Integer
z ((Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2) Integer
p Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
                                    case nonres of
                                        [] -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Tonelli-Shanks nonresidue" Int
budget)
                                        Integer
z : [Integer]
_ -> do
                                            r <- Int -> Integer -> Integer -> Integer -> Either MathError Integer
forall {t}.
Integral t =>
t -> Integer -> Integer -> Integer -> Either MathError Integer
loop Int
s (Integer -> Integer -> Integer -> Integer
powMod Integer
z Integer
q Integer
p) (Integer -> Integer -> Integer -> Integer
powMod Integer
a Integer
q Integer
p) (Integer -> Integer -> Integer -> Integer
powMod Integer
a ((Integer
q Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2) Integer
p)
                                            Right (min r (p - r))
  where
    loop :: t -> Integer -> Integer -> Integer -> Either MathError Integer
loop t
m Integer
c Integer
t Integer
r
        | Integer
t Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right Integer
r
        | Bool
otherwise = case [t
i | t
i <- [t
1 .. t
m t -> t -> t
forall a. Num a => a -> a -> a
- t
1], Integer -> Integer -> Integer -> Integer
powMod Integer
t (Integer
2 Integer -> t -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ t
i) Integer
p Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1] of
            [] -> MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Tonelli-Shanks" Int
budget)
            t
i : [t]
_ ->
                let b :: Integer
b = Integer -> Integer -> Integer -> Integer
powMod Integer
c (Integer
2 Integer -> t -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (t
m t -> t -> t
forall a. Num a => a -> a -> a
- t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1)) Integer
p
                 in t -> Integer -> Integer -> Integer -> Either MathError Integer
loop t
i (Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
p) (Integer
t Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
p) (Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
p)

-- | Finite simple continued fraction of a rational, using floor division.
continuedFraction :: Rational -> [Integer]
continuedFraction :: Rational -> [Integer]
continuedFraction Rational
r = Integer -> Integer -> [Integer]
forall {t}. Integral t => t -> t -> [t]
go (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r) (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r)
  where
    go :: t -> t -> [t]
go t
a t
b = let (t
q, t
s) = t
a t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
`divMod` t
b in t
q t -> [t] -> [t]
forall a. a -> [a] -> [a]
: if t
s t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 then [] else t -> t -> [t]
go t
b t
s

-- | Evaluate a finite continued fraction. Empty lists and zero tails fail.
fromContinuedFraction :: [Integer] -> Either MathError Rational
fromContinuedFraction :: [Integer] -> Either MathError Rational
fromContinuedFraction [] = MathError -> Either MathError Rational
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"empty continued fraction")
fromContinuedFraction [Integer
a] = Rational -> Either MathError Rational
forall a b. b -> Either a b
Right (Integer
a Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% Integer
1)
fromContinuedFraction (Integer
a : [Integer]
xs) = do
    r <- [Integer] -> Either MathError Rational
fromContinuedFraction [Integer]
xs
    if r == 0
        then Left (DomainError "zero continued-fraction tail")
        else Right (fromInteger a + 1 / r)

-- | Successive convergents, rejecting invalid finite prefixes.
convergents :: [Integer] -> Either MathError [Rational]
convergents :: [Integer] -> Either MathError [Rational]
convergents [Integer]
xs = (Int -> Either MathError Rational)
-> [Int] -> Either MathError [Rational]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Integer] -> Either MathError Rational
fromContinuedFraction ([Integer] -> Either MathError Rational)
-> (Int -> [Integer]) -> Int -> Either MathError Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
`take` [Integer]
xs)) [Int
1 .. [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
xs]

-- | Farey sequence of positive order, including zero and one.
farey :: Integer -> Either MathError [Rational]
farey :: Integer -> Either MathError [Rational]
farey Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1 = MathError -> Either MathError [Rational]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Farey order must be positive")
    | Bool
otherwise = [Rational] -> Either MathError [Rational]
forall a b. b -> Either a b
Right (Rational
0 Rational -> [Rational] -> [Rational]
forall a. a -> [a] -> [a]
: Integer -> Integer -> Integer -> Integer -> [Rational]
go Integer
0 Integer
1 Integer
1 Integer
n)
  where
    go :: Integer -> Integer -> Integer -> Integer -> [Rational]
go Integer
a Integer
b Integer
c Integer
d
        | Integer
c Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
d = []
        | Bool
otherwise = (Integer
c Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
% Integer
d) Rational -> [Rational] -> [Rational]
forall a. a -> [a] -> [a]
: let k :: Integer
k = (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
d in Integer -> Integer -> Integer -> Integer -> [Rational]
go Integer
c Integer
d (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
a) (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b)

-- | Minimal positive solution to @x^2-d*y^2=1@ via continued fractions of sqrt d.
pellFundamental :: Int -> Integer -> Either MathError (Integer, Integer)
pellFundamental :: Int -> Integer -> Either MathError (Integer, Integer)
pellFundamental Int
budget Integer
d
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError (Integer, Integer)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative Pell budget")
    | Integer
d Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = MathError -> Either MathError (Integer, Integer)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"Pell d must be positive nonsquare")
    | Bool
otherwise = do
        a0 <- Integer -> Either MathError Integer
integerSquareRoot Integer
d
        if a0 * a0 == d
            then Left (DomainError "Pell d must be nonsquare")
            else go budget a0 0 1 a0 1 0 0 1
  where
    go :: t
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Either MathError (Integer, Integer)
go t
0 Integer
_ Integer
_ Integer
_ Integer
_ Integer
_ Integer
_ Integer
_ Integer
_ = MathError -> Either MathError (Integer, Integer)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Pell equation" Int
budget)
    go t
left Integer
a0 Integer
m Integer
den Integer
a Integer
hm1 Integer
hm2 Integer
km1 Integer
km2 =
        let h :: Integer
h = Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
hm1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
hm2
            k :: Integer
k = Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
km1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
km2
            m' :: Integer
m' = Integer
den Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
m
            den' :: Integer
den' = (Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
m' Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
m') Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
den
            a' :: Integer
a' = (Integer
a0 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
m') Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
den'
         in if Integer
h Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
h Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1
                then (Integer, Integer) -> Either MathError (Integer, Integer)
forall a b. b -> Either a b
Right (Integer
h, Integer
k)
                else t
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Integer
-> Either MathError (Integer, Integer)
go (t
left t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Integer
a0 Integer
m' Integer
den' Integer
a' Integer
h Integer
hm1 Integer
k Integer
km1

-- | Perfect-number predicate on positive integers.
isPerfectNumber :: Int -> Integer -> Either MathError Bool
isPerfectNumber :: Int -> Integer -> Either MathError Bool
isPerfectNumber Int
budget Integer
n = (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
n) (Integer -> Bool)
-> Either MathError Integer -> Either MathError Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Integer -> Either MathError Integer
divisorSum Int
budget Integer
n

-- | Amicable-number predicate; equal numbers are not an amicable pair.
areAmicable :: Int -> Integer -> Integer -> Either MathError Bool
areAmicable :: Int -> Integer -> Integer -> Either MathError Bool
areAmicable Int
budget Integer
a Integer
b = do
    sa <- Int -> Integer -> Either MathError Integer
divisorSum Int
budget Integer
a
    sb <- divisorSum budget b
    pure (a /= b && sa - a == b && sb - b == a)

-- | Integer partitions in descending part order. Zero has one empty partition.
partitions :: Int -> Either MathError [[Int]]
partitions :: Int -> Either MathError [[Int]]
partitions Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError [[Int]]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"negative partition argument")
    | Bool
otherwise = [[Int]] -> Either MathError [[Int]]
forall a b. b -> Either a b
Right (Int -> Int -> [[Int]]
forall {a}. (Num a, Ord a, Enum a) => a -> a -> [[a]]
go Int
n Int
n)
  where
    go :: a -> a -> [[a]]
go a
0 a
_ = [[]]
    go a
k a
top = [a
p a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rest | a
p <- [a -> a -> a
forall a. Ord a => a -> a -> a
min a
top a
k, a -> a -> a
forall a. Ord a => a -> a -> a
min a
top a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
1 .. a
1], [a]
rest <- a -> a -> [[a]]
go (a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
p) a
p]

-- | Partition count by exact dynamic programming, O(n^2) work and O(n) storage.
partitionCount :: Int -> Either MathError Integer
partitionCount :: Int -> Either MathError Integer
partitionCount Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"negative partition argument")
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
drop Int
n (([Integer] -> Int -> [Integer]) -> [Integer] -> [Int] -> [Integer]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' [Integer] -> Int -> [Integer]
forall {a}. Num a => [a] -> Int -> [a]
addCoin (Integer
1 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Int -> Integer -> [Integer]
forall a. Int -> a -> [a]
replicate Int
n Integer
0) [Int
1 .. Int
n])))
  where
    addCoin :: [a] -> Int -> [a]
addCoin [a]
xs Int
k = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
k [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(+) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
k [a]
xs) [a]
result
      where
        result :: [a]
result = [a] -> Int -> [a]
addCoin [a]
xs Int
k

-- | Signed-index Fibonacci by fast doubling; @F(-n)=(-1)^(n+1)F(n)@.
fibonacciFast :: Integer -> Integer
fibonacciFast :: Integer -> Integer
fibonacciFast Integer
n
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
&& Integer -> Bool
forall a. Integral a => a -> Bool
even Integer
n = -(Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst (Integer -> (Integer, Integer)
forall {a} {b}. (Integral a, Num b) => a -> (b, b)
pair (-Integer
n))
    | Bool
otherwise = (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst (Integer -> (Integer, Integer)
forall {a} {b}. (Integral a, Num b) => a -> (b, b)
pair (Integer -> Integer
forall a. Num a => a -> a
abs Integer
n))
  where
    pair :: a -> (b, b)
pair a
0 = (b
0, b
1)
    pair a
k =
        let (b
a, b
b) = a -> (b, b)
pair (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2); c :: b
c = b
a b -> b -> b
forall a. Num a => a -> a -> a
* (b
2 b -> b -> b
forall a. Num a => a -> a -> a
* b
b b -> b -> b
forall a. Num a => a -> a -> a
- b
a); d :: b
d = b
a b -> b -> b
forall a. Num a => a -> a -> a
* b
a b -> b -> b
forall a. Num a => a -> a -> a
+ b
b b -> b -> b
forall a. Num a => a -> a -> a
* b
b
         in if a -> Bool
forall a. Integral a => a -> Bool
even a
k then (b
c, b
d) else (b
d, b
c b -> b -> b
forall a. Num a => a -> a -> a
+ b
d)