From d1aba29b573e080b46fbd9ca6c544ac350a5d523 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Fri, 21 Feb 2025 08:17:32 -0800
Subject: [PATCH] feat: model construction for divisibility constraints in
 cutsat (#7176)

This PR implements model construction for divisibility constraints in
the cutsat procedure.
---
 .../Tactic/Grind/Arith/Cutsat/DvdCnstr.lean   |  5 +++
 .../Meta/Tactic/Grind/Arith/Cutsat/Proof.lean |  2 +
 .../Tactic/Grind/Arith/Cutsat/Search.lean     | 40 ++++++++++++++++++-
 .../Meta/Tactic/Grind/Arith/Cutsat/Types.lean |  1 +
 tests/lean/run/grind_cutsat_div_1.lean        | 39 ++++++++++++++++++
 5 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DvdCnstr.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DvdCnstr.lean
index 6e8ac98abb..1ed83c0d2c 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DvdCnstr.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/DvdCnstr.lean
@@ -16,6 +16,9 @@ abbrev DvdCnstrWithProof.isUnsat (cₚ : DvdCnstrWithProof) : Bool :=
 abbrev DvdCnstrWithProof.isTrivial (cₚ : DvdCnstrWithProof) : Bool :=
   cₚ.c.isTrivial
 
+abbrev DvdCnstrWithProof.satisfied (cₚ : DvdCnstrWithProof) : GoalM LBool :=
+  cₚ.c.satisfied
+
 def mkDvdCnstrWithProof (c : DvdCnstr) (h : DvdCnstrProof) : GoalM DvdCnstrWithProof := do
   return { c, h, id := (← mkCnstrId) }
 
@@ -46,6 +49,8 @@ partial def assertDvdCnstr (cₚ : DvdCnstrWithProof) : GoalM Unit := withIncRec
     let d₁ := cₚ.c.k
     let .add a₁ x p₁ := cₚ.c.p
       | throwError "internal `grind` error, unexpected divisibility constraint {indentExpr (← cₚ.denoteExpr)}"
+    if (← cₚ.satisfied) == .false then
+      resetAssignmentFrom x
     if let some cₚ' := (← get').dvdCnstrs[x]! then
       trace[grind.cutsat.dvd.solve] "{← cₚ.denoteExpr}, {← cₚ'.denoteExpr}"
       let d₂ := cₚ'.c.k
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
index 0f676d1fa0..344c695396 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Proof.lean
@@ -20,6 +20,8 @@ partial def DvdCnstrWithProof.toExprProof' (cₚ : DvdCnstrWithProof) : ProofM E
     return h
   | .norm cₚ' =>
     return mkApp5 (mkConst ``Int.Linear.DvdCnstr.of_isNorm) (← getContext) (toExpr cₚ'.c) (toExpr cₚ.c) reflBoolTrue (← toExprProof' cₚ')
+  | .elim cₚ' =>
+    return mkApp5 (mkConst ``Int.Linear.DvdCnstr.elim) (← getContext) (toExpr cₚ'.c) (toExpr cₚ.c) reflBoolTrue (← toExprProof' cₚ')
   | .divCoeffs cₚ' =>
     let k := cₚ'.c.p.gcdCoeffs cₚ'.c.k
     return mkApp6 (mkConst ``Int.Linear.DvdCnstr.of_isEqv) (← getContext) (toExpr cₚ'.c) (toExpr cₚ.c) (toExpr k) reflBoolTrue (← toExprProof' cₚ')
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Search.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Search.lean
index a7773aa260..0c2ea8a471 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Search.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Search.lean
@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.Var
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.Util
+import Lean.Meta.Tactic.Grind.Arith.Cutsat.DvdCnstr
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.RelCnstr
 
 namespace Lean.Meta.Grind.Arith.Cutsat
@@ -45,6 +46,29 @@ def getBestUpper? (x : Var) : GoalM (Option (Int × RelCnstrWithProof)) := do
       best? := some (upper', cₚ)
   return best?
 
+def getDvdSolutions? (cₚ : DvdCnstrWithProof) : GoalM (Option (Int × Int)) := do
+  let d := cₚ.c.k
+  let .add a _ p := cₚ.c.p
+    | throwError "`grind` internal error, unexpected divisibility constraint{indentExpr (← cₚ.denoteExpr)}"
+  let some b ← p.eval?
+    | throwError "`grind` internal error, divisibility constraint is not ready to be solved{indentExpr (← cₚ.denoteExpr)}"
+  -- We must solve `d ∣ a*x + b`
+  let g := d.gcd a
+  if b % g != 0 then
+    return none -- no solutions
+  let d := d / g
+  let a := a / g
+  let b := b / g
+  -- `α*a + β*d = 1`
+  -- `α*a = 1 (mod d)`
+  let (_, α, _β) := gcdExt a d
+  -- `a'*a = 1 (mod d)`
+  let a' := if α < 0 then α % d else α
+  -- `a*x = -b (mod d)`
+  -- `x = -b*a' (mod d)`
+  -- `x = k*d + -b*a'` for any k
+  return some (d, -b*a')
+
 private partial def setAssignment (x : Var) (v : Int) : GoalM Unit := do
   if x == (← get').assignment.size then
     trace[grind.cutsat.assign] "{(← getVar x)} := {v}"
@@ -60,11 +84,18 @@ def resolveLowerUpperConflict (c₁ c₂ : RelCnstrWithProof) : GoalM Unit := do
   trace[grind.cutsat.conflict] "{← c₁.denoteExpr}, {← c₂.denoteExpr}"
   return ()
 
+def resolveDvdConflict (cₚ : DvdCnstrWithProof) : GoalM Unit := do
+  trace[grind.cutsat.conflict] "{← cₚ.denoteExpr}"
+  let d := cₚ.c.k
+  let .add a _ p := cₚ.c.p
+    | throwError "`grind` internal error, unexpected divisibility constraint{indentExpr (← cₚ.denoteExpr)}"
+  assertDvdCnstr (← mkDvdCnstrWithProof { k := a.gcd d, p } (.elim cₚ))
+
 def decideVar (x : Var) : GoalM Unit := do
   let lower? ← getBestLower? x
   let upper? ← getBestUpper? x
-  let div? := (← get').dvdCnstrs[x]!
-  match lower?, upper?, div? with
+  let dvd? := (← get').dvdCnstrs[x]!
+  match lower?, upper?, dvd? with
   | none, none, none =>
     setAssignment x 0
   | some (lower, _), none, none =>
@@ -79,6 +110,11 @@ def decideVar (x : Var) : GoalM Unit := do
       resolveLowerUpperConflict c₁ c₂
       -- TODO: remove the following
       setAssignment x 0
+  | none, none, some cₚ =>
+    if let some (_, v) ← getDvdSolutions? cₚ then
+      setAssignment x v
+    else
+      resolveDvdConflict cₚ
   | _, _, _ =>
     -- TODO: cases containing a divisibility constraint.
     -- TODO: remove the following
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
index 409a782061..50f8a4a4cb 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Types.lean
@@ -28,6 +28,7 @@ inductive DvdCnstrProof where
   | divCoeffs (c : DvdCnstrWithProof)
   | solveCombine (c₁ c₂ : DvdCnstrWithProof)
   | solveElim (c₁ c₂ : DvdCnstrWithProof)
+  | elim (c : DvdCnstrWithProof)
 
 structure RelCnstrWithProof where
   c  : RelCnstr
diff --git a/tests/lean/run/grind_cutsat_div_1.lean b/tests/lean/run/grind_cutsat_div_1.lean
index fc491fc1b3..c6879e2e6c 100644
--- a/tests/lean/run/grind_cutsat_div_1.lean
+++ b/tests/lean/run/grind_cutsat_div_1.lean
@@ -18,3 +18,42 @@ theorem ex₄ (f : Int → Int) (a b : Int) (_ : 2 ∣ f (f a) + 1) (h₁ : 3 
 #print ex₂
 #print ex₃
 #print ex₄
+
+/--
+info: [grind.cutsat.assign] a := -3
+[grind.cutsat.assign] b := 7
+-/
+#guard_msgs (info) in -- finds the model without any backtracking
+set_option trace.grind.cutsat.assign true in
+set_option trace.grind.cutsat.conflict true in
+example (a b : Int) (_ : 2 ∣ a + 3) (_ : 3 ∣ a + b - 4) : False := by
+  fail_if_success grind
+  sorry
+
+
+/--
+info: [grind.cutsat.dvd.update] 2 ∣ a + 3
+[grind.cutsat.dvd.update] 3 ∣ 3 * b + a + -4
+[grind.cutsat.assign] a := -3
+[grind.cutsat.conflict] 3 ∣ 3 * b + a + -4
+[grind.cutsat.dvd.solve] 3 ∣ a + -4, 2 ∣ a + 3
+[grind.cutsat.dvd.update] 6 ∣ a + 17
+[grind.cutsat.assign] a := -17
+[grind.cutsat.assign] b := 0
+-/
+#guard_msgs (info) in
+set_option trace.grind.cutsat.assign true in
+set_option trace.grind.cutsat.dvd true in
+set_option trace.grind.cutsat.dvd.solve.elim false in
+set_option trace.grind.cutsat.dvd.solve.combine false in
+set_option trace.grind.cutsat.dvd.trivial false in
+set_option trace.grind.cutsat.conflict true in
+/-
+In this example, cutsat fails to extend the model to `b` after assigning `a := - 3`.
+Then, it learns a new constraaint `6 ∣ a + 17`, finds a new assignment for `a := -17`
+and then satisfies `3 ∣ a + 3*b - 4`, but assigning `b := 0`.
+So model (aka counter-example) is `a := -17` and `b := 0`.
+-/
+example (a b : Int) (_ : 2 ∣ a + 3) (_ : 3 ∣ a + 3*b - 4) : False := by
+  fail_if_success grind
+  sorry