From 2946ba04d5c7d314c2d849742a2b9d71f73ee444 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Tue, 18 Mar 2025 17:52:04 -0700
Subject: [PATCH] fix: assert that nonlinear `Nat` terms are nonneg in cutsat
 (#7561)

This PR fixes the support for nonlinear `Nat` terms in cutsat. For
example, cutsat was failing in the following example
```lean
example (i j k l : Nat) : i / j + k + l - k = i / j + l := by grind
```
because we were not adding the fact that `i / j` is non negative when we
inject the `Nat` expression into `Int`.
---
 src/Init/Data/Int/OfNat.lean                  |  3 ++
 .../Meta/Tactic/Grind/Arith/Cutsat/Nat.lean   | 47 +++++++++++++++++++
 .../Meta/Tactic/Grind/Arith/Cutsat/Var.lean   |  4 +-
 src/Lean/Meta/Tactic/Grind/Simp.lean          | 10 ++--
 tests/lean/run/grind_canon_insts.lean         |  8 ++++
 tests/lean/run/grind_cutsat_nat_dvd.lean      |  9 ++++
 6 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/src/Init/Data/Int/OfNat.lean b/src/Init/Data/Int/OfNat.lean
index 91275e7e8d..9c37b81e6a 100644
--- a/src/Init/Data/Int/OfNat.lean
+++ b/src/Init/Data/Int/OfNat.lean
@@ -84,4 +84,7 @@ theorem ofNat_toNat (a : Int) : (NatCast.natCast a.toNat : Int) = if a ≤ 0 the
     have := Int.toNat_of_nonneg (Int.le_of_lt h)
     assumption
 
+theorem Expr.denoteAsInt_nonneg (ctx : Context) (e : Expr) : e.denoteAsInt ctx ≥ 0 := by
+  simp [Expr.denoteAsInt_eq]
+
 end Int.OfNat
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
index 79ea4cfd39..ed004c0ac7 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Init.Data.Int.OfNat
+import Lean.Meta.Tactic.Grind.Simp
 import Lean.Meta.Tactic.Simp.Arith.Nat.Basic
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.Util
 
@@ -101,4 +102,50 @@ where
   conv (lhs rhs : Lean.Expr) : M (Expr × Expr) :=
     return (← toOfNatExpr lhs, ← toOfNatExpr rhs)
 
+/--
+Given `e` of type `Int`, tries to compute `a : Int.OfNat.Expr` s.t.
+`a.denoteAsInt ctx` is `e`
+-/
+partial def ofDenoteAsIntExpr? (e : Lean.Expr) : OptionT M Expr := do
+  match_expr e with
+  | OfNat.ofNat _ _ _ =>
+    let some n ← getIntValue? e | failure
+    guard (n ≥ 0)
+    return .num n.toNat
+  | HAdd.hAdd _ _ _ i a b =>
+    guard (← isInstHAddInt i)
+    return .add (← ofDenoteAsIntExpr? a) (← ofDenoteAsIntExpr? b)
+  | HMul.hMul _ _ _ i a b =>
+    guard (← isInstHMulInt i)
+    return .mul (← ofDenoteAsIntExpr? a) (← ofDenoteAsIntExpr? b)
+  | HDiv.hDiv _ _ _ i a b =>
+    guard (← isInstHDivInt i)
+    return .div (← ofDenoteAsIntExpr? a) (← ofDenoteAsIntExpr? b)
+  | HMod.hMod _ _ _ i a b =>
+    guard (← isInstHModInt i)
+    return .mod (← ofDenoteAsIntExpr? a) (← ofDenoteAsIntExpr? b)
+  | _ =>
+    let_expr NatCast.natCast _ inst a ← e | failure
+    let_expr instNatCastInt := inst | failure
+    if let some x := (← get).map.get? { expr := a } then
+      return .var x
+    else
+      let x := (← get).ctx.size
+      modify fun s => { s with ctx := s.ctx.push a, map := s.map.insert { expr := a } x }
+      return .var x
+
 end Int.OfNat
+
+namespace Lean.Meta.Grind.Arith.Cutsat
+/--
+If `e` is of the form `a.denoteAsInt ctx` for some `a` and `ctx`,
+assert that `e` is nonnegative.
+-/
+def assertDenoteAsIntNonneg (e : Expr) : GoalM Unit := do
+  if e.isAppOf ``NatCast.natCast then return ()
+  let (some a, s) ← Int.OfNat.ofDenoteAsIntExpr? e |>.run |>.run {} | return ()
+  let ctx := Simp.Arith.Nat.toContextExpr s.ctx
+  let h := mkApp2 (mkConst ``Int.OfNat.Expr.denoteAsInt_nonneg) ctx (toExpr a)
+  pushNewFact' (mkIntLE (mkIntLit 0) e) h (← getGeneration e)
+
+end Lean.Meta.Grind.Arith.Cutsat
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
index 1c33667731..b5916bcddb 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
@@ -5,8 +5,9 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.Meta.IntInstTesters
-import Lean.Meta.Tactic.Grind.Arith.Cutsat.Util
 import Lean.Meta.Tactic.Grind.Simp
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.Util
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.Nat
 
 namespace Lean.Meta.Grind.Arith.Cutsat
 
@@ -29,6 +30,7 @@ private def assertNatCast (e : Expr) : GoalM Unit := do
 
 private def assertHelpers (e : Expr) : GoalM Unit := do
   assertNatCast e
+  assertDenoteAsIntNonneg e
 
 /-- Creates a new variable in the cutsat module. -/
 def mkVar (expr : Expr) : GoalM Var := do
diff --git a/src/Lean/Meta/Tactic/Grind/Simp.lean b/src/Lean/Meta/Tactic/Grind/Simp.lean
index 1ba7c2fa0d..e52e85526a 100644
--- a/src/Lean/Meta/Tactic/Grind/Simp.lean
+++ b/src/Lean/Meta/Tactic/Grind/Simp.lean
@@ -46,9 +46,7 @@ def preprocess (e : Expr) : GoalM Simp.Result := do
   trace_goal[grind.simp] "{e}\n===>\n{e'}"
   return { r with expr := e' }
 
-/-- Infers the type of the proof, preprocess it, and adds it to todo list. -/
-def pushNewFact (proof : Expr) (generation : Nat := 0) : GoalM Unit := do
-  let prop ← inferType proof
+def pushNewFact' (prop : Expr) (proof : Expr) (generation : Nat := 0) : GoalM Unit := do
   let r ← preprocess prop
   let prop' := r.expr
   let proof := if let some h := r.proof? then
@@ -57,4 +55,10 @@ def pushNewFact (proof : Expr) (generation : Nat := 0) : GoalM Unit := do
     proof
   modify fun s => { s with newFacts := s.newFacts.push <| .fact prop' proof generation }
 
+
+/-- Infers the type of the proof, preprocess it, and adds it to todo list. -/
+def pushNewFact (proof : Expr) (generation : Nat := 0) : GoalM Unit := do
+  let prop ← inferType proof
+  pushNewFact' prop proof generation
+
 end Lean.Meta.Grind
diff --git a/tests/lean/run/grind_canon_insts.lean b/tests/lean/run/grind_canon_insts.lean
index 5603df9839..8f31c8e0d6 100644
--- a/tests/lean/run/grind_canon_insts.lean
+++ b/tests/lean/run/grind_canon_insts.lean
@@ -81,6 +81,14 @@ info: [Meta.debug] [@HMul.hMul Int Int Int (@instHMul Int Int.instMul) (@NatCast
        (@NatCast.natCast Int instNatCastInt a),
      @HMul.hMul Int Int Int (@instHMul Int Int.instMul) (@NatCast.natCast Int instNatCastInt b)
        (@NatCast.natCast Int instNatCastInt d),
+     @HMul.hMul Int Int Int (@instHMul Int Int.instMul)
+       (@Neg.neg Int Int.instNegInt (@OfNat.ofNat Int 1 (@instOfNat 1)))
+       (@HMul.hMul Int Int Int (@instHMul Int Int.instMul) (@NatCast.natCast Int instNatCastInt b)
+         (@NatCast.natCast Int instNatCastInt a)),
+     @HMul.hMul Int Int Int (@instHMul Int Int.instMul)
+       (@Neg.neg Int Int.instNegInt (@OfNat.ofNat Int 1 (@instOfNat 1)))
+       (@HMul.hMul Int Int Int (@instHMul Int Int.instMul) (@NatCast.natCast Int instNatCastInt b)
+         (@NatCast.natCast Int instNatCastInt d)),
      @HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b a,
      @HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b d]
 -/
diff --git a/tests/lean/run/grind_cutsat_nat_dvd.lean b/tests/lean/run/grind_cutsat_nat_dvd.lean
index f9741ed9ea..b6f47197f8 100644
--- a/tests/lean/run/grind_cutsat_nat_dvd.lean
+++ b/tests/lean/run/grind_cutsat_nat_dvd.lean
@@ -42,3 +42,12 @@ example (x y : Nat) (_ : 2 ≤ x) (_ : x ≤ 3) (_ : 2 ≤ y) (_ : y ≤ 3) :
 example (x y : Nat) (_ : 2 ≤ x) (_ : x ≤ 3) (_ : 2 ≤ y) (_ : y ≤ 3) :
     4 ≤ (y + x) % 8 ∧ (y + x) % 8 ≤ 6 := by
   grind
+
+example (i j k l : Nat) : i / j + k + l - k = i / j + l := by
+  grind
+
+example (i j k l : Nat) : i % j + k + l - k = i % j + l := by
+  grind
+
+example (i j k l : Nat) : i * j + k + l - k = i * j + l := by
+  grind