Use Nat.nthRoot instead of custom helper and add formula tests

Copilot · ckrause · Copilot · commit bbf4743ea33a · 2025-11-08T14:27:48.000Z
Co-authored-by: ckrause &lt;840744+ckrause@users.noreply.github.com&gt;
diff --git a/src/form/lean.cpp b/src/form/lean.cpp
@@ -141,7 +141,7 @@ bool LeanFormula::convertToLean(Expression& expr, int64_t offset,
         break;
       }
       if (expr.name == "sqrtnint") {
-        // sqrtnint(x, n) -> sqrtnint x n (custom helper function)
+        // sqrtnint(x, n) -> Int.ofNat (Nat.nthRoot (Int.toNat n) (Int.toNat x))
         if (expr.children.size() != 2) {
           return false;
         }
@@ -153,9 +153,13 @@ bool LeanFormula::convertToLean(Expression& expr, int64_t offset,
         if (ExpressionUtil::canBeNegative(root, offset)) {
           return false;  // negative root index not supported
         }
-        // Use helper function sqrtnint
-        expr.name = "sqrtnint";
-        helper_funcs.insert("sqrtnint");
+        // Use Nat.nthRoot from Mathlib
+        Expression baseToNat(Expression::Type::FUNCTION, "Int.toNat", {base});
+        Expression rootToNat(Expression::Type::FUNCTION, "Int.toNat", {root});
+        Expression nthRoot(Expression::Type::FUNCTION, "Nat.nthRoot", {rootToNat, baseToNat});
+        Expression ofNat(Expression::Type::FUNCTION, "Int.ofNat", {nthRoot});
+        expr = ofNat;
+        imports.insert("Mathlib.Data.Nat.NthRoot.Defs");
         break;
       }
       // Allow calls to locally defined functions and sequences
@@ -256,36 +260,6 @@ std::string LeanFormula::toString() const {
   auto functions = FormulaUtil::getDefinitions(main_formula);
 
   std::stringstream buf;
-  
-  // Print helper function definitions first
-  if (!helper_funcs.empty()) {
-    for (const auto& func : helper_funcs) {
-      if (func == "sqrtnint") {
-        // Integer nth root: floor(x^(1/n))
-        // We compute it by binary search to find largest y such that y^n <= x
-        buf << "def sqrtnint (x : Int) (n : Int) : Int :=\n";
-        buf << "  if n <= 0 then 0\n";
-        buf << "  else if x < 0 then 0\n";
-        buf << "  else\n";
-        buf << "    let xNat := Int.toNat x\n";
-        buf << "    let nNat := Int.toNat n\n";
-        buf << "    -- Binary search for largest y such that y^n <= x\n";
-        buf << "    let rec search (low high : Nat) (fuel : Nat) : Nat :=\n";
-        buf << "      match fuel with\n";
-        buf << "      | 0 => low\n";
-        buf << "      | fuel' + 1 =>\n";
-        buf << "        if low >= high then low\n";
-        buf << "        else\n";
-        buf << "          let mid := (low + high + 1) / 2\n";
-        buf << "          if mid ^ nNat <= xNat then\n";
-        buf << "            search mid high fuel'\n";
-        buf << "          else\n";
-        buf << "            search low (mid - 1) fuel'\n";
-        buf << "    Int.ofNat (search 0 xNat 64)\n\n";
-      }
-    }
-  }
-  
   if (functions.size() == 1) {
     buf << printFunction(functions[0]);
   } else {
diff --git a/src/form/lean.hpp b/src/form/lean.hpp
@@ -28,7 +28,6 @@ class LeanFormula {
   Formula main_formula;
   std::string domain;  // Int or Nat
   std::set<std::string> imports;
-  std::set<std::string> helper_funcs;
   std::vector<std::string> funcNames;
 
   static bool initializeLeanProject();
diff --git a/tests/formula/formula.txt b/tests/formula/formula.txt
@@ -82,6 +82,7 @@ A007531: a(n) = n*(n-2)*(n-1)
 A007583: a(n) = 2*floor((4^n)/3)+1
 A008785: a(n) = (n+4)^n
 A008999: a(n) = 2*a(n-1)+a(n-4), a(4) = 17, a(3) = 8, a(2) = 4, a(1) = 2, a(0) = 1
+A010052: a(n) = (sqrtint(n)^2)==n
 A010170: a(n) = (gcd(-(n==0)+n,2)+7)*((n==0)+gcd(-(n==0)+n,2)+7)-63
 A010873: a(n) = b(n+3), b(n) = b(n-4), b(4) = 1, b(3) = 0, b(2) = 3, b(1) = 2, b(0) = 1
 A014731: a(n) = 4*b(n)^2, b(n) = 4*b(n-1)+b(n-2), b(1) = -2, b(0) = -1
@@ -122,3 +123,4 @@ A356334: a(n) = 3, a(2) = 4, a(1) = 3, a(0) = 1
 A356464: a(n) = sign(n-1)*((n-2)%2+1)+1
 A368047: a(n) = floor((n*(n+1)*(sign(n)*((n-1)%2+1)+1)*(n*sign(n)*((n-1)%2+1)+2))/12)
 A382019: a(n) = floor((-sign(-n)*((abs(-n)-1)%6+1)+2*n)/3)
+A999998: a(n) = (sqrtnint(n,3)^3)==n
diff --git a/tests/formula/lean.txt b/tests/formula/lean.txt
@@ -65,6 +65,7 @@ A008369: def a (n : Int) : Int := Int.fdiv ((Int.xor (n^2) 3)*((Int.xor (n^2) 3)
 A008531: def a (n : Int) : Int := n*(5*n^2+5)+(Bool.toInt (n==0))
 A008576: def a (n : Int) : Int := (Int.tdiv (8*n-2) 3)+1
 A008577: def a (n : Int) : Int := Int.fdiv ((2*n+1)^2+3) 3
+A010052: def a (n : Int) : Int := Bool.toInt (((Int.ofNat (Nat.sqrt (Int.toNat n)))^2)==n)
 A010170: def a (n : Int) : Int := ((Int.gcd (-(Bool.toInt (n==0))+n) 2)+7)*((Bool.toInt (n==0))+(Int.gcd (-(Bool.toInt (n==0))+n) 2)+7)-63
 A016729: def a (n : Int) : Int := (Int.fdiv ((Int.gcd ((Int.xor (n-1) 97)-1) (n+96))+n+96) 3)-31
 A021019: def a (n : Int) : Int := 6*(min n 1)
@@ -73,3 +74,4 @@ A080457: def a (n : Int) : Int := (Int.lor (3*n-6) 3)+4
 A106249: def a (n : Int) : Int := Int.tdiv ((Int.lor (n+1) 2)-1) 2
 A109008: def a (n : Int) : Int := Int.gcd n 4
 A129760: def a (n : Int) : Int := Int.land n (n-1)
+A999998: def a (n : Int) : Int := Bool.toInt (((Int.ofNat (Nat.nthRoot (Int.toNat 3) (Int.toNat n)))^3)==n)
diff --git a/tests/programs/oeis/999/A999998.asm b/tests/programs/oeis/999/A999998.asm
@@ -0,0 +1,7 @@
+; A999998: Characteristic function of perfect cubes: a(n) = 1 if n is a cube, otherwise 0.
+; Test case for cube root in formulas
+
+mov $1,$0
+nrt $0,3
+pow $0,3
+equ $0,$1
