From b0ac56117a1c6a11c4b9a4f6ca921d37a9856e72 Mon Sep 17 00:00:00 2001
From: Xavier Leroy <xavier.leroy@college-de-france.fr>
Date: Thu, 16 Jan 2025 17:51:09 +0100
Subject: [PATCH] Add `Z.solve_linear_congruences`

Solves linear Diophantine equations of one variable.
Generalization of the Chinese remainder theorem.
---
 tests/zq.ml       | 14 ++++++++++++++
 tests/zq.output32 |  4 ++++
 tests/zq.output64 |  4 ++++
 z.ml              | 45 +++++++++++++++++++++++++++++++++++++++++++++
 z.mli             | 14 ++++++++++++++
 5 files changed, 81 insertions(+)

diff --git a/tests/zq.ml b/tests/zq.ml
index fb007ca..31005bd 100644
--- a/tests/zq.ml
+++ b/tests/zq.ml
@@ -626,6 +626,20 @@ let test_Z() =
   Printf.printf "lcm -2^120 -2^300 = %a\n" pr (I.lcm (I.neg p120) (I.neg p300));
   Printf.printf "lcm 2^120 0 = %a\n" pr (I.lcm p120 I.zero);
   Printf.printf "lcm -2^120 0 = %a\n" pr (I.lcm (I.neg p120) I.zero);
+  begin
+    let print_solve l =
+      match I.solve_linear_congruences l with
+      | (c, p) -> Printf.printf "    x = %a + %a k\n" pr c pr p
+      | exception I.No_solution -> Printf.printf "    no solution\n" in
+    Printf.printf "solve_linear_congruences 2x = 5 mod 7 & 3x = 1 mod 11\n";
+    print_solve
+        [(I.of_int 2, I.of_int 5, I.of_int 7);
+         (I.of_int 3, I.of_int 1, I.of_int 11)];
+    Printf.printf "solve_linear_congruences x = 1 mod 6 & x = 5 mod 9\n";
+    print_solve
+        [(I.of_int 1, I.of_int 1, I.of_int 6);
+         (I.of_int 1, I.of_int 5, I.of_int 9)]
+  end;
   Printf.printf "is_odd 0\n = %b\n" (I.is_odd (Z.of_int 0));
   Printf.printf "is_odd 1\n = %b\n" (I.is_odd (Z.of_int 1));
   Printf.printf "is_odd 2\n = %b\n" (I.is_odd (Z.of_int 2));
diff --git a/tests/zq.output32 b/tests/zq.output32
index d0912dd..e63e889 100644
--- a/tests/zq.output32
+++ b/tests/zq.output32
@@ -841,6 +841,10 @@ lcm -2^120 2^300 = 2037035976334486086268445688409378161051468393665936250636140
 lcm -2^120 -2^300 = 2037035976334486086268445688409378161051468393665936250636140449354381299763336706183397376
 lcm 2^120 0 = 0
 lcm -2^120 0 = 0
+solve_linear_congruences 2x = 5 mod 7 & 3x = 1 mod 11
+    x = 48 + 77 k
+solve_linear_congruences x = 1 mod 6 & x = 5 mod 9
+    no solution
 is_odd 0
  = false
 is_odd 1
diff --git a/tests/zq.output64 b/tests/zq.output64
index a49ec06..416787c 100644
--- a/tests/zq.output64
+++ b/tests/zq.output64
@@ -841,6 +841,10 @@ lcm -2^120 2^300 = 2037035976334486086268445688409378161051468393665936250636140
 lcm -2^120 -2^300 = 2037035976334486086268445688409378161051468393665936250636140449354381299763336706183397376
 lcm 2^120 0 = 0
 lcm -2^120 0 = 0
+solve_linear_congruences 2x = 5 mod 7 & 3x = 1 mod 11
+    x = 48 + 77 k
+solve_linear_congruences x = 1 mod 6 & x = 5 mod 9
+    no solution
 is_odd 0
  = false
 is_odd 1
diff --git a/z.ml b/z.ml
index 58faadb..8e3a450 100644
--- a/z.ml
+++ b/z.ml
@@ -447,6 +447,51 @@ let sprint () x = to_string x
 let bprint b x = Buffer.add_string b (to_string x)
 let pp_print f x = Format.pp_print_string f (to_string x)
 
+(* Diophantine equations *)
+
+exception No_solution
+
+(* Chinese remainder theorem.
+   Solve the system [ x = a mod m /\ x = b mod n ]
+   Assume 0 <= a < m, 0 <= b < n.
+   Result is (c, p) with 0 <= c < p.
+   The solutions are x = c + pk for k in Z. *)
+
+let crt2 (a, m) (b, n) =
+  let (d, u, _) = gcdext m n in
+  let (c, r) = (ediv_rem (sub b a) d) in
+  if r <> zero then raise No_solution;
+  let p = mul (div m d) n in
+  let x = add a (mul (mul c u) m) in
+  (erem x p, p)
+
+(* Solve the equation [ a * x = b mod m ].
+   Assume m <> 0.
+   Result is (c, p) with 0 <= c < p.
+   The solutions are x = c + pk for k in Z *)
+
+let solve_linear_congruence (a, b, m) =
+  let m = abs m in
+  let (d, inv_a', _) = gcdext a m in
+  let (b', r) = ediv_rem b d in
+  if r <> zero then raise No_solution;
+  let m' = div m d in
+  let c = mul inv_a' b' in
+  (erem c m', m')
+
+(* Solve the set of equations [ a_i * x = b_i mod m_i ].
+   The m_i must not be 0.
+   Result is (c, p) with 0 <= c < p.
+   The solutions are x = c + pk for k in Z *)
+
+let solve_linear_congruences = function
+  | [] -> (of_int 0, of_int 1)
+  | first :: rem ->
+      List.fold_left
+        (fun cur next -> crt2 cur (solve_linear_congruence next))
+        (solve_linear_congruence first)
+        rem
+
 (* Pseudo-random generation *)
 
 let rec raw_bits_random ?(rng: Random.State.t option) nbits =
diff --git a/z.mli b/z.mli
index 1794942..51f6c53 100644
--- a/z.mli
+++ b/z.mli
@@ -553,6 +553,20 @@ external invert: t -> t -> t = "ml_z_invert"
     Raises a [Division_by_zero] if [base] is not invertible modulo [mod].
  *)
 
+val solve_linear_congruences: (t * t * t) list -> t * t
+(** Solve a set of linear congruence (Diophantine) equations of one unknown [x].
+    [solve_linear_congruences [(a₁, b₁, m₁); ...; (aₙ, bₙ, mₙ)] solves
+    the equations [a₁·x = b₁ mod m₁ /\ ... /\ aₙ·x = bₙ mod mₙ].
+    The moduli [mᵢ] must be nonzero.
+    The Chinese remainder theorem is a special case, taking [aᵢ = 1].
+    @return a pair [(c, p)] with [0 <= c < p]. The solutions are [x = c + p·k] for [k ∈ Z].
+    @raise No_solution if the system has no solution.
+    @since 1.15
+*)
+
+exception No_solution
+(** Exception raised by [solve_linear_congruences] when no solution exists. *)
+
 external probab_prime: t -> int -> int = "ml_z_probab_prime"
 (** [probab_prime x r] returns 0 if [x] is definitely composite,
     1 if [x] is probably prime, and 2 if [x] is definitely prime.