feat: proof production for divisibility constraint solver in grind (#7138)

This PR implements proof generation for the divisibility constraint
solver in `grind`.

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

@@ -10,36 +10,30 @@ import Lean.Meta.Tactic.Grind.Arith.Cutsat.Var
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.Proof
 
 namespace Lean.Meta.Grind.Arith.Cutsat
-/--
-`gcdExt a b` returns the triple `(g, α, β)` such that
-- `g = gcd a b` (with `g ≥ 0`), and
-- `g = α * a + β * β`.
--/
-partial def gcdExt (a b : Int) : Int × Int × Int :=
-  if b = 0 then
-    (a.natAbs, if a = 0 then 0 else a / a.natAbs, 0)
-  else
-    let (g, α, β) := gcdExt b (a % b)
-    (g, β, α - (a / b) * β)
-
 abbrev DvdCnstrWithProof.isUnsat (cₚ : DvdCnstrWithProof) : Bool :=
   cₚ.c.isUnsat
 
 abbrev DvdCnstrWithProof.isTrivial (cₚ : DvdCnstrWithProof) : Bool :=
   cₚ.c.isTrivial
 
-def DvdCnstrWithProof.norm (cₚ : DvdCnstrWithProof) : DvdCnstrWithProof :=
-  let cₚ := if cₚ.c.isSorted then cₚ else { cₚ with c.p := cₚ.c.p.norm, h := .norm cₚ }
-  let g := cₚ.c.p.gcdCoeffs cₚ.c.k
-  if cₚ.c.p.getConst % g == 0 then
-    { cₚ with c := cₚ.c.div g, h := .divCoeffs cₚ }
+def mkDvdCnstrWithProof (c : DvdCnstr) (h : DvdCnstrProof) : GoalM DvdCnstrWithProof := do
+  return { c, h, id := (← mkCnstrId) }
+
+def DvdCnstrWithProof.norm (cₚ : DvdCnstrWithProof) : GoalM DvdCnstrWithProof := do
+  let cₚ ← if cₚ.c.isSorted then
+    pure cₚ
   else
-    cₚ
+    mkDvdCnstrWithProof { k := cₚ.c.k, p := cₚ.c.p.norm } (.norm cₚ)
+  let g := cₚ.c.p.gcdCoeffs cₚ.c.k
+  if cₚ.c.p.getConst % g == 0 && g != 1 then
+    mkDvdCnstrWithProof (cₚ.c.div g) (.divCoeffs cₚ)
+  else
+    return cₚ
 
 /-- Asserts divisibility constraint. -/
 partial def assertDvdCnstr (cₚ : DvdCnstrWithProof) : GoalM Unit := withIncRecDepth do
   if (← isInconsistent) then return ()
-  let cₚ := cₚ.norm
+  let cₚ ← cₚ.norm
   if cₚ.isUnsat then
     trace[grind.cutsat.dvd.unsat] "{← cₚ.denoteExpr}"
     withProofContext do
@@ -71,13 +65,13 @@ partial def assertDvdCnstr (cₚ : DvdCnstrWithProof) : GoalM Unit := withIncRec
       -/
       let α_d₂_p₁ := p₁.mul (α*d₂)
       let β_d₁_p₂ := p₂.mul (β*d₁)
-      let combine := { c.k := d₁*d₂, c.p := .add d x (α_d₂_p₁.combine β_d₁_p₂), h := .solveCombine cₚ cₚ' }
+      let combine ← mkDvdCnstrWithProof { k := d₁*d₂, p := .add d x (α_d₂_p₁.combine β_d₁_p₂) } (.solveCombine cₚ cₚ')
       trace[grind.cutsat.dvd.solve.combine] "{← combine.denoteExpr}"
       modify' fun s => { s with dvdCnstrs := s.dvdCnstrs.set x none}
       assertDvdCnstr combine
       let a₂_p₁ := p₁.mul a₂
       let a₁_p₂ := p₂.mul (-a₁)
-      let elim := { c.k := d, c.p := a₂_p₁.combine a₁_p₂, h := .solveElim cₚ cₚ' }
+      let elim ← mkDvdCnstrWithProof { k := d, p := a₂_p₁.combine a₁_p₂ } (.solveElim cₚ cₚ')
       trace[grind.cutsat.dvd.solve.elim] "{← elim.denoteExpr}"
       assertDvdCnstr elim
     else
@@ -92,7 +86,7 @@ builtin_grind_propagator propagateDvd ↓Dvd.dvd := fun e => do
       return ()
   if (← isEqTrue e) then
     let p ← toPoly b
-    let cₚ := { c.k := k, c.p := p, h := .expr (← mkOfEqTrue (← mkEqTrueProof e)) }
+    let cₚ ← mkDvdCnstrWithProof { k, p } (.expr (← mkOfEqTrue (← mkEqTrueProof e)))
     trace[grind.cutsat.assert.dvd] "{← cₚ.denoteExpr}"
     assertDvdCnstr cₚ
   else if (← isEqFalse e) then

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

@@ -8,8 +8,39 @@ import Lean.Meta.Tactic.Grind.Arith.Cutsat.Util
 
 namespace Lean.Meta.Grind.Arith.Cutsat
 
-def DvdCnstrWithProof.toExprProof (cₚ : DvdCnstrWithProof) : ProofM Expr := do
-  -- TODO
-  mkSorry (← cₚ.denoteExpr) false
+private def DvdCnstrWithProof.get_d_a (cₚ : DvdCnstrWithProof) : GoalM (Int × Int) := do
+  let d := cₚ.c.k
+  let .add a _ _ := cₚ.c.p
+    | throwError "internal `grind` error, unexpected divisibility constraint {indentExpr (← cₚ.denoteExpr)}"
+  return (d, a)
+
+partial def DvdCnstrWithProof.toExprProof' (cₚ : DvdCnstrWithProof) : ProofM Expr := cₚ.caching do
+  match cₚ.h with
+  | .expr h =>
+    return h
+  | .norm cₚ' =>
+    return mkApp5 (mkConst ``Int.Linear.DvdCnstr.of_isNorm) (← 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ₚ')
+  | .solveCombine cₚ₁ cₚ₂ =>
+    let (d₁, a₁) ← cₚ₁.get_d_a
+    let (d₂, a₂) ← cₚ₂.get_d_a
+    let (d, α, β) := gcdExt (a₁*d₂) (a₂*d₁)
+    return mkApp10 (mkConst ``Int.Linear.DvdCnstr.solve_combine)
+      (← getContext) (toExpr cₚ₁.c) (toExpr cₚ₂.c) (toExpr cₚ.c)
+      (toExpr d) (toExpr α) (toExpr β) reflBoolTrue
+      (← toExprProof' cₚ₁) (← toExprProof' cₚ₂)
+  | .solveElim cₚ₁ cₚ₂ =>
+    let (d₁, a₁) ← cₚ₁.get_d_a
+    let (d₂, a₂) ← cₚ₂.get_d_a
+    let (d, _, _) := gcdExt (a₁*d₂) (a₂*d₁)
+    return mkApp8 (mkConst ``Int.Linear.DvdCnstr.solve_elim)
+      (← getContext) (toExpr cₚ₁.c) (toExpr cₚ₂.c) (toExpr cₚ.c)
+      (toExpr d) reflBoolTrue
+      (← toExprProof' cₚ₁) (← toExprProof' cₚ₂)
+
+partial def DvdCnstrWithProof.toExprProof (cₚ : DvdCnstrWithProof) : ProofM Expr := do
+  mkExpectedTypeHint (← toExprProof' cₚ) (← cₚ.denoteExpr)
 
 end Lean.Meta.Grind.Arith.Cutsat

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

@@ -17,8 +17,10 @@ export Int.Linear (Var Poly RelCnstr DvdCnstr)
 mutual
 /-- A divisibility constraint and its justification/proof. -/
 structure DvdCnstrWithProof where
-  c : DvdCnstr
-  h : DvdCnstrProof
+  c  : DvdCnstr
+  h  : DvdCnstrProof
+  /-- Unique id for caching proofs in `ProofM` -/
+  id : Nat
 
 inductive DvdCnstrProof where
   | expr (h : Expr)
@@ -38,6 +40,8 @@ structure State where
   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. -/
   dvdCnstrs : PArray (Option DvdCnstrWithProof) := {}
+  /-- Next unique id for a constraint. -/
+  nextCnstrId : Nat := 0
   deriving Inhabited
 
 end Lean.Meta.Grind.Arith.Cutsat

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

@@ -7,7 +7,6 @@ prelude
 import Lean.Meta.Tactic.Grind.Types
 
 namespace Int.Linear
-
 def Poly.isZero : Poly → Bool
   | .num 0 => true
   | _ => false
@@ -35,6 +34,17 @@ def DvdCnstr.isTrivial (c : DvdCnstr) : Bool :=
 end Int.Linear
 
 namespace Lean.Meta.Grind.Arith.Cutsat
+/--
+`gcdExt a b` returns the triple `(g, α, β)` such that
+- `g = gcd a b` (with `g ≥ 0`), and
+- `g = α * a + β * β`.
+-/
+partial def gcdExt (a b : Int) : Int × Int × Int :=
+  if b = 0 then
+    (a.natAbs, if a = 0 then 0 else a / a.natAbs, 0)
+  else
+    let (g, α, β) := gcdExt b (a % b)
+    (g, β, α - (a / b) * β)
 
 def get' : GoalM State := do
   return (← get).arith.cutsat
@@ -45,6 +55,11 @@ def get' : GoalM State := do
 def getVars : GoalM (PArray Expr) :=
   return (← get').vars
 
+def mkCnstrId : GoalM Nat := do
+  let id := (← get').nextCnstrId
+  modify' fun s => { s with nextCnstrId := id + 1 }
+  return id
+
 def DvdCnstrWithProof.denoteExpr (cₚ : DvdCnstrWithProof) : GoalM Expr := do
   let vars ← getVars
   cₚ.c.denoteExpr (vars[·]!)
@@ -56,14 +71,28 @@ def toContextExpr : GoalM Expr := do
   else
     return RArray.toExpr (mkConst ``Int) id (RArray.leaf (mkIntLit 0))
 
+structure ProofM.State where
+  cache : Std.HashMap Nat Expr := {}
+
 /-- Auxiliary monad for constructing cutsat proofs. -/
-abbrev ProofM := ReaderT Expr GoalM
+abbrev ProofM := ReaderT Expr (StateRefT ProofM.State GoalM)
 
 /-- Returns a Lean expression representing the variable context used to construct cutsat proofs. -/
 abbrev getContext : ProofM Expr := do
   read
 
+abbrev caching (id : Nat) (k : ProofM Expr) : ProofM Expr := do
+  if let some h := (← get).cache[id]? then
+    return h
+  else
+    let h ← k
+    modify fun s => { s with cache := s.cache.insert id h }
+    return h
+
+abbrev DvdCnstrWithProof.caching (c : DvdCnstrWithProof) (k : ProofM Expr) : ProofM Expr :=
+  Cutsat.caching c.id k
+
 abbrev withProofContext (x : ProofM α) : GoalM α := do
-  x (← toContextExpr)
+  x (← toContextExpr) |>.run' {}
 
 end Lean.Meta.Grind.Arith.Cutsat

tests/lean/run/grind_cutsat_div_1.lean

@@ -0,0 +1,16 @@
+set_option grind.warning false
+set_option pp.structureInstances false
+open Int.Linear
+
+theorem ex₁ (a : Int) (h₁ : 2 ∣ a) (h₂ : 2 ∣ 2*a + 1 - a) : False := by
+  grind
+
+theorem ex₂ (a b : Int) (h₀ : 2 ∣ a + 1) (h₁ : 2 ∣ b + a) (h₂ : 2 ∣ b + 2*a) : False := by
+  grind
+
+theorem ex₃ (a b : Int) (_ : 2 ∣ a + 1) (h₁ : 3 ∣ a + 3*b + a) (h₂ : 2 ∣ 3*b + a + 3 - b) (h₃ : 3 ∣ 3 * b + 2 * a + 1) : False := by
+  grind
+
+#print ex₁
+#print ex₂
+#print ex₃