feat: nonlinear monomials in grind cutsat (#9996)

This PR improves support for nonlinear monomials in `grind cutsat`. For
example, given a monomial `a * b`, if `cutsat` discovers that `a = 2`,
it now propagates that `a * b = 2 * b`.
Recall that nonlinear monomials like `a * b` are treated as variables in
`cutsat`, a procedure designed for linear integer arithmetic.

Example:
```lean
example (a : Nat) (ha : a < 8) (b c : Nat) : 2 ≤ b → c = 1 → b ≤ c + 1 → a * b < 8 * b := by
  grind

example (x y z w : Int) : z * x * y = 4 → x = z + w → z = 1 → w = 2 → False := by
  grind
```

src/Init/Data/Int/Linear.lean

@@ -2144,6 +2144,45 @@ theorem natCast_sub (x y : Nat)
     rw [Int.ofNat_le] at this
     rw [Lean.Omega.Int.ofNat_sub_eq_zero this]
 
+/-! Helper theorem for linearizing nonlinear terms -/
+
+@[expose] noncomputable def var_eq_cert (x : Var) (k : Int) (p : Poly) : Bool :=
+  Poly.rec (fun _ => false)
+    (fun k₁ x' p' _ => Poly.rec (fun k₂ => k₁ != 0 && x == x' && k == -k₂/k₁) (fun _ _ _ _ => false) p')
+    p
+
+theorem var_eq (ctx : Context) (x : Var) (k : Int) (p : Poly) : var_eq_cert x k p → p.denote' ctx = 0 → x.denote ctx = k := by
+  simp [var_eq_cert]; cases p <;> simp; next k₁ x' p' =>
+  cases p' <;> simp; next k₂ =>
+  intro h₁ _ _; subst x' k; intro h₂
+  replace h₂ := Int.neg_eq_of_add_eq_zero h₂
+  rw [Int.neg_eq_comm] at h₂
+  rw [← h₂]; clear h₂; simp; rw [Int.mul_comm, Int.mul_ediv_cancel]
+  assumption
+
+@[expose] noncomputable def of_var_eq_mul_cert (x : Var) (k : Int) (y : Var) (p : Poly) : Bool :=
+  p.beq' (.add 1 x (.add (-k) y (.num 0)))
+
+theorem of_var_eq_mul (ctx : Context) (x : Var) (k : Int) (y : Var) (p : Poly) : of_var_eq_mul_cert x k y p → x.denote ctx = k * y.denote ctx → p.denote' ctx = 0 := by
+  simp [of_var_eq_mul_cert]; intro _ h; subst p; simp [h]
+  rw [Int.neg_mul, ← Int.sub_eq_add_neg, Int.sub_self]
+
+@[expose] noncomputable def of_var_eq_cert (x : Var) (k : Int) (p : Poly) : Bool :=
+  p.beq' (.add 1 x (.num (-k)))
+
+theorem of_var_eq (ctx : Context) (x : Var) (k : Int) (p : Poly) : of_var_eq_cert x k p → x.denote ctx = k → p.denote' ctx = 0 := by
+  simp [of_var_eq_cert]; intro _ h; subst p; simp [h]
+  rw [← Int.sub_eq_add_neg, Int.sub_self]
+
+theorem eq_one_mul (a : Int) : a = 1*a := by simp
+theorem mul_eq_kk (a b k₁ k₂ k : Int) (h₁ : a = k₁) (h₂ : b = k₂) (h₃ : k₁*k₂ == k) : a*b = k := by simp_all
+theorem mul_eq_kkx (a b k₁ k₂ c k : Int) (h₁ : a = k₁) (h₂ : b = k₂*c) (h₃ : k₁*k₂ == k) : a*b = k*c := by
+  simp at h₃; rw [h₁, h₂, ← Int.mul_assoc, h₃]
+theorem mul_eq_kxk (a b k₁ c k₂ k : Int) (h₁ : a = k₁*c) (h₂ : b = k₂) (h₃ : k₁*k₂ == k) : a*b = k*c := by
+  simp at h₃; rw [h₁, h₂, Int.mul_comm, ← Int.mul_assoc, Int.mul_comm k₂, h₃]
+theorem mul_eq_zero_left (a b : Int) (h : a = 0) : a*b = 0 := by simp [*]
+theorem mul_eq_zero_right (a b : Int) (h : b = 0) : a*b = 0 := by simp [*]
+
 end Int.Linear
 
 theorem Int.not_le_eq (a b : Int) : (¬a ≤ b) = (b + 1 ≤ a) := by

src/Lean/Meta/Tactic/Grind/Arith/Cutsat.lean

@@ -28,6 +28,7 @@ public section
 namespace Lean
 
 builtin_initialize registerTraceClass `grind.cutsat
+builtin_initialize registerTraceClass `grind.cutsat.nonlinear
 builtin_initialize registerTraceClass `grind.cutsat.model
 builtin_initialize registerTraceClass `grind.cutsat.assert
 builtin_initialize registerTraceClass `grind.cutsat.assert.trivial

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/EqCnstr.lean

@@ -169,6 +169,78 @@ private def updateDiseqs (a : Int) (x : Var) (c : EqCnstr) (y : Var) : GoalM Uni
     c₂.assert
     if (← inconsistent) then return ()
 
+/-- Returns `some k` if the given variable has been eliminate with equality `y = k` -/
+private def isVarEqConst? (y : Var) : GoalM (Option (Int × EqCnstr)) := do
+  let some c := (← get').elimEqs[y]! | return none
+  let .add k₁ _ (.num k₂) := c.p | return none
+  if k₂ % k₁ != 0 then return none
+  return some (- k₂/k₁, c)
+
+/-- Returns `some k` if `e` is represented by a variable `y` that has been eliminate with equality `y = k` -/
+private def isExprEqConst? (e : Expr) : GoalM (Option (Int × EqCnstr)) := do
+  let some x := (← get').varMap.find? { expr := e } | return none
+  isVarEqConst? x
+
+structure PropagateMul.State where
+  a? : Option Expr := none
+  k  : Int := 1
+  cs : Array (Expr × Int × EqCnstr) := #[]
+  n  : Nat := 0 -- num unknowns
+
+private partial def propagateNonlinearMul (x : Var) : GoalM Bool := do
+  let e ← getVar x
+  let (_, { a?, k, cs, n }) ← go e |>.run {}
+  if k == 0 || n == 0 then
+    let c := { p := .add 1 x (.num (-k)), h := .mul none cs : EqCnstr }
+    c.assert
+    return true
+  else if n > 1 then
+    return false
+  else
+    let some a := a? | unreachable!
+    -- x = k*a
+    let y ← mkVar a
+    let c := { p := .add 1 x (.add (-k) y (.num 0)), h := .mul a? cs : EqCnstr }
+    c.assert
+    return true
+where
+  goVar (e : Expr) : StateT PropagateMul.State GoalM Unit := do
+    if let some (k', c) ← isExprEqConst? e then
+      modify fun { a?, k, cs, n } => { a?, k := k*k', cs := cs.push (e, k', c), n }
+    else if (← get).n == 0 then
+      modify fun { k, cs, .. } => { a? := some e, k, cs, n := 1 }
+    else
+      modify fun { k, cs, a?, n } => { a?, k, cs, n := n + 1 }
+
+  go (e : Expr) : StateT PropagateMul.State GoalM Unit := do
+    let_expr HMul.hMul _ _ _ i a b := e | goVar e
+    if (← isInstHMulInt i) then
+      go a; go b
+    else
+      goVar e
+
+private def propagateNonlinearDiv (x : Var) : GoalM Bool := do
+  trace[Meta.debug] "{← getVar x}" -- TODO
+  return true
+
+private def propagateNonlinearMod (x : Var) : GoalM Bool := do
+  trace[Meta.debug] "{← getVar x}" -- TODO
+  return true
+
+@[export lean_cutsat_propagate_nonlinear]
+def propagateNonlinearTermImpl (y : Var) (x : Var) : GoalM Bool := do
+  unless (← isVarEqConst? y).isSome do return false
+  match_expr (← getVar x) with
+  | HMul.hMul _ _ _ _ _ _ => propagateNonlinearMul x
+  | HDiv.hDiv _ _ _ _ _ _ => propagateNonlinearDiv x
+  | HMod.hMod _ _ _ _ _ _ => propagateNonlinearMod x
+  | _ => return false
+
+def propagateNonlinearTerms (y : Var) : GoalM Unit := do
+  let some occs := (← get').nonlinearOccs.find? y | return ()
+  let occs ← occs.filterM fun x => return !(← propagateNonlinearTermImpl y x)
+  modify' fun s => { s with nonlinearOccs := s.nonlinearOccs.insert y occs }
+
 private def updateElimEqs (a : Int) (x : Var) (c : EqCnstr) (y : Var) : GoalM Unit := do
   if (← inconsistent) then return ()
   assert! x != y
@@ -178,6 +250,7 @@ private def updateElimEqs (a : Int) (x : Var) (c : EqCnstr) (y : Var) : GoalM Un
   let c₂ := { p := c₂.p.mul a |>.combine (c.p.mul (-b)), h := .subst x c₂ c : EqCnstr }
   trace[grind.debug.cutsat.elimEq] "updated: {← getVar y}, {← c₂.pp}"
   modify' fun s => { s with elimEqs := s.elimEqs.set y (some c₂) }
+  propagateNonlinearTerms y
 
 private def updateOccsAt (k : Int) (x : Var) (c : EqCnstr) (y : Var) : GoalM Unit := do
   updateDvdCnstr k x c y
@@ -251,6 +324,8 @@ def EqCnstr.assertImpl (c : EqCnstr) : GoalM Unit := do
     let p := c.p.insert (-k) x
     let d := Int.ofNat k.natAbs
     { d, p, h := .ofEq x c : DvdCnstr }.assert
+  if (← inconsistent) then return ()
+  propagateNonlinearTerms x
 
 private def exprAsPoly (a : Expr) : GoalM Poly := do
   if let some k ← getIntValue? a then

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean

@@ -102,6 +102,10 @@ private def getVarMap : ProofM (PHashMap ExprPtr Var) := do
   else
     return (← get').varMap
 
+private def getVarOf (e : Expr) : ProofM Var := do
+  let some x := (← getVarMap).find? { expr := e } | throwError "`grind` internal error, missing cutsat variable{indentExpr e}"
+  return x
+
 private abbrev caching (c : α) (k : ProofM Expr) : ProofM Expr := do
   let addr := unsafe (ptrAddrUnsafe c).toUInt64 >>> 2
   if let some h := (← get).cache[addr]? then
@@ -212,14 +216,22 @@ private def DvdCnstr.get_d_a (c : DvdCnstr) : GoalM (Int × Int) := do
   let .add a _ _ := c.p | c.throwUnexpected
   return (d, a)
 
+private def getCurrVars : ProofM (PArray Expr) := do
+  if (← read).unordered then pure (← get').vars' else getVars
+
 /--
 Similar to `denoteExpr'`, but takes into account the `unordered` flag in the `ProofM` context.
 Recall that if `unordered` is `true`, we should use `vars'`
 -/
 private def _root_.Int.Linear.Poly.denoteExprUsingCurrVars (p : Poly) : ProofM Expr := do
-  let vars ← if (← read).unordered then pure (← get').vars' else getVars
+  let vars ← getCurrVars
   return (← p.denoteExpr (vars[·]!))
 
+inductive MulEqProof where
+  | const (k : Int) (h : Expr)
+  | mulVar (k : Int) (a : Expr) (h : Expr)
+  | none
+
 mutual
 partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
   trace[grind.debug.cutsat.proof] "{← c'.pp}"
@@ -233,7 +245,7 @@ partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
     let h := mkApp toIntThm (← mkEqProof a b)
     return mkApp6 (mkConst ``Int.Linear.eq_norm_expr) (← getContext) (← mkExprDecl lhs) (← mkExprDecl rhs) (← mkPolyDecl c'.p) eagerReflBoolTrue h
   | .defn e p =>
-    let some x := (← getVarMap).find? { expr := e } | throwError "`grind` internal error, missing cutsat variable{indentExpr e}"
+    let x ← getVarOf e
     return mkApp6 (mkConst ``Int.Linear.eq_def) (← getContext) (← mkVarDecl x) (← mkPolyDecl p) (← mkPolyDecl c'.p) eagerReflBoolTrue (← mkEqRefl e)
   | .defnNat h x e =>
     return mkApp6 (mkConst ``Int.Linear.eq_def') (← getContext) (← mkVarDecl x) (← mkExprDecl e) (← mkPolyDecl c'.p) eagerReflBoolTrue h
@@ -258,7 +270,7 @@ partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
     return mkApp5 (mkConst ``Int.Linear.eq_norm_poly) (← getContext) (← mkPolyDecl c.p) (← mkPolyDecl c'.p) h (← c.toExprProof)
   | .defnCommRing e p re rp p' =>
     let h := mkApp4 (mkConst ``Grind.CommRing.norm_int) (← getRingContext) (← mkRingExprDecl re) (← mkRingPolyDecl rp) eagerReflBoolTrue
-    let some x := (← getVarMap).find? { expr := e } | throwError "`grind` internal error, missing cutsat variable{indentExpr e}"
+    let x ← getVarOf e
     return mkApp8 (mkConst ``Int.Linear.eq_def_norm) (← getContext)
       (← mkVarDecl x) (← mkPolyDecl p) (← mkPolyDecl p') (← mkPolyDecl c'.p)
       eagerReflBoolTrue (← mkEqRefl e) h
@@ -266,6 +278,54 @@ partial def EqCnstr.toExprProof (c' : EqCnstr) : ProofM Expr := caching c' do
     let h₂ := mkApp4 (mkConst ``Grind.CommRing.norm_int) (← getRingContext) (← mkRingExprDecl re) (← mkRingPolyDecl rp) eagerReflBoolTrue
     return mkApp9 (mkConst ``Int.Linear.eq_def'_norm) (← getContext) (← mkVarDecl x) (← mkExprDecl e')
       (← mkPolyDecl p) (← mkPolyDecl p') (← mkPolyDecl c'.p) eagerReflBoolTrue h₁ h₂
+  | .mul a? cs =>
+    let .add _ x _ := c'.p | c'.throwUnexpected
+    mkMulEqProof x a? cs c'
+
+partial def mkMulEqProof (x : Var) (a? : Option Expr) (cs : Array (Expr × Int × EqCnstr)) (c' : EqCnstr) : ProofM Expr := do
+  let h ← go (← getCurrVars)[x]!
+  match h with
+  | .const k h =>
+    return mkApp6 (mkConst ``Int.Linear.of_var_eq) (← getContext) (← mkVarDecl x) (toExpr k) (← mkPolyDecl c'.p) eagerReflBoolTrue h
+  | .mulVar k a h =>
+    assert! a? == some a
+    let y ← getVarOf a
+    return mkApp7 (mkConst ``Int.Linear.of_var_eq_mul) (← getContext) (← mkVarDecl x) (toExpr k) (← mkVarDecl y) (← mkPolyDecl c'.p) eagerReflBoolTrue h
+  | .none =>
+    throwError "`grind` internal error, cutsat failed to construct proof for multiplication equality"
+where
+  goVar (e : Expr) : ProofM MulEqProof := do
+    if some e == a? then
+      return .mulVar 1 e (mkApp (mkConst ``Int.Linear.eq_one_mul) e)
+    else
+      let some (_, k, c) := cs.find? fun (e', _, _) => e' == e | return .none
+      let x ← getVarOf e
+      let h := mkApp6 (mkConst ``Int.Linear.var_eq) (← getContext) (← mkVarDecl x) (toExpr k) (← mkPolyDecl c.p) eagerReflBoolTrue (← c.toExprProof)
+      return .const k h
+
+  go (e : Expr) : ProofM MulEqProof := do
+    let_expr HMul.hMul _ _ _ i a b := e | goVar e
+    if !(← isInstHMulInt i) then goVar e else
+    let ha ← go a
+    if let .const 0 h := ha then
+      return .const 0 (mkApp3 (mkConst ``Int.Linear.mul_eq_zero_left) a b h)
+    let hb ← go b
+    if let .const 0 h := hb then
+      return .const 0 (mkApp3 (mkConst ``Int.Linear.mul_eq_zero_right) a b h)
+    match ha, hb with
+    | .const k₁ h₁, .const k₂ h₂ =>
+      let k := k₁*k₂
+      let h := mkApp8 (mkConst ``Int.Linear.mul_eq_kk) a b (toExpr k₁) (toExpr k₂) (toExpr k) h₁ h₂ eagerReflBoolTrue
+      return .const k h
+    | .const k₁ h₁, .mulVar k₂ c h₂ =>
+      let k := k₁*k₂
+      let h := mkApp9 (mkConst ``Int.Linear.mul_eq_kkx) a b (toExpr k₁) (toExpr k₂) c (toExpr k) h₁ h₂ eagerReflBoolTrue
+      return .mulVar k c h
+    | .mulVar k₁ c h₁, .const k₂ h₂ =>
+      let k := k₁*k₂
+      let h := mkApp9 (mkConst ``Int.Linear.mul_eq_kxk) a b (toExpr k₁) c (toExpr k₂) (toExpr k) h₁ h₂ eagerReflBoolTrue
+      return .mulVar k c h
+    | _, _ => return .none
 
 partial def DvdCnstr.toExprProof (c' : DvdCnstr) : ProofM Expr := caching c' do
   trace[grind.debug.cutsat.proof] "{← c'.pp}"
@@ -522,6 +582,7 @@ partial def EqCnstr.collectDecVars (c' : EqCnstr) : CollectDecVarsM Unit := do u
   | .defnCommRing .. | .defnNatCommRing .. | .coreToInt .. => return () -- Equalities coming from the core never contain cutsat decision variables
   | .commRingNorm c .. | .reorder c | .norm c | .divCoeffs c => c.collectDecVars
   | .subst _ c₁ c₂ | .ofLeGe c₁ c₂ => c₁.collectDecVars; c₂.collectDecVars
+  | .mul _ cs => cs.forM fun (_, _, c) => c.collectDecVars
 
 partial def CooperSplit.collectDecVars (s : CooperSplit) : CollectDecVarsM Unit := do unless (← alreadyVisited s) do
   s.pred.c₁.collectDecVars

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean

@@ -102,6 +102,7 @@ inductive EqCnstrProof where
   | commRingNorm (c : EqCnstr) (e : CommRing.RingExpr) (p : CommRing.Poly)
   | defnCommRing (e : Expr) (p : Poly) (re : CommRing.RingExpr) (rp : CommRing.Poly) (p' : Poly)
   | defnNatCommRing (h : Expr) (x : Var) (e' : Int.Linear.Expr) (p : Poly) (re : CommRing.RingExpr) (rp : CommRing.Poly) (p' : Poly)
+  | mul (a? : Option Expr) (cs : Array (Expr × Int × EqCnstr))
 
 /-- A divisibility constraint and its justification/proof. -/
 structure DvdCnstr where
@@ -343,6 +344,16 @@ structure State where
   `usedCommRing` is `true` if the `CommRing` has been used to normalize expressions.
   -/
   usedCommRing : Bool := false
+  /--
+  Mapping from terms to variables representing nonlinear terms.
+  For example, suppose the denotation of variable `x` is the nonlinear term `a*b*c`,
+  and `y` is the nonlinear term `a / d`. Then the mapping contains the entries
+  - `a ↦ [x, y]`
+  - `b ↦ [x]`
+  - `c ↦ [x]`
+  - `d ↦ [y]`
+  -/
+  nonlinearOccs : PHashMap Var (List Var) := {}
   deriving Inhabited
 
 end Lean.Meta.Grind.Arith.Cutsat

src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Var.lean

@@ -16,6 +16,42 @@ public section
 
 namespace Lean.Meta.Grind.Arith.Cutsat
 
+@[extern "lean_cutsat_propagate_nonlinear"]
+opaque propagateNonlinearTerm (y : Var) (x : Var) : GoalM Bool
+
+private def isNonlinearTerm (e : Expr) : MetaM Bool := do
+  match_expr e with
+  | HMul.hMul _ _ _ i _ _ => isInstHMulInt i
+  | HDiv.hDiv _ _ _ i _ b => pure (← getIntValue? b).isNone <&&> isInstHDivInt i
+  | HMod.hMod _ _ _ i _ b => pure (← getIntValue? b).isNone <&&> isInstHModInt i
+  | _ => return false
+
+private def registerNonlinearOcc (arg : Expr) (x : Var) : GoalM Unit := do
+  let y ← mkVar arg
+  if (← get').elimEqs[y]!.isSome then
+    if (← propagateNonlinearTerm y x) then
+      return ()
+  let y ← mkVar arg
+  let occs := (← get').nonlinearOccs.find? y |>.getD []
+  unless x ∈ occs do
+  modify' fun s => { s with nonlinearOccs := s.nonlinearOccs.insert y (x::occs) }
+
+private partial def registerNonlinearOccsAt (e : Expr) (x : Var) : GoalM Unit := do
+  match_expr e with
+  | HMul.hMul _ _ _ _ a b => go a; go b
+  | HDiv.hDiv _ _ _ _ a b => registerNonlinearOcc a x; registerNonlinearOcc b x
+  | HMod.hMod _ _ _ _ a b => registerNonlinearOcc a x; registerNonlinearOcc b x
+  | _ => return ()
+where
+  go (e : Expr) : GoalM Unit := do
+    match_expr e with
+    | HMul.hMul _ _ _ i a b =>
+      if (← isInstHMulInt i) then
+        go a; go b
+      else
+        registerNonlinearOcc e x
+    | _ => registerNonlinearOcc e x
+
 @[export lean_grind_cutsat_mk_var]
 def mkVarImpl (expr : Expr) : GoalM Var := do
   if let some var := (← get').varMap.find? { expr } then
@@ -36,6 +72,8 @@ def mkVarImpl (expr : Expr) : GoalM Var := do
   assertNatCast expr var
   assertNonneg expr var
   assertToIntBounds expr var
+  if (← isNonlinearTerm expr) then
+    registerNonlinearOccsAt expr var
   return var
 
 def isInt (e : Expr) : GoalM Bool := do

tests/lean/run/grind_linearize.lean

@@ -0,0 +1,32 @@
+example (x y z w : Int) : z * x * y = 4 → x = z + w → z = 1 → w = 2 → False := by
+  grind -ring
+
+example (x y z w : Int) : y * z * x = 4 → x = z + w → z = 1 → w = 2 → False := by
+  grind -ring
+
+example (x y z w : Int) : y * z * x * w = 4 → x = 0 → False := by
+  grind -ring
+
+example (x y z w : Int) : x = z + w → z = 1 → w = 2 → z * x * y = 4 → False := by
+  grind -ring
+
+example (x y z : Int) : x = 3 → z = 1 → z * x * y = 4 → False := by
+  grind -ring
+
+example (x y z w : Int) : x = 1 → y = 2 → z = 3 → w = 4 → x * y * z * w = 24 := by
+  grind -ring
+
+example (x y z w : Int) : 1 ≤ x → x < 2 → y = 2 → z = 3 → w = 4 → x * y * z * w = 24 := by
+  grind -ring
+
+example (x y z w : Int) : x = 1 → y = 2 → z = 3 → w = 4 → x * (y * z) * w = 24 := by
+  grind -ring
+
+example (x y : Int) : 1 ≤ x → x ≤ 1 → 2 ≤ y → y ≤ 2 → x * y = 2 := by
+  grind -ring
+
+example (a : Nat) (ha : a < 8) (b : Nat) (hb : b = 2) : a * b < 8 * b := by
+  grind -ring
+
+example (a : Nat) (ha : a < 8) (b c : Nat) : 2 ≤ b → c = 1 → b ≤ c + 1 → a * b < 8 * b := by
+  grind -ring