From b7354aacaa01af1599d463f330772ec9e44bd742 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 15 Mar 2025 14:24:04 -0700
Subject: [PATCH] feat: `Nat` equalities and disequalities in cutsat (#7501)

This PR implements support for `Nat` equalities and disequalities in the
cutsat procedure.
---
 src/Init/Data/Int/Linear.lean                 | 10 +++
 src/Init/Data/Int/OfNat.lean                  |  8 ++
 src/Init/Grind/Norm.lean                      |  3 +
 src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean  |  1 +
 .../Tactic/Grind/Arith/Cutsat/DivMod.lean     | 14 +--
 .../Tactic/Grind/Arith/Cutsat/EqCnstr.lean    | 87 ++++++++++++++-----
 .../Meta/Tactic/Grind/Arith/Cutsat/Nat.lean   |  7 ++
 .../Meta/Tactic/Grind/Arith/Cutsat/Proof.lean | 12 ++-
 .../Meta/Tactic/Grind/Arith/Cutsat/Types.lean | 20 +++--
 .../Meta/Tactic/Grind/Arith/Cutsat/Var.lean   | 12 +++
 src/Lean/Meta/Tactic/Grind/Internalize.lean   |  2 +-
 src/Lean/Meta/Tactic/Grind/Types.lean         |  6 +-
 tests/lean/run/grind_canon_insts.lean         | 12 ++-
 tests/lean/run/grind_cutsat_div_mod.lean      |  3 +
 tests/lean/run/grind_cutsat_nat_eq.lean       | 16 ++++
 tests/lean/run/grind_cutsat_nat_le.lean       | 12 +++
 tests/lean/run/grind_eq.lean                  |  1 +
 tests/lean/run/grind_lazy_ite.lean            |  8 +-
 tests/lean/run/grind_offset.lean              |  3 +
 tests/lean/run/grind_regression.lean          |  5 +-
 tests/lean/run/grind_t1.lean                  |  2 -
 21 files changed, 193 insertions(+), 51 deletions(-)
 create mode 100644 tests/lean/run/grind_cutsat_nat_eq.lean

diff --git a/src/Init/Data/Int/Linear.lean b/src/Init/Data/Int/Linear.lean
index 7bff1ccfc0..84032233f3 100644
--- a/src/Init/Data/Int/Linear.lean
+++ b/src/Init/Data/Int/Linear.lean
@@ -1757,6 +1757,16 @@ theorem dvd_norm_expr (ctx : Context) (d : Int) (e : Expr) (p : Poly)
     : p == e.norm → d ∣ e.denote ctx → d ∣ p.denote' ctx := by
   simp; intro; subst p; simp
 
+theorem eq_norm_expr (ctx : Context) (lhs rhs : Expr) (p : Poly)
+    : norm_eq_cert lhs rhs p → lhs.denote ctx = rhs.denote ctx → p.denote' ctx = 0 := by
+  intro h₁ h₂; rwa [norm_eq ctx lhs rhs p h₁] at h₂
+
+theorem not_eq_norm_expr (ctx : Context) (lhs rhs : Expr) (p : Poly)
+    : norm_eq_cert lhs rhs p → ¬ lhs.denote ctx = rhs.denote ctx → ¬ p.denote' ctx = 0 := by
+  simp [norm_eq_cert]
+  intro; subst p; simp
+  intro; rwa [Int.sub_eq_zero]
+
 end Int.Linear
 
 theorem Int.not_le_eq (a b : Int) : (¬a ≤ b) = (b + 1 ≤ a) := by
diff --git a/src/Init/Data/Int/OfNat.lean b/src/Init/Data/Int/OfNat.lean
index e8edd3824c..b01c341d89 100644
--- a/src/Init/Data/Int/OfNat.lean
+++ b/src/Init/Data/Int/OfNat.lean
@@ -67,4 +67,12 @@ theorem of_dvd (ctx : Context) (d : Nat) (e : Expr)
     : d ∣ e.denote ctx → Int.ofNat d ∣ e.denoteAsInt ctx := by
   simp [Expr.denoteAsInt_eq, Int.ofNat_dvd]
 
+theorem of_eq (ctx : Context) (lhs rhs : Expr)
+    : lhs.denote ctx = rhs.denote ctx → lhs.denoteAsInt ctx = rhs.denoteAsInt ctx := by
+  rw [Expr.eq ctx lhs rhs]; simp
+
+theorem of_not_eq (ctx : Context) (lhs rhs : Expr)
+    : ¬ lhs.denote ctx = rhs.denote ctx → ¬ lhs.denoteAsInt ctx = rhs.denoteAsInt ctx := by
+  rw [Expr.eq ctx lhs rhs]; simp
+
 end Int.OfNat
diff --git a/src/Init/Grind/Norm.lean b/src/Init/Grind/Norm.lean
index a057333ab9..4fd8368491 100644
--- a/src/Init/Grind/Norm.lean
+++ b/src/Init/Grind/Norm.lean
@@ -9,6 +9,7 @@ import Init.PropLemmas
 import Init.Classical
 import Init.ByCases
 import Init.Data.Int.Linear
+import Init.Data.Int.Pow
 
 namespace Lean.Grind
 /-!
@@ -124,6 +125,8 @@ init_grind_norm
   -- Int
   Int.lt_eq
   Int.emod_neg Int.ediv_zero Int.emod_zero
+  Int.natCast_add Int.natCast_mul Int.natCast_pow
+  Int.natCast_succ Int.natCast_zero
   -- GT GE
   ge_eq gt_eq
   -- Int op folding
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean
index a7785412d8..abadb8eb36 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean
@@ -64,5 +64,6 @@ builtin_initialize registerTraceClass `grind.debug.cutsat.subst
 builtin_initialize registerTraceClass `grind.debug.cutsat.getBestLower
 builtin_initialize registerTraceClass `grind.debug.cutsat.nat
 builtin_initialize registerTraceClass `grind.debug.cutsat.proof
+builtin_initialize registerTraceClass `grind.debug.cutsat.internalize
 
 end Lean
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DivMod.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DivMod.lean
index 0d3891ff12..0b280a8932 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DivMod.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DivMod.lean
@@ -22,15 +22,15 @@ private def expandDivMod (a : Expr) (b : Int) : GoalM Unit := do
 
 builtin_grind_propagator propagateDiv ↑HDiv.hDiv := fun e => do
   let_expr HDiv.hDiv _ _ _ inst a b ← e | return ()
-  unless (← isInstHDivInt inst) do return ()
-  let some b ← getIntValue? b | return ()
-  -- Remark: we currently do not consider the case where `b` is in the equivalence class of a numeral.
-  expandDivMod a b
+  if (← isInstHDivInt inst) then
+    let some b ← getIntValue? b | return ()
+    -- Remark: we currently do not consider the case where `b` is in the equivalence class of a numeral.
+    expandDivMod a b
 
 builtin_grind_propagator propagateMod ↑HMod.hMod := fun e => do
   let_expr HMod.hMod _ _ _ inst a b ← e | return ()
-  unless (← isInstHModInt inst) do return ()
-  let some b ← getIntValue? b | return ()
-  expandDivMod a b
+  if (← isInstHModInt inst) then
+    let some b ← getIntValue? b | return ()
+    expandDivMod a b
 
 end Lean.Meta.Grind.Arith.Cutsat
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/EqCnstr.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/EqCnstr.lean
index 6c4dcc08b7..96318c4b66 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/EqCnstr.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/EqCnstr.lean
@@ -234,17 +234,31 @@ private def exprAsPoly (a : Expr) : GoalM Poly := do
   else
     throwError "internal `grind` error, expression is not relevant to cutsat{indentExpr a}"
 
-@[export lean_process_cutsat_eq]
-def processNewEqImpl (a b : Expr) : GoalM Unit := do
-  trace[grind.debug.cutsat.eq] "{a} = {b}"
+private def processNewIntEq (a b : Expr) : GoalM Unit := do
   let p₁ ← exprAsPoly a
   let p₂ ← exprAsPoly b
   let p := p₁.combine (p₂.mul (-1))
   { p, h := .core a b p₁ p₂ : EqCnstr }.assert
 
-@[export lean_process_cutsat_eq_lit]
-def processNewEqLitImpl (a ke : Expr) : GoalM Unit := do
-  trace[grind.debug.cutsat.eq] "{a} = {ke}"
+private def processNewNatEq (a b : Expr) : GoalM Unit := do
+  let (lhs, rhs, ctx) ← Int.OfNat.toIntEq a b
+  let gen ← getGeneration a
+  let lhs' ← toLinearExpr (lhs.denoteAsIntExpr ctx) gen
+  let rhs' ← toLinearExpr (rhs.denoteAsIntExpr ctx) gen
+  let p := lhs'.sub rhs' |>.norm
+  let c := { p, h := .coreNat a b ctx lhs rhs lhs' rhs' : EqCnstr }
+  trace[grind.cutsat.assert.eq] "{← c.pp}"
+  c.assert
+
+@[export lean_process_cutsat_eq]
+def processNewEqImpl (a b : Expr) : GoalM Unit := do
+  trace[grind.debug.cutsat.eq] "{a} = {b}"
+  match (← foreignTerm? a), (← foreignTerm? b) with
+  | none, none => processNewIntEq a b
+  | some .nat, some .nat => processNewNatEq a b
+  | _, _ => return ()
+
+private def processNewIntLitEq (a ke : Expr) : GoalM Unit := do
   let some k ← getIntValue? ke | return ()
   let p₁ ← exprAsPoly a
   let c ← if k == 0 then
@@ -255,9 +269,14 @@ def processNewEqLitImpl (a ke : Expr) : GoalM Unit := do
     pure { p, h := .core a ke p₁ p₂ : EqCnstr }
   c.assert
 
-@[export lean_process_cutsat_diseq]
-def processNewDiseqImpl (a b : Expr) : GoalM Unit := do
-  trace[grind.debug.cutsat.diseq] "{a} ≠ {b}"
+@[export lean_process_cutsat_eq_lit]
+def processNewEqLitImpl (a ke : Expr) : GoalM Unit := do
+  trace[grind.debug.cutsat.eq] "{a} = {ke}"
+  match (← foreignTerm? a) with
+  | none => processNewIntLitEq a ke
+  | some .nat => processNewNatEq a ke
+
+private def processNewIntDiseq (a b : Expr) : GoalM Unit := do
   let p₁ ← exprAsPoly a
   let c ← if let some 0 ← getIntValue? b then
     pure { p := p₁, h := .core0 a b : DiseqCnstr }
@@ -267,10 +286,38 @@ def processNewDiseqImpl (a b : Expr) : GoalM Unit := do
     pure {p, h := .core a b p₁ p₂ : DiseqCnstr }
   c.assert
 
+private def processNewNatDiseq (a b : Expr) : GoalM Unit := do
+  let (lhs, rhs, ctx) ← Int.OfNat.toIntEq a b
+  let gen ← getGeneration a
+  let lhs' ← toLinearExpr (lhs.denoteAsIntExpr ctx) gen
+  let rhs' ← toLinearExpr (rhs.denoteAsIntExpr ctx) gen
+  let p := lhs'.sub rhs' |>.norm
+  let c := { p, h := .coreNat a b ctx lhs rhs lhs' rhs' : DiseqCnstr }
+  trace[grind.cutsat.assert.eq] "{← c.pp}"
+  c.assert
+
+@[export lean_process_cutsat_diseq]
+def processNewDiseqImpl (a b : Expr) : GoalM Unit := do
+  trace[grind.debug.cutsat.diseq] "{a} ≠ {b}"
+  match (← foreignTerm? a), (← foreignTermOrLit? b) with
+  | none, none => processNewIntDiseq a b
+  | some .nat, some .nat => processNewNatDiseq a b
+  | _, _ => return ()
+
 /-- Different kinds of terms internalized by this module. -/
 private inductive SupportedTermKind where
   | add | mul | num
 
+private def getKindAndType? (e : Expr) : Option (SupportedTermKind × Expr) :=
+  match_expr e with
+  | HAdd.hAdd α _ _ _ _ _ => some (.add, α)
+  | HMul.hMul α _ _ _ _ _ => some (.mul, α)
+  | OfNat.ofNat α _ _ => some (.num, α)
+  | Neg.neg α _ a =>
+    let_expr OfNat.ofNat _ _ _ := a | none
+    some (.num, α)
+  | _ => none
+
 private def isForbiddenParent (parent? : Option Expr) (k : SupportedTermKind) : Bool := Id.run do
   let some parent := parent? | return false
   let .const declName _ := parent.getAppFn | return false
@@ -280,16 +327,15 @@ private def isForbiddenParent (parent? : Option Expr) (k : SupportedTermKind) :
   | .mul => return declName == ``HMul.hMul
   | .num => return declName == ``HMul.hMul || declName == ``Eq
 
-private def internalizeCore (e : Expr) (parent? : Option Expr) (k : SupportedTermKind) : GoalM Unit := do
+private def internalizeInt (e : Expr) : GoalM Unit := do
   if (← get').terms.contains { expr := e } then return ()
-  if isForbiddenParent parent? k then return ()
   let p ← toPoly e
   markAsCutsatTerm e
   trace[grind.cutsat.internalize] "{aquote e}:= {← p.pp}"
   modify' fun s => { s with terms := s.terms.insert { expr := e } p }
 
 /--
-Internalizes an integer expression. Here are the different cases that are handled.
+Internalizes an integer (and `Nat`) expression. Here are the different cases that are handled.
 
 - `a + b` when `parent?` is not `+`, `≤`, or `∣`
 - `k * a` when `k` is a numeral and `parent?` is not `+`, `*`, `≤`, `∣`
@@ -298,14 +344,13 @@ Internalizes an integer expression. Here are the different cases that are handle
   back to the congruence closure module. Example: we have `f 5`, `f x`, `x - y = 3`, `y = 2`.
 -/
 def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
-  let k ← if (← isAdd e) then
-    pure .add
-  else if (← isMul e) then
-    pure .mul
-  else if (← getIntValue? e).isSome then
-    pure .num
-  else
-    return ()
-  internalizeCore e parent? k
+  let some (k, type) := getKindAndType? e | return ()
+  if isForbiddenParent parent? k then return ()
+  trace[grind.debug.cutsat.internalize] "{e} : {type}"
+  if type.isConstOf ``Int then
+    internalizeInt e
+  else if type.isConstOf ``Nat then
+    markForeignTerm e .nat
+
 
 end Lean.Meta.Grind.Arith.Cutsat
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
index 6d9c9a8a3d..79ea4cfd39 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Nat.lean
@@ -94,4 +94,11 @@ def toIntDvd? (e : Lean.Expr) : GoalM (Option (Nat × Expr × Array Lean.Expr))
   let (b, s) ← toOfNatExpr b |>.run {}
   return some (d, b, s.ctx)
 
+def toIntEq (lhs rhs : Lean.Expr) : MetaM (Expr × Expr × Array Lean.Expr) := do
+  let ((lhs, rhs), s) ← conv lhs rhs |>.run {}
+  return (lhs, rhs, s.ctx)
+where
+  conv (lhs rhs : Lean.Expr) : M (Expr × Expr) :=
+    return (← toOfNatExpr lhs, ← toOfNatExpr rhs)
+
 end Int.OfNat
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
index a121c335ae..4e9de0e483 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
@@ -109,6 +109,10 @@ partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
   | .core a b p₁ p₂ =>
     let h ← mkEqProof a b
     return mkApp6 (mkConst ``Int.Linear.eq_of_core) (← getContext) (← mkPolyDecl p₁) (← mkPolyDecl p₂) (← mkPolyDecl c'.p) reflBoolTrue h
+  | .coreNat a b ctx lhs rhs lhs' rhs' =>
+    let ctx ← mkNatCtxDecl ctx
+    let h := mkApp4 (mkConst ``Int.OfNat.of_eq) ctx (← mkNatExprDecl lhs) (← mkNatExprDecl rhs) (← mkEqProof a b)
+    return mkApp6 (mkConst ``Int.Linear.eq_norm_expr) (← getContext) (← mkExprDecl lhs') (← mkExprDecl rhs') (← mkPolyDecl c'.p) reflBoolTrue h
   | .norm c =>
     return mkApp5 (mkConst ``Int.Linear.eq_norm) (← getContext) (← mkPolyDecl c.p) (← mkPolyDecl c'.p) reflBoolTrue (← c.toExprProof)
   | .divCoeffs c =>
@@ -268,6 +272,10 @@ partial def DiseqCnstr.toExprProof (c' : DiseqCnstr) : ProofM Expr := caching c'
   | .core a b p₁ p₂ =>
     let h ← mkDiseqProof a b
     return mkApp6 (mkConst ``Int.Linear.diseq_of_core) (← getContext) (← mkPolyDecl p₁) (← mkPolyDecl p₂) (← mkPolyDecl c'.p) reflBoolTrue h
+  | .coreNat a b ctx lhs rhs lhs' rhs' =>
+    let ctx ← mkNatCtxDecl ctx
+    let h := mkApp4 (mkConst ``Int.OfNat.of_not_eq) ctx (← mkNatExprDecl lhs) (← mkNatExprDecl rhs) (← mkDiseqProof a b)
+    return mkApp6 (mkConst ``Int.Linear.not_eq_norm_expr) (← getContext) (← mkExprDecl lhs') (← mkExprDecl rhs') (← mkPolyDecl c'.p) reflBoolTrue h
   | .norm c =>
     return mkApp5 (mkConst ``Int.Linear.diseq_norm) (← getContext) (← mkPolyDecl c.p) (← mkPolyDecl c'.p) reflBoolTrue (← c.toExprProof)
   | .divCoeffs c =>
@@ -378,7 +386,7 @@ private def markAsFound (fvarId : FVarId) : CollectDecVarsM Unit := do
 mutual
 partial def EqCnstr.collectDecVars (c' : EqCnstr) : CollectDecVarsM Unit := do unless (← alreadyVisited c') do
   match c'.h with
-  | .core0 .. | .core .. => return () -- Equalities coming from the core never contain cutsat decision variables
+  | .core0 .. | .core .. | .coreNat .. => return () -- Equalities coming from the core never contain cutsat decision variables
   | .norm c | .divCoeffs c => c.collectDecVars
   | .subst _ c₁ c₂ | .ofLeGe c₁ c₂ => c₁.collectDecVars; c₂.collectDecVars
 
@@ -410,7 +418,7 @@ partial def LeCnstr.collectDecVars (c' : LeCnstr) : CollectDecVarsM Unit := do u
 
 partial def DiseqCnstr.collectDecVars (c' : DiseqCnstr) : CollectDecVarsM Unit := do unless (← alreadyVisited c') do
   match c'.h with
-  | .core0 .. | .core .. => return () -- Disequalities coming from the core never contain cutsat decision variables
+  | .core0 .. | .core .. | .coreNat .. => return () -- Disequalities coming from the core never contain cutsat decision variables
   | .norm c | .divCoeffs c | .neg c => c.collectDecVars
   | .subst _ c₁ c₂  => c₁.collectDecVars; c₂.collectDecVars
 
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
index 48bbf07f08..ea5fe07474 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
@@ -84,6 +84,7 @@ inductive EqCnstrProof where
     `p₁` and `p₂` are the polynomials corresponding to `a` and `b`.
     -/
     core (a b : Expr) (p₁ p₂ : Poly)
+  | coreNat (a b : Expr) (ctx : Array Expr) (lhs rhs : Int.OfNat.Expr) (lhs' rhs' : Int.Linear.Expr)
   | norm (c : EqCnstr)
   | divCoeffs (c : EqCnstr)
   | subst (x : Var) (c₁ : EqCnstr) (c₂ : EqCnstr)
@@ -184,6 +185,7 @@ inductive DiseqCnstrProof where
     `p₁` and `p₂` are the polynomials corresponding to `a` and `b`.
     -/
     core (a b : Expr) (p₁ p₂ : Poly)
+  | coreNat (a b : Expr) (ctx : Array Expr) (lhs rhs : Int.OfNat.Expr) (lhs' rhs' : Int.Linear.Expr)
   | norm (c : DiseqCnstr)
   | divCoeffs (c : DiseqCnstr)
   | neg (c : DiseqCnstr)
@@ -216,6 +218,10 @@ instance : Inhabited CooperSplit where
 
 abbrev VarSet := RBTree Var compare
 
+inductive ForeignType where
+  | nat
+  deriving BEq
+
 /-- State of the cutsat procedure. -/
 structure State where
   /-- Mapping from variables to their denotations. -/
@@ -223,6 +229,15 @@ structure State where
   /-- Mapping from `Expr` to a variable representing it. -/
   varMap  : PHashMap ENodeKey Var := {}
   /--
+  Mapping from terms (e.g., `x + 2*y + 2`, `3*x`, `5`) to polynomials representing them.
+  These are terms used to propagate equalities between this module and the congruence closure module.
+  -/
+  terms : PHashMap ENodeKey Poly := {}
+  /--
+  Foreign terms (e.g., `Nat`). They are also marked using `markAsCutsatTerm`.
+  -/
+  foreignTerms : PHashMap ENodeKey ForeignType := {}
+  /--
   Mapping from variables to divisibility constraints. Recall that we keep the divisibility constraint in solved form.
   Thus, we have at most one divisibility per variable. -/
   dvds : PArray (Option DvdCnstr) := {}
@@ -251,11 +266,6 @@ structure State where
   -/
   elimStack : List Var := []
   /--
-  Mapping from terms (e.g., `x + 2*y + 2`, `3*x`, `5`) to polynomials representing them.
-  These are terms used to propagate equalities between this module and the congruence closure module.
-  -/
-  terms : PHashMap ENodeKey Poly := {}
-  /--
   Mapping from variable to occurrences. For example, an entry `x ↦ {y, z}` means that `x` may occur in `dvdCnstrs`, `lowers`, or `uppers` of
   variables `y` and `z`.
   If `x` occurs in `dvdCnstrs[y]`, `lowers[y]`, or `uppers[y]`, then `y` is in `occurs[x]`, but the reverse is not true.
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
index fb059802b7..4745e70cf4 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean
@@ -10,10 +10,22 @@ import Lean.Meta.Tactic.Grind.Canon
 
 namespace Lean.Meta.Grind.Arith.Cutsat
 
+def markForeignTerm (e : Expr) (t : ForeignType) : GoalM Unit := do
+  modify' fun s => { s with foreignTerms := s.foreignTerms.insert { expr := e} t }
+  markAsCutsatTerm e
+
+def foreignTerm? (e : Expr) : GoalM (Option ForeignType) := do
+  return (← get').foreignTerms.find? { expr := e }
+
+def foreignTermOrLit? (e : Expr) : GoalM (Option ForeignType) := do
+  if isNatNum e then return some .nat
+  foreignTerm? e
+
 private def assertNatCast (e : Expr) : GoalM Unit := do
   let_expr NatCast.natCast _ inst a := e | return ()
   let_expr instNatCastInt := inst | return ()
   pushNewProof <| mkApp (mkConst ``Int.Linear.natCast_nonneg) a
+  markForeignTerm a .nat
 
 private def assertHelpers (e : Expr) : GoalM Unit := do
   assertNatCast e
diff --git a/src/Lean/Meta/Tactic/Grind/Internalize.lean b/src/Lean/Meta/Tactic/Grind/Internalize.lean
index 6c0dbe1b11..cd25a4fcf5 100644
--- a/src/Lean/Meta/Tactic/Grind/Internalize.lean
+++ b/src/Lean/Meta/Tactic/Grind/Internalize.lean
@@ -247,6 +247,7 @@ private partial def internalizeImpl (e : Expr) (generation : Nat) (parent? : Opt
     else if e.isAppOfArity ``Grind.MatchCond 1 then
       internalizeMatchCond e generation
     else e.withApp fun f args => do
+      mkENode e generation
       checkAndAddSplitCandidate e
       pushCastHEqs e
       addMatchEqns f generation
@@ -270,7 +271,6 @@ private partial def internalizeImpl (e : Expr) (generation : Nat) (parent? : Opt
           let arg := args[i]
           internalize arg generation e
           registerParent e arg
-      mkENode e generation
       addCongrTable e
       updateAppMap e
       Arith.internalize e parent?
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index aa8a2ee666..f9c27747f6 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -936,12 +936,12 @@ def propagateCutsatDiseq (lhs rhs : Expr) : GoalM Unit := do
   let some lhs ← get? lhs | return ()
   let some rhs ← get? rhs | return ()
   -- Recall that core can take care of disequalities of the form `1≠2`.
-  unless isIntNum lhs && isIntNum rhs do
+  unless isNum lhs && isNum rhs do
     Arith.Cutsat.processNewDiseq lhs rhs
 where
   get? (a : Expr) : GoalM (Option Expr) := do
     let root ← getRootENode a
-    if isIntNum root.self then
+    if isNum root.self then
       return some root.self
     return root.cutsat?
 
@@ -961,7 +961,7 @@ def markAsCutsatTerm (e : Expr) : GoalM Unit := do
   let root ← getRootENode e
   if let some e' := root.cutsat? then
     Arith.Cutsat.processNewEq e e'
-  else if isIntNum root.self && !isSameExpr e root.self then
+  else if isNum root.self && !isSameExpr e root.self then
     Arith.Cutsat.processNewEqLit e root.self
   else
     setENode root.self { root with cutsat? := some e }
diff --git a/tests/lean/run/grind_canon_insts.lean b/tests/lean/run/grind_canon_insts.lean
index 287f735ed6..08ef53c91f 100644
--- a/tests/lean/run/grind_canon_insts.lean
+++ b/tests/lean/run/grind_canon_insts.lean
@@ -53,12 +53,14 @@ theorem left_comm [CommMonoid α] (a b c : α) : a * (b * c) = b * (a * c) := by
 open Lean Meta Elab Tactic Grind in
 def fallback : Fallback := do
   let nodes ← filterENodes fun e => return e.self.isApp && e.self.isAppOf ``HMul.hMul
-  trace[Meta.debug] "{nodes.toList.map (·.self)}"
+  trace[Meta.debug] "{nodes.map (·.self) |>.qsort Expr.lt}"
   (← get).mvarId.admit
 
 set_option trace.Meta.debug true
 
-/-- info: [Meta.debug] [b * c, a * (b * c), d * (b * c)] -/
+/--
+info: [Meta.debug] [-1 * NatCast.natCast (a * (b * c)), -1 * NatCast.natCast (d * (b * c)), a * (b * c), b * c, d * (b * c)]
+-/
 #guard_msgs (info) in
 example (a b c d : Nat) : b * (a * c) = d * (b * c) → False := by
   rw [left_comm] -- Introduces a new (non-canonical) instance for `Mul Nat`
@@ -68,7 +70,11 @@ example (a b c d : Nat) : b * (a * c) = d * (b * c) → False := by
 set_option pp.notation false in
 set_option pp.explicit true in
 /--
-info: [Meta.debug] [@HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b a,
+info: [Meta.debug] [@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) (@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]
 -/
 #guard_msgs (info) in
diff --git a/tests/lean/run/grind_cutsat_div_mod.lean b/tests/lean/run/grind_cutsat_div_mod.lean
index 5c2023431d..0c97fd1613 100644
--- a/tests/lean/run/grind_cutsat_div_mod.lean
+++ b/tests/lean/run/grind_cutsat_div_mod.lean
@@ -23,3 +23,6 @@ set_option trace.grind.cutsat.model true in
 example (x y : Int) : x % 2 + y = 3 → x ≤ 5 → x > 4 → y = 1 := by
   fail_if_success grind
   sorry
+
+example (x y : Int) : x = y / 2 → y % 2 = 0 → y - 2*x = 0 := by
+  grind
diff --git a/tests/lean/run/grind_cutsat_nat_eq.lean b/tests/lean/run/grind_cutsat_nat_eq.lean
new file mode 100644
index 0000000000..70a826d9eb
--- /dev/null
+++ b/tests/lean/run/grind_cutsat_nat_eq.lean
@@ -0,0 +1,16 @@
+set_option grind.warning false
+
+example (a b c : Nat) : a = 0 → b = 0 → c ≥ a + b := by
+  grind
+
+example (a b c : Nat) : a + b = 0 → a ≤ b + c + a → a ≤ c := by
+  grind
+
+example (a b : Nat) (_ : 2*a + 3*b = 0) (_ : 2 ∣ 3*b + 1) : False := by
+  grind
+
+example (a b c : Nat) : a + 2*b = 0 → b + c + b = 0 → a = c := by
+  grind
+
+example (a : Nat) : a ≤ 2 → a ≠ 0 → a ≠ 1 → a ≠ 2 → False := by
+  grind
diff --git a/tests/lean/run/grind_cutsat_nat_le.lean b/tests/lean/run/grind_cutsat_nat_le.lean
index c1c8b9643b..2dbf4d0749 100644
--- a/tests/lean/run/grind_cutsat_nat_le.lean
+++ b/tests/lean/run/grind_cutsat_nat_le.lean
@@ -15,3 +15,15 @@ open Int.Linear Int.OfNat
 #print ex1
 #print ex2
 #print ex3
+
+example (a b c : Nat) : Int.ofNat (a + b) = 0 → a + b + b ≤ c → b ≤ c := by
+  grind
+
+example (a b c : Nat) : Int.ofNat (a + b) = 0 → a + b + b ≤ c → Int.ofNat b ≤ c := by
+  grind
+
+example (a b c : Nat) : a + b < c → c ≥ 0 := by
+  grind
+
+example (a b : Int) : a + b = Int.ofNat 2 → a - 2 = -b := by
+  grind
diff --git a/tests/lean/run/grind_eq.lean b/tests/lean/run/grind_eq.lean
index 28ba3f95ed..bb785be30b 100644
--- a/tests/lean/run/grind_eq.lean
+++ b/tests/lean/run/grind_eq.lean
@@ -72,6 +72,7 @@ info: [grind.assert] x1 = appV a_2 b
 [grind.assert] ¬HEq x2 x4
 [grind.ematch.instance] appV_assoc: HEq (appV a_2 (appV b c)) (appV (appV a_2 b) c)
 [grind.assert] HEq (appV a_2 (appV b c)) (appV (appV a_2 b) c)
+[grind.assert] -1 * NatCast.natCast a ≤ 0
 -/
 #guard_msgs (info) in
 example : x1 = appV a b → x2 = appV x1 c → x3 = appV b c → x4 = appV a x3 → HEq x2 x4 := by
diff --git a/tests/lean/run/grind_lazy_ite.lean b/tests/lean/run/grind_lazy_ite.lean
index d7571794ba..7aafb0df92 100644
--- a/tests/lean/run/grind_lazy_ite.lean
+++ b/tests/lean/run/grind_lazy_ite.lean
@@ -13,12 +13,8 @@ example : f 5 m > 0 := by
   fail_if_success grind (splits := 0) [f.eq_def]
   sorry
 
-/--
-info: [grind.ematch.instance] f.eq_def: f 5 m = if 5 < m then f (5 + 1) m + 5 else 5
-[grind.ematch.instance] f.eq_def: f 6 m = if 6 < m then f (6 + 1) m + 6 else 6
--/
+/-- info: [grind.ematch.instance] f.eq_def: f 5 m = if 5 < m then f (5 + 1) m + 5 else 5 -/
 #guard_msgs (info) in
 set_option trace.grind.ematch.instance true in
 example : f 5 m > 0 := by
-  fail_if_success grind (splits := 1) [f.eq_def]
-  sorry
+  grind (splits := 1) [f.eq_def]
diff --git a/tests/lean/run/grind_offset.lean b/tests/lean/run/grind_offset.lean
index 1de16dc996..29335f370b 100644
--- a/tests/lean/run/grind_offset.lean
+++ b/tests/lean/run/grind_offset.lean
@@ -51,6 +51,7 @@ info: [grind.assert] f (c + 2) = a
 [grind.assert] ¬a = g (g (f c))
 [grind.ematch.instance] f.eq_2: f (c + 1).succ = g (f (c + 1))
 [grind.assert] f (c + 2) = g (f (c + 1))
+[grind.assert] -1 * NatCast.natCast c ≤ 0
 [grind.ematch.instance] f.eq_2: f c.succ = g (f c)
 [grind.assert] f (c + 1) = g (f c)
 -/
@@ -76,6 +77,8 @@ info: [grind.assert] foo (c + 1) = a
 [grind.assert] ¬a = g (foo b)
 [grind.ematch.instance] foo.eq_3: foo b.succ.succ = g (foo b)
 [grind.assert] foo (b + 2) = g (foo b)
+[grind.assert] -1 * NatCast.natCast b ≤ 0
+[grind.assert] -1 * NatCast.natCast c ≤ 0
 -/
 #guard_msgs (info) in
 example : foo (c + 1) = a → c = b + 1 → a = g (foo b) := by
diff --git a/tests/lean/run/grind_regression.lean b/tests/lean/run/grind_regression.lean
index c82f47f19b..7423f81e9d 100644
--- a/tests/lean/run/grind_regression.lean
+++ b/tests/lean/run/grind_regression.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Marcus Rossel, Kim Morrison
 -/
 import Lean.Elab.Term
-
+set_option grind.warning false
 /-!
 These tests are originally from the `lean-egg` repository at
 https://github.com/marcusrossel/lean-egg and were written by Marcus Rossel.
@@ -97,6 +97,9 @@ example {x : Nat} (h₁ : x = y) (h₂ : x = y → 1 = 2) : 1 = 2 := by
 example (h : ∀ p : Prop, p → 1 = id 1) : 1 = id 1 := by
   grind only [id]
 
+example (h : ∀ p : Prop, p → (1 : Int) = id 1) : (1 : Int) = id 1 := by
+  grind only [id]
+
 example {p q r : Prop} (h₁ : p) (h₂ : p ↔ q) (h₃ : q → (p ↔ r)) : p ↔ r := by
   grind
 
diff --git a/tests/lean/run/grind_t1.lean b/tests/lean/run/grind_t1.lean
index e29c5f6184..7e4fe39633 100644
--- a/tests/lean/run/grind_t1.lean
+++ b/tests/lean/run/grind_t1.lean
@@ -117,7 +117,6 @@ end dite_propagator_test
 info: [grind.eqc] x = 2 * a
 [grind.eqc] y = x
 [grind.eqc] (y = 2 * a) = False
-[grind.eqc] (y = 2 * a) = True
 -/
 #guard_msgs (info) in
 set_option trace.grind.eqc true in
@@ -128,7 +127,6 @@ example (a : Nat) : let x := a + a; y = x → y = a + a := by
 info: [grind.eqc] x = 2 * a
 [grind.eqc] y = x
 [grind.eqc] (y = 2 * a) = False
-[grind.eqc] (y = 2 * a) = True
 -/
 #guard_msgs (info) in
 set_option trace.grind.eqc true in