Diophantine.hs

-- Module for Diophantine Equations and related functions

module Math.NumberTheory.Diophantine
  ( cornacchiaPrimitive
  , cornacchia
  )
where

import Data.List.Infinite (Infinite(..))
import qualified Data.List.Infinite as Inf

import           Math.NumberTheory.Moduli.Sqrt  ( sqrtsModFactorisation )
import           Math.NumberTheory.Primes       ( factorise
                                                , unPrime
                                                , UniqueFactorisation
                                                )
import           Math.NumberTheory.Roots        ( integerSquareRoot )
import           Math.NumberTheory.Utils.FromIntegral

-- | See `cornacchiaPrimitive`, this is the internal algorithm implementation
-- as described at https://en.wikipedia.org/wiki/Cornacchia%27s_algorithm
cornacchiaPrimitive' :: Integer -> Integer -> [(Integer, Integer)]
cornacchiaPrimitive' d m = concatMap
  (findSolution . Inf.head . Inf.dropWhile (\r -> r * r >= m) . gcdSeq m)
  roots
 where
  roots :: [Integer]
  roots = filter (<= m `div` 2) $ sqrtsModFactorisation (m - d) (factorise m)

  gcdSeq :: Integer -> Integer -> Infinite Integer
  gcdSeq a b = a :< gcdSeq b (mod a b)

  -- If s = sqrt((m - r*r) / d) is an integer then (r, s) is a solution
  findSolution :: Integer -> [(Integer, Integer)]
  findSolution r = [ (r, s) | rem1 == 0 && s * s == s2 ]
   where
    (s2, rem1) = divMod (m - r * r) d
    s          = integerSquareRoot s2

-- | @cornacchiaPrimitive d m@ finds all primitive solutions \((x, y)\) (i.e. with \(\gcd(x, y) = 1\))
-- to the Diophantine equation
-- \[ x^2 + d \cdot y^2 = m \]
-- Preconditions: \(1 \le d < m\) and \(\gcd(d, m) = 1\).
--
-- Throws error if the preconditions are not met.
--
-- When \(m\) is square-free these are all the positive integer solutions;
-- use 'cornacchia' to find all solutions for arbitrary \(m\).
cornacchiaPrimitive :: Integer -> Integer -> [(Integer, Integer)]
cornacchiaPrimitive d m
  | not (1 <= d && d < m) = error "precondition failed: 1 <= d < m"
  | gcd d m /= 1          = error "precondition failed: d and m coprime"
  |
  -- If d=1 then the algorithm doesn't generate symmetrical pairs
    d == 1                = concatMap genPairs solutions
  | otherwise             = solutions
 where
  solutions = cornacchiaPrimitive' d m
  genPairs (x, y) = if x == y then [(x, y)] else [(x, y), (y, x)]

-- Find numbers whose square is a factor of the input
squareFactors :: UniqueFactorisation a => a -> [a]
squareFactors = foldl squareProducts [1] . factorise
 where
  squareProducts acc f = [ a * b | a <- acc, b <- squarePowers f ]
  squarePowers (p, a) = map (unPrime p ^) [0 .. wordToInt a `div` 2]

-- | @cornacchia d m@ finds all positive integer solutions \((x, y)\) to the Diophantine equation
-- \[ x^2 + d \cdot y^2 = m \]
-- Preconditions: \(1 \le d < m\) and \(\gcd(d, m) = 1\).
--
-- Throws error if the preconditions are not met.
--
-- Unlike 'cornacchiaPrimitive', this also finds non-primitive solutions (where \(\gcd(x, y) > 1\))
-- by solving the equation for each square divisor of \(m\) and scaling the results.
cornacchia :: Integer -> Integer -> [(Integer, Integer)]
cornacchia d m
  | not (1 <= d && d < m) = error "precondition failed: 1 <= d < m"
  | gcd d m /= 1 = error "precondition failed: d and m coprime"
  | otherwise = concatMap solve $ filter ((> d) . snd) candidates
 where
  candidates = map (\sf -> (sf, m `div` (sf * sf))) (squareFactors m)
  solve (sf, m') = map (\(x, y) -> (x * sf, y * sf)) (cornacchiaPrimitive d m')