From 26946ddc7f3bcd20fb8bff94e1dbeb39e3fb08d3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 15 Jun 2025 13:50:43 -0400
Subject: [PATCH] feat: `Inv.lean` for grind linarith (#8800)

---
 src/Lean/Meta/Tactic/Grind/Arith/Inv.lean     |  2 +
 .../Meta/Tactic/Grind/Arith/Linear/Inv.lean   | 93 +++++++++++++++++++
 .../Meta/Tactic/Grind/Arith/Linear/Util.lean  |  5 +
 tests/lean/run/grind_field_div.lean           |  1 +
 tests/lean/run/grind_field_norm.lean          |  1 +
 tests/lean/run/grind_linarith_1.lean          |  1 +
 tests/lean/run/grind_linarith_2.lean          |  1 +
 .../lean/run/grind_module_normalization.lean  |  1 +
 8 files changed, 105 insertions(+)
 create mode 100644 src/Lean/Meta/Tactic/Grind/Arith/Linear/Inv.lean

diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Inv.lean b/src/Lean/Meta/Tactic/Grind/Arith/Inv.lean
index 8b04cff5eb..de7a1f4a9e 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Inv.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Inv.lean
@@ -6,11 +6,13 @@ Authors: Leonardo de Moura
 prelude
 import Lean.Meta.Tactic.Grind.Arith.Offset
 import Lean.Meta.Tactic.Grind.Arith.Cutsat.Inv
+import Lean.Meta.Tactic.Grind.Arith.Linear.Inv
 
 namespace Lean.Meta.Grind.Arith
 
 def checkInvariants : GoalM Unit := do
   Offset.checkInvariants
   Cutsat.checkInvariants
+  Linear.checkInvariants
 
 end Lean.Meta.Grind.Arith
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/Inv.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Inv.lean
new file mode 100644
index 0000000000..d41afcf083
--- /dev/null
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Inv.lean
@@ -0,0 +1,93 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Meta.Tactic.Grind.Arith.Linear.Util
+
+namespace Lean.Meta.Grind.Arith.Linear
+
+/--
+Returns `true` if the variables in the given polynomial are sorted
+in decreasing order.
+-/
+def _root_.Lean.Grind.Linarith.Poly.isSorted (p : Poly) : Bool :=
+  go none p
+where
+  go : Option Var → Poly → Bool
+  | _,      .nil       => true
+  | none,   .add _ y p => go (some y) p
+  | some x, .add _ y p => x > y && go (some y) p
+
+/-- Returns `true` if all coefficients are not `0`. -/
+def _root_.Lean.Grind.Linarith.Poly.checkCoeffs : Poly → Bool
+  | .nil => true
+  | .add k _ p => k != 0 && checkCoeffs p
+
+def _root_.Lean.Grind.Linarith.Poly.checkCnstrOf (p : Poly) (x : Var) : LinearM Unit := do
+  assert! p.isSorted
+  assert! p.checkCoeffs
+  unless (← inconsistent) do
+    -- p.checkNoElimVars
+    -- p.checkOccs
+    pure ()
+  let .add _ y _ := p | unreachable!
+  assert! x == y
+
+def checkLeCnstrs (css : PArray (PArray IneqCnstr)) (isLower : Bool) : LinearM Unit := do
+  let mut x := 0
+  for cs in css do
+    for c in cs do
+      c.p.checkCnstrOf x
+      let .add a _ _ := c.p | unreachable!
+      assert! isLower == (a < 0)
+    x := x + 1
+  return ()
+
+def checkLowers : LinearM Unit := do
+  let s ← getStruct
+  assert! s.lowers.size == s.vars.size
+  checkLeCnstrs s.lowers (isLower := true)
+
+def checkUppers : LinearM Unit := do
+  let s ← getStruct
+  assert! s.uppers.size == s.vars.size
+  checkLeCnstrs s.uppers (isLower := false)
+
+def checkDiseqCnstrs : LinearM Unit := do
+  let s ← getStruct
+  assert! s.vars.size == s.diseqs.size
+  let mut x := 0
+  for cs in s.diseqs do
+    for c in cs do
+      c.p.checkCnstrOf x
+    x := x + 1
+  return ()
+
+def checkVars : LinearM Unit := do
+  let s ← getStruct
+  let mut num := 0
+  for ({ expr }, var) in s.varMap do
+    if h : var < s.vars.size then
+      let expr' := s.vars[var]
+      assert! isSameExpr expr expr'
+    else
+      unreachable!
+    num := num + 1
+  assert! s.vars.size == num
+
+def checkStructInvs : LinearM Unit := do
+  checkVars
+  checkLowers
+  checkUppers
+  checkDiseqCnstrs
+
+def checkInvariants : GoalM Unit := do
+  unless grind.debug.get (← getOptions) do return ()
+  for structId in [: (← get').structs.size] do
+    LinearM.run structId do
+      assert! (← getStructId) == structId
+      checkStructInvs
+
+end Lean.Meta.Grind.Arith.Linear
diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Linear/Util.lean b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Util.lean
index 5794059bb0..8950e89da8 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Linear/Util.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Linear/Util.lean
@@ -189,4 +189,9 @@ def resetAssignmentFrom (x : Var) : LinearM Unit := do
 def getVar (x : Var) : LinearM Expr :=
   return (← getStruct).vars[x]!
 
+/-- Returns `true` if the linarith state is inconsistent. -/
+def inconsistent : LinearM Bool := do
+  if (← isInconsistent) then return true
+  return (← getStruct).conflict?.isSome
+
 end Lean.Meta.Grind.Arith.Linear
diff --git a/tests/lean/run/grind_field_div.lean b/tests/lean/run/grind_field_div.lean
index 91dc4cd9d7..690fe3e458 100644
--- a/tests/lean/run/grind_field_div.lean
+++ b/tests/lean/run/grind_field_div.lean
@@ -1,4 +1,5 @@
 open Lean Grind
+set_option grind.debug true
 
 example [Field α] [IsCharP α 0] (a b c : α) : a/3 = b → c = a/3 → a/2 + a/2 = b + 2*c  := by
   grind
diff --git a/tests/lean/run/grind_field_norm.lean b/tests/lean/run/grind_field_norm.lean
index e7ea50296f..4a8674613c 100644
--- a/tests/lean/run/grind_field_norm.lean
+++ b/tests/lean/run/grind_field_norm.lean
@@ -1,4 +1,5 @@
 open Lean.Grind
+set_option grind.debug true
 
 variable (R : Type u) [Field R]
 
diff --git a/tests/lean/run/grind_linarith_1.lean b/tests/lean/run/grind_linarith_1.lean
index 0457e5461a..6f4d1f9e60 100644
--- a/tests/lean/run/grind_linarith_1.lean
+++ b/tests/lean/run/grind_linarith_1.lean
@@ -1,4 +1,5 @@
 open Lean.Grind
+set_option grind.debug true
 
 example [IntModule α] [Preorder α] [IntModule.IsOrdered α] (a b : α)
     : (2:Int)*a + b < b + a + a → False := by
diff --git a/tests/lean/run/grind_linarith_2.lean b/tests/lean/run/grind_linarith_2.lean
index 7b563a9f1a..76b120bf3c 100644
--- a/tests/lean/run/grind_linarith_2.lean
+++ b/tests/lean/run/grind_linarith_2.lean
@@ -1,4 +1,5 @@
 open Lean Grind
+set_option grind.debug true
 
 example [IntModule α] [Preorder α] [IntModule.IsOrdered α] (a b : α)
     : a + b = b + a  := by
diff --git a/tests/lean/run/grind_module_normalization.lean b/tests/lean/run/grind_module_normalization.lean
index 6a88a3f81a..f848d45b2a 100644
--- a/tests/lean/run/grind_module_normalization.lean
+++ b/tests/lean/run/grind_module_normalization.lean
@@ -1,5 +1,6 @@
 open Lean Grind
 variable (R : Type u) [IntModule R]
+set_option grind.debug true
 
 example (a b : R) : a + b = b + a := by grind
 example (a : R) : a + 0 = a := by grind