feat: add a public extended Euclidean algorithm for natural numbers

src/Init/Data/Nat.lean

@@ -10,6 +10,7 @@ public import Init.Data.Nat.Basic
 public import Init.Data.Nat.Div
 public import Init.Data.Nat.Dvd
 public import Init.Data.Nat.Gcd
+public import Init.Data.Nat.ExtendedGcd
 public import Init.Data.Nat.Coprime
 public import Init.Data.Nat.MinMax
 public import Init.Data.Nat.Order

src/Init/Data/Nat/ExtendedGcd.lean

@@ -0,0 +1,145 @@
+/-
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+module
+
+prelude
+public import Init.Data.Nat.Gcd
+public import Init.Data.Int.Basic
+public import Init.Data.Int.Repr
+import Init.Data.Int.Lemmas
+import Init.Data.AC
+import Init.TacticsExtra
+
+/-!
+# Extended Euclidean algorithm
+
+`Nat.extendedGcd` computes a greatest common divisor together with signed Bézout coefficients in
+one Euclidean pass. `Nat.extendedGcd_gcd` identifies the gcd, and `Nat.extendedGcd_bezout` certifies
+the coefficients. The Lean implementation uses arbitrary-precision `Nat` and `Int` arithmetic;
+compiled evaluation uses a native implementation backed by GMP when available.
+-/
+
+@[expose] public section
+
+namespace Nat
+
+/-- The gcd and the two Bézout coefficients returned by `Nat.extendedGcd`. -/
+structure ExtendedGcdResult where
+  /-- The greatest common divisor of the inputs. -/
+  gcd : Nat
+  /-- The signed coefficient of the first input. -/
+  coeffA : Int
+  /-- The signed coefficient of the second input. -/
+  coeffB : Int
+deriving Repr, DecidableEq, Inhabited
+
+/--
+Computes the greatest common divisor of `a` and `b` together with signed Bézout coefficients in
+one pass of the extended Euclidean algorithm. The result `r` satisfies `r.gcd = Nat.gcd a b` and
+`(r.gcd : Int) = a * r.coeffA + b * r.coeffB`.
+
+For a zero first input, the result is `⟨b, 0, 1⟩`, including `⟨0, 0, 1⟩` when both inputs are zero.
+For a nonzero first input and a zero second input, the result is `⟨a, 1, 0⟩`.
+The coefficients are not unique; no minimality or symmetry of the returned coefficients is
+specified.
+
+Examples:
+* `Nat.extendedGcd 240 46 = ⟨2, -9, 47⟩`
+* `Nat.extendedGcd 0 19 = ⟨19, 0, 1⟩`
+* `Nat.extendedGcd 19 0 = ⟨19, 1, 0⟩`
+-/
+@[extern "lean_nat_extended_gcd"]
+def extendedGcd (a b : @& Nat) : ExtendedGcdResult :=
+  go a 1 0 b 0 1
+where
+  /-- Implementation detail of `Nat.extendedGcd`: the two rows represent the remainders as
+  linear combinations of the original inputs. -/
+  go (r : Nat) (s t : Int) (r' : Nat) (s' t' : Int) : ExtendedGcdResult :=
+    if _h : r = 0 then
+      ⟨r', s', t'⟩
+    else
+      let q : Int := r' / r
+      go (r' % r) (s' - q * s) (t' - q * t) r s t
+  termination_by r
+  decreasing_by exact Nat.mod_lt _ (Nat.pos_of_ne_zero _h)
+
+@[export lean_nat_extended_gcd_fallback]
+private def extendedGcdFallback (a b : Nat) : ExtendedGcdResult :=
+  extendedGcd.go a 1 0 b 0 1
+
+@[simp] theorem extendedGcd_zero_left (b : Nat) : extendedGcd 0 b = ⟨b, 0, 1⟩ := by
+  rw [extendedGcd, extendedGcd.go]
+  rfl
+
+@[simp] theorem extendedGcd_zero_right {a : Nat} (h : a ≠ 0) :
+    extendedGcd a 0 = ⟨a, 1, 0⟩ := by
+  rw [extendedGcd, extendedGcd.go, dite_eq_right h]
+  simp only [Nat.zero_div, Nat.zero_mod, Int.natCast_zero, Int.zero_mul, Int.sub_zero]
+  rw [extendedGcd.go]
+  rfl
+
+@[simp] theorem extendedGcd_self {a : Nat} (h : a ≠ 0) : extendedGcd a a = ⟨a, 1, 0⟩ := by
+  rw [extendedGcd, extendedGcd.go, dite_eq_right h]
+  simp only [Nat.mod_self]
+  rw [extendedGcd.go]
+  rfl
+
+private theorem extendedGcd_go_gcd (r r' : Nat) : ∀ s t s' t',
+    (extendedGcd.go r s t r' s' t').gcd = Nat.gcd r r' := by
+  induction r, r' using Nat.gcd.induction with
+  | H0 r' =>
+    intro s t s' t'
+    rw [extendedGcd.go]
+    simp
+  | H1 r r' hr ih =>
+    intro s t s' t'
+    rw [extendedGcd.go, dite_eq_right (Nat.ne_of_gt hr), ih, ← Nat.gcd_rec]
+
+private theorem extendedGcd_go_bezout (a b r r' : Nat) : ∀ s t s' t',
+    (r : Int) = a * s + b * t → (r' : Int) = a * s' + b * t' →
+    let out := extendedGcd.go r s t r' s' t'
+    (out.gcd : Int) = a * out.coeffA + b * out.coeffB := by
+  induction r, r' using Nat.gcd.induction with
+  | H0 r' =>
+    intro s t s' t' _ hr'
+    simpa [extendedGcd.go] using hr'
+  | H1 r r' h ih =>
+    intro s t s' t' hr hr'
+    rw [extendedGcd.go, dite_eq_right (Nat.ne_of_gt h)]
+    apply ih _ _ _ _ ?_ hr
+    calc
+      (r' % r : Int) = (r' : Int) - (r' / r : Nat) * (r : Int) := by
+        have hd := congrArg (fun n : Nat => (n : Int)) (Nat.mod_add_div r' r)
+        simp only [Int.natCast_add, Int.natCast_mul] at hd
+        rw [← hd, Int.mul_comm (r : Int), Int.add_sub_cancel]
+      _ = a * (s' - (r' / r : Nat) * s) + b * (t' - (r' / r : Nat) * t) := by
+        rw [Int.mul_sub, Int.mul_sub, hr', hr, Int.mul_add]
+        simp only [Int.sub_eq_add_neg, Int.neg_add]
+        ac_rfl
+
+/-- The gcd component of `Nat.extendedGcd a b` is `Nat.gcd a b`. -/
+@[simp] theorem extendedGcd_gcd (a b : Nat) : (extendedGcd a b).gcd = Nat.gcd a b :=
+  extendedGcd_go_gcd a b 1 0 0 1
+
+/-- The coefficients returned by `Nat.extendedGcd` satisfy Bézout's identity. -/
+theorem extendedGcd_bezout (a b : Nat) :
+    (Nat.gcd a b : Int) = a * (extendedGcd a b).coeffA + b * (extendedGcd a b).coeffB := by
+  have h := extendedGcd_go_bezout a b a b 1 0 0 1 (by simp) (by simp)
+  change ((extendedGcd a b).gcd : Int) =
+    a * (extendedGcd a b).coeffA + b * (extendedGcd a b).coeffB at h
+  simpa only [extendedGcd_gcd] using h
+
+/-- The gcd returned by `Nat.extendedGcd a b` divides `a`. -/
+theorem extendedGcd_dvd_left (a b : Nat) : (extendedGcd a b).gcd ∣ a := by
+  rw [extendedGcd_gcd]
+  exact Nat.gcd_dvd_left a b
+
+/-- The gcd returned by `Nat.extendedGcd a b` divides `b`. -/
+theorem extendedGcd_dvd_right (a b : Nat) : (extendedGcd a b).gcd ∣ b := by
+  rw [extendedGcd_gcd]
+  exact Nat.gcd_dvd_right a b
+
+end Nat

src/include/lean/lean.h

@@ -1781,6 +1781,7 @@ LEAN_EXPORT lean_obj_res lean_nat_shiftl(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_big_shiftr(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_pow(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_gcd(b_lean_obj_arg a1, b_lean_obj_arg a2);
+LEAN_EXPORT lean_obj_res lean_nat_extended_gcd(b_lean_obj_arg a1, b_lean_obj_arg a2);
 LEAN_EXPORT lean_obj_res lean_nat_log2(b_lean_obj_arg a);
 /* Upper bound on the size in bytes of the representation of `a` (one word for scalars). Returns a raw `size_t`, not a boxed `Nat`. */
 LEAN_EXPORT size_t lean_nat_size_in_bytes(b_lean_obj_arg a);

src/runtime/mpz.h

@@ -67,6 +67,8 @@ class LEAN_EXPORT mpz {
 
 #ifdef LEAN_USE_GMP
     void set(mpz_t r) const;
+    mpz_srcptr get_mpz_t() const { return m_val; }
+    mpz_ptr get_mpz_t() { return m_val; }
 #endif
 
     friend void swap(mpz & a, mpz & b) noexcept;

src/runtime/object.cpp

@@ -1379,6 +1379,8 @@ static inline obj_res mpz_to_nat(mpz const & m) {
         return mpz_to_nat_core(m);
 }
 
+static object * mpz_to_int(mpz const & m);
+
 extern "C" LEAN_EXPORT object * lean_cstr_to_nat(char const * n) {
     return mpz_to_nat(mpz(n));
 }
@@ -1621,6 +1623,32 @@ extern "C" LEAN_EXPORT lean_obj_res lean_nat_gcd(b_lean_obj_arg a1, b_lean_obj_a
     }
 }
 
+extern "C" LEAN_EXPORT lean_obj_res lean_nat_extended_gcd_fallback(lean_obj_arg a, lean_obj_arg b);
+
+extern "C" LEAN_EXPORT lean_obj_res lean_nat_extended_gcd(b_lean_obj_arg a, b_lean_obj_arg b) {
+#ifdef LEAN_USE_GMP
+    if (a != lean_box(0) && b != lean_box(0) && !lean_nat_dec_eq(a, b) &&
+        !(lean_is_scalar(a) && lean_is_scalar(b) && lean_unbox(a) <= UINT16_MAX && lean_unbox(b) <= UINT16_MAX)) {
+        mpz aa = lean_is_scalar(a) ? mpz::of_size_t(lean_unbox(a)) : mpz_value(a);
+        mpz bb = lean_is_scalar(b) ? mpz::of_size_t(lean_unbox(b)) : mpz_value(b);
+        mpz g, s, t;
+        // Outside the zero/equal cases, the final Euclidean quotient is at least two.
+        // The alternating-sign coefficient recurrence then gives GMP's half-size bounds;
+        // equality occurs only when a/g or b/g is two, with the corresponding coefficient +1.
+        mpz_gcdext(g.get_mpz_t(), s.get_mpz_t(), t.get_mpz_t(), aa.get_mpz_t(), bb.get_mpz_t());
+        // Match the field order of Nat.ExtendedGcdResult.
+        obj_res r = lean_alloc_ctor(0, 3, 0);
+        lean_ctor_set(r, 0, mpz_to_nat(g));
+        lean_ctor_set(r, 1, mpz_to_int(s));
+        lean_ctor_set(r, 2, mpz_to_int(t));
+        return r;
+    }
+#endif
+    lean_inc(a);
+    lean_inc(b);
+    return lean_nat_extended_gcd_fallback(a, b);
+}
+
 extern "C" LEAN_EXPORT lean_obj_res lean_nat_log2(b_lean_obj_arg a) {
     if (lean_is_scalar(a)) {
       unsigned res = 0;

tests/compile/nat_extended_gcd.lean

@@ -0,0 +1,29 @@
+module
+
+import Init.Data.Nat.ExtendedGcd
+
+/-!
+Tests `Nat.extendedGcd` from compiled and interpreted callers, with retained operands and results
+whose gcd and both signed coefficients require heap allocation.
+-/
+
+private def values (seed : Nat) : Array Nat :=
+  #[0, 1, 2, 3, 65535, 65536, 65537,
+    2 ^ 31 - 1, 2 ^ 31, 2 ^ 31 + 1, 2 ^ 31 + 3,
+    2 ^ 32 - 1, 2 ^ 32, 2 ^ 32 + 1, 2 ^ 32 + 3,
+    2 ^ 63 - 1, 2 ^ 63, 2 ^ 63 + 1, 2 ^ 64 - 1, 2 ^ 64, 2 ^ 64 + 1,
+    2 ^ 128 - 1, 2 ^ 128, 2 ^ 128 + 1, 2 ^ 256 - 1, 2 ^ 256 + 1].map (· + seed)
+
+public def main (args : List String) : IO Unit := do
+  -- Depend on runtime input so the calls are not evaluated during module initialization.
+  let inputs := values args.length
+  for g in #[1, 2 ^ 64] do
+    for x in inputs do
+      for y in inputs do
+        let a := x * g
+        let b := y * g
+        let r := Nat.extendedGcd a b
+        let expected := Nat.extendedGcd.go a 1 0 b 0 1
+        unless r = expected && r.gcd = Nat.gcd a b &&
+            (r.gcd : Int) = a * r.coeffA + b * r.coeffB do
+          throw <| IO.userError s!"extendedGcd disagrees on ({a}, {b}): {repr r}, expected {repr expected}"

tests/compile/nat_extended_gcd_fallback.lean

@@ -0,0 +1,26 @@
+module
+
+import Init.Data.Nat.ExtendedGcd
+
+/-!
+Tests the exported Lean fallback for `Nat.extendedGcd`, including its ownership convention,
+on small and large inputs even when GMP is enabled. Inputs remain live after each call.
+-/
+
+@[extern "lean_nat_extended_gcd_fallback"]
+private def fallback (a b : Nat) : Nat.ExtendedGcdResult :=
+  Nat.extendedGcd.go a 1 0 b 0 1
+
+public def main (args : List String) : IO Unit := do
+  let inputs := #[0, 1, 2, 3, 65535, 65536, 65537,
+    2^31 - 1, 2^31, 2^32 + 1, 2^63 - 1, 2^63, 2^64 + 1,
+    2^128 - 1, 2^128 + 1, 2^256 + 1].map (· + args.length)
+  for g in #[1, 2^64] do
+    for x in inputs do
+      for y in inputs do
+        let a := x * g
+        let b := y * g
+        let r := fallback a b
+        unless r = Nat.extendedGcd a b && r.gcd = Nat.gcd a b &&
+            (r.gcd : Int) = a * r.coeffA + b * r.coeffB do
+          throw <| IO.userError s!"extendedGcd fallback disagrees on ({a}, {b}): {repr r}"

tests/elab/nat_extended_gcd.lean

@@ -0,0 +1,62 @@
+module
+
+import Init.Data.Nat.ExtendedGcd
+
+/-!
+Tests the extended Euclidean algorithm on natural numbers: zero conventions, signed Bézout
+coefficients, arbitrary-precision inputs, long Euclidean chains, and the public correctness lemmas.
+-/
+
+#guard Nat.extendedGcd 0 0 = ⟨0, 0, 1⟩
+#guard Nat.extendedGcd 0 19 = ⟨19, 0, 1⟩
+#guard Nat.extendedGcd 19 0 = ⟨19, 1, 0⟩
+#guard Nat.extendedGcd 19 19 = ⟨19, 1, 0⟩
+#guard Nat.extendedGcd 240 46 = ⟨2, -9, 47⟩
+#guard Nat.extendedGcd 46 240 = ⟨2, 47, -9⟩
+
+-- These examples exercise kernel reduction rather than only compiled evaluation.
+example : Nat.extendedGcd 240 46 = ⟨2, -9, 47⟩ := by decide +kernel
+example : Nat.extendedGcd 46 240 = ⟨2, 47, -9⟩ := by decide +kernel
+example : (Nat.extendedGcd (2 ^ 64) (2 ^ 64 + 1)).gcd = 1 := by decide +kernel
+
+private def values : List Nat :=
+  [0, 1, 2, 3, 6, 15, 46, 97, 240, 65536, 2 ^ 31 - 1, 2 ^ 32, 2 ^ 63 - 1, 2 ^ 63,
+    2 ^ 64 - 1, 2 ^ 64, 2 ^ 64 + 1, 2 ^ 128 - 1, 2 ^ 128 + 1, 2 ^ 256 - 1]
+
+private def correct (a b : Nat) : Bool :=
+  let r := Nat.extendedGcd a b
+  r.gcd == Nat.gcd a b && decide ((r.gcd : Int) = a * r.coeffA + b * r.coeffB) &&
+    decide (r.gcd ∣ a ∧ r.gcd ∣ b)
+
+#guard values.all fun a => values.all fun b => correct a b
+
+private def fibPair (n : Nat) : Nat × Nat :=
+  match n with
+  | 0 => (0, 1)
+  | n + 1 => let (a, b) := fibPair n; (b, a + b)
+
+-- Consecutive Fibonacci numbers force many quotient-one Euclidean steps.
+#guard [100, 256, 512].all fun n =>
+  let (a, b) := fibPair n
+  correct a b && (Nat.extendedGcd a b).gcd == 1
+
+#guard (Nat.extendedGcd (240 * 2 ^ 128) (46 * 2 ^ 128)).gcd == 2 * 2 ^ 128
+
+example (a b : Nat) : (Nat.extendedGcd a b).gcd = Nat.gcd a b := by simp
+
+example (a b : Nat) :
+    (Nat.gcd a b : Int) = a * (Nat.extendedGcd a b).coeffA + b * (Nat.extendedGcd a b).coeffB :=
+  Nat.extendedGcd_bezout a b
+
+example (b : Nat) : Nat.extendedGcd 0 b = ⟨b, 0, 1⟩ := by simp
+example (a : Nat) (h : a ≠ 0) : Nat.extendedGcd a 0 = ⟨a, 1, 0⟩ := by simp [h]
+example (a : Nat) (h : a ≠ 0) : Nat.extendedGcd a a = ⟨a, 1, 0⟩ := by simp [h]
+
+example (a b : Nat) : (Nat.extendedGcd a b).gcd ∣ a := Nat.extendedGcd_dvd_left a b
+example (a b : Nat) : (Nat.extendedGcd a b).gcd ∣ b := Nat.extendedGcd_dvd_right a b
+
+-- A caller can retain one computed result and reason about it without exposing the algorithm.
+example (a b : Nat) (r : Nat.ExtendedGcdResult) (h : r = Nat.extendedGcd a b)
+    (hc : Nat.gcd a b = 1) : a * r.coeffA + b * r.coeffB = (1 : Int) := by
+  subst r
+  rw [← Nat.extendedGcd_bezout, hc, Int.natCast_one]