feat: use nat_gcd in the kernel (#2533)

* feat: use nat_gcd in the kernel

---------

Co-authored-by: Sebastian Ullrich <sebasti@nullri.ch>

src/Lean/Meta/WHNF.lean

@@ -823,6 +823,8 @@ def reduceNat? (e : Expr) : MetaM (Option Expr) :=
       else if fn == ``Nat.mul then reduceBinNatOp Nat.mul a1 a2
       else if fn == ``Nat.div then reduceBinNatOp Nat.div a1 a2
       else if fn == ``Nat.mod then reduceBinNatOp Nat.mod a1 a2
+      else if fn == ``Nat.pow then reduceBinNatOp Nat.pow a1 a2
+      else if fn == ``Nat.gcd then reduceBinNatOp Nat.gcd a1 a2
       else if fn == ``Nat.beq then reduceBinNatPred Nat.beq a1 a2
       else if fn == ``Nat.ble then reduceBinNatPred Nat.ble a1 a2
       else return none

src/kernel/type_checker.cpp

@@ -30,6 +30,7 @@ static expr * g_nat_add      = nullptr;
 static expr * g_nat_sub      = nullptr;
 static expr * g_nat_mul      = nullptr;
 static expr * g_nat_pow      = nullptr;
+static expr * g_nat_gcd      = nullptr;
 static expr * g_nat_mod      = nullptr;
 static expr * g_nat_div      = nullptr;
 static expr * g_nat_beq      = nullptr;
@@ -603,6 +604,7 @@ optional<expr> type_checker::reduce_nat(expr const & e) {
         if (f == *g_nat_sub) return reduce_bin_nat_op(nat_sub, e);
         if (f == *g_nat_mul) return reduce_bin_nat_op(nat_mul, e);
         if (f == *g_nat_pow) return reduce_bin_nat_op(nat_pow, e);
+        if (f == *g_nat_gcd) return reduce_bin_nat_op(nat_gcd, e);
         if (f == *g_nat_mod) return reduce_bin_nat_op(nat_mod, e);
         if (f == *g_nat_div) return reduce_bin_nat_op(nat_div, e);
         if (f == *g_nat_beq) return reduce_bin_nat_pred(nat_eq, e);
@@ -1151,6 +1153,8 @@ void initialize_type_checker() {
     mark_persistent(g_nat_mul->raw());
     g_nat_pow      = new expr(mk_constant(name{"Nat", "pow"}));
     mark_persistent(g_nat_pow->raw());
+    g_nat_gcd      = new expr(mk_constant(name{"Nat", "gcd"}));
+    mark_persistent(g_nat_gcd->raw());
     g_nat_div      = new expr(mk_constant(name{"Nat", "div"}));
     mark_persistent(g_nat_div->raw());
     g_nat_mod      = new expr(mk_constant(name{"Nat", "mod"}));
@@ -1177,6 +1181,7 @@ void finalize_type_checker() {
     delete g_nat_sub;
     delete g_nat_mul;
     delete g_nat_pow;
+    delete g_nat_gcd;
     delete g_nat_div;
     delete g_nat_mod;
     delete g_nat_beq;

src/runtime/object.h

@@ -325,6 +325,7 @@ inline obj_res nat_add(b_obj_arg a1, b_obj_arg a2) { return lean_nat_add(a1, a2)
 inline obj_res nat_sub(b_obj_arg a1, b_obj_arg a2) { return lean_nat_sub(a1, a2); }
 inline obj_res nat_mul(b_obj_arg a1, b_obj_arg a2) { return lean_nat_mul(a1, a2); }
 inline obj_res nat_pow(b_obj_arg a1, b_obj_arg a2) { return lean_nat_pow(a1, a2); }
+inline obj_res nat_gcd(b_obj_arg a1, b_obj_arg a2) { return lean_nat_gcd(a1, a2); }
 inline obj_res nat_div(b_obj_arg a1, b_obj_arg a2) { return lean_nat_div(a1, a2); }
 inline obj_res nat_mod(b_obj_arg a1, b_obj_arg a2) { return lean_nat_mod(a1, a2); }
 inline bool nat_eq(b_obj_arg a1, b_obj_arg a2) { return lean_nat_eq(a1, a2); }

tests/lean/run/lean_nat_gcd.lean

@@ -0,0 +1,67 @@
+import Lean
+
+/-! Basic tests, including edge cases. -/
+example : Nat.gcd 0  0  = 0  := rfl
+example : Nat.gcd 0  1  = 1  := rfl
+example : Nat.gcd 0  17 = 17 := rfl
+example : Nat.gcd 1  0  = 1  := rfl
+example : Nat.gcd 17 0  = 17 := rfl
+example : Nat.gcd 1  1  = 1  := rfl
+example : Nat.gcd 1  17 = 1  := rfl
+example : Nat.gcd 17 1  = 1  := rfl
+example : Nat.gcd 2  2  = 2  := rfl
+example : Nat.gcd 2  3  = 1  := rfl
+example : Nat.gcd 2  4  = 2  := rfl
+example : Nat.gcd 9  6  = 3  := rfl
+
+/-!
+We check that `Nat.gcd` is evaluated using bignum functions in the kernel.
+
+Because of variations in run time on different operating systems during CI,
+for the larger calculations we don't attempt to do any timing.
+
+All of the "large" calculations below used to fail with
+```
+maximum recursion depth has been reached (use `set_option maxRecDepth <num>` to increase limit)
+```
+prior to lean4#2533.
+-/
+
+open Lean Elab Command in
+elab "#fast" c:command : command => do
+  let start ← IO.monoMsNow
+  elabCommand c
+  let elapsed := (← IO.monoMsNow) - start
+  if elapsed > 1000 then throwError m!"Too slow! {elapsed}ms"
+
+-- Used to take ~1500ms, now takes ~3ms.
+#fast example : Nat.gcd 45 (Nat.gcd 15 85) = 5 := rfl
+
+example : Nat.gcd (115249 * 180811) (115249*181081) = 115249 := rfl
+
+/-- Mersenne primes. -/
+def m (p : Nat) := 2^p - 1
+
+/-- Largish primes, targeting <100ms GCD calculations. -/
+def p_29 := 110503
+def p_30 := 132049
+def p_31 := 216091
+def p_32 := 756839
+def p_33 := 859433
+
+/- GCD with large prime factors on one side, and small primes on the other. -/
+example : Nat.gcd (p_29 * p_30 * p_31 * p_32 * p_33) 2^(2^20) = 1 := rfl
+/- GCD with two prime factors on both sides, including one in common. -/
+example : Nat.gcd (m p_31 * m p_33) (m p_32 * m p_33) - m p_33 = 0 := rfl
+/- GCD with many small prime factors. -/
+example :
+    Nat.gcd (2^1 * 3^1 * 5^2 * 7^3 * 11^5 * 13^8) (2^8 * 3^5 * 5^3 * 7^2 * 11^1 * 13^1) =
+      2 * 3 * 5^2 * 7^2 * 11 * 13 := rfl
+
+-- #eval Lean.maxSmallNat -- 9223372036854775807
+def maxSmallNat := 9223372036854775807
+
+example : maxSmallNat = 7^2 * 73 * 127 * 337 * 92737 * 649657 := rfl
+-- Calculate GCDs of numbers on either side of `maxSmallNat`.
+example : Nat.gcd (maxSmallNat - 92737) (maxSmallNat + 92737) = 185474 := rfl
+example : Nat.gcd (maxSmallNat / 649657) (maxSmallNat * 649657) = 14197294936951 := rfl