linear diophantine equations

Math/NumberTheory/Diophantine.hs

@@ -1,8 +1,14 @@
 -- Module for Diophantine Equations and related functions
 
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
+
 module Math.NumberTheory.Diophantine
   ( cornacchiaPrimitive
   , cornacchia
+  , LinearSolution (..)
+  , linear
+  , runLinearSolution
   )
 where
 
@@ -14,6 +20,9 @@ import           Math.NumberTheory.Primes       ( factorise
 import           Math.NumberTheory.Roots        ( integerSquareRoot )
 import           Math.NumberTheory.Utils.FromIntegral
 
+import           Control.Monad                  (guard)
+import           Data.Euclidean                 (gcdExt)
+
 -- | See `cornacchiaPrimitive`, this is the internal algorithm implementation
 -- | as described at https://en.wikipedia.org/wiki/Cornacchia%27s_algorithm 
 cornacchiaPrimitive' :: Integer -> Integer -> [(Integer, Integer)]
@@ -64,3 +73,31 @@ cornacchia d m
  where
   candidates = map (\sf -> (sf, m `div` (sf * sf))) (squareFactors m)
   solve (sf, m') = map (\(x, y) -> (x * sf, y * sf)) (cornacchiaPrimitive d m')
+
+----
+
+-- | A solution to a linear equation
+data LinearSolution a = LS { base1,base2,scale1,scale2 :: a }
+  deriving
+    ( Show, Eq, Ord
+    )
+
+-- | Solves a linear diophantine equation
+-- |   ax + by = c
+-- | where `x` and `y` are unknown
+linear :: _ => a -> a -> a -> Maybe (LinearSolution a)
+linear a b c =
+    LS x y v u <$ guard (b /= 0 && q == 0)
+  where
+    (d, e) = gcdExt a b
+    (h, q) = divMod c d
+    f      = div (a*e-d) (-b)
+    (x, y) = (e*h, f*h)
+    (u, v) = (quot a d, quot b d)
+
+-- | Produces an unique solution given any
+-- | arbitrary number k
+runLinearSolution :: _ => LinearSolution a -> a -> (a, a)
+runLinearSolution (LS x y v u) k =
+  ( x + k*v, y - k*u )
+

test-suite/Math/NumberTheory/DiophantineTests.hs

@@ -1,6 +1,7 @@
 -- Tests for Math.NumberTheory.Diophantine
 
 {-# LANGUAGE CPP       #-}
+{-# LANGUAGE GADTs     #-}
 
 {-# OPTIONS_GHC -fno-warn-type-defaults #-}
 
@@ -15,6 +16,7 @@ import Test.Tasty
 import Math.NumberTheory.Diophantine
 import Math.NumberTheory.Roots (integerSquareRoot)
 import Math.NumberTheory.TestUtils
+import Math.NumberTheory.Primes
 
 cornacchiaTest :: Positive Integer -> Positive Integer -> Bool
 cornacchiaTest (Positive d) (Positive a) = gcd d m /= 1 || all checkSoln (cornacchia d m)
@@ -33,8 +35,30 @@ cornacchiaBruteForce (Positive d) (Positive a) = gcd d m /= 1 || findSolutions [
             where x2 = m - d*y*y
                   x = integerSquareRoot x2
 
+linearTest :: (a ~ Integer) => a -> a -> a -> a -> Bool
+linearTest a b c k =
+  case linear a b c of
+    Nothing -> True -- Disproving this would require a counter example
+    Just ls | (x, y) <- runLinearSolution ls k
+            -> a*x + b*y == c
+
+linearTest' :: (a ~ Integer) => Prime a -> Prime a -> a -> a -> Bool
+linearTest' l c' d k =
+  case linear a b c of
+    Nothing -> l == c'
+    Just ls | (x, y) <- runLinearSolution ls k
+            -> a*x + b*y == c
+  where
+    a = unPrime l
+    b = unPrime c'
+    c = d
+
+
 testSuite :: TestTree
 testSuite = testGroup "Diophantine"
   [ testSmallAndQuick "Cornacchia correct" cornacchiaTest
   , testSmallAndQuick "Cornacchia same solutions as brute force" cornacchiaBruteForce
-  ]
+  , testSmallAndQuick "Linear correct" linearTest
+  , testSmallAndQuick "Linear correct #2" linearTest'
+  ]
+