From de2f4e2199f2309f2864271909c0df6f239351ca Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 9 Jan 2025 12:07:12 -0800
Subject: [PATCH] feat: improve inequality offset support theorems for `grind`

This PR improves the theorems used to justify the steps performed by
the inequality offset module. See new test for examples of how they are
going to be used.
---
 src/Init/Grind/Offset.lean                 | 81 +++++++++++++++-------
 tests/lean/run/grind_offset_ineq_thms.lean | 31 +++++++++
 2 files changed, 88 insertions(+), 24 deletions(-)
 create mode 100644 tests/lean/run/grind_offset_ineq_thms.lean

diff --git a/src/Init/Grind/Offset.lean b/src/Init/Grind/Offset.lean
index 56626662488..8d35dff1e81 100644
--- a/src/Init/Grind/Offset.lean
+++ b/src/Init/Grind/Offset.lean
@@ -7,7 +7,7 @@ prelude
 import Init.Core
 import Init.Omega
 
-namespace Lean.Grind
+namespace Lean.Grind.Offset
 
 abbrev Var := Nat
 abbrev Context := Lean.RArray Nat
@@ -22,7 +22,7 @@ structure Cnstr where
   y : Var
   k : Nat  := 0
   l : Bool := true
-  deriving Repr, BEq, Inhabited
+  deriving Repr, DecidableEq, Inhabited
 
 def Cnstr.denote (c : Cnstr) (ctx : Context) : Prop :=
   if c.l then
@@ -82,51 +82,84 @@ def Cnstr.isFalse (c : Cnstr) : Bool := c.x == c.y && c.k != 0 && c.l == true
 theorem Cnstr.of_isFalse (ctx : Context) {c : Cnstr} : c.isFalse = true → ¬c.denote ctx := by
   cases c; simp [isFalse]; intros; simp [*, denote]; omega
 
-def Certificate := List Cnstr
+def Cnstrs := List Cnstr
 
-def Certificate.denote' (ctx : Context) (c₁ : Cnstr) (c₂ : Certificate) : Prop :=
+def Cnstrs.denoteAnd' (ctx : Context) (c₁ : Cnstr) (c₂ : Cnstrs) : Prop :=
   match c₂ with
   | [] => c₁.denote ctx
-  | c::cs => c₁.denote ctx ∧ Certificate.denote' ctx c cs
+  | c::cs => c₁.denote ctx ∧ Cnstrs.denoteAnd' ctx c cs
 
-theorem Certificate.denote'_trans (ctx : Context) (c₁ c : Cnstr) (cs : Certificate) : c₁.denote ctx → denote' ctx c cs → denote' ctx (c₁.trans c) cs := by
+theorem Cnstrs.denote'_trans (ctx : Context) (c₁ c : Cnstr) (cs : Cnstrs) : c₁.denote ctx → denoteAnd' ctx c cs → denoteAnd' ctx (c₁.trans c) cs := by
   induction cs
-  next => simp [denote', *]; apply Cnstr.denote_trans
-  next c cs ih => simp [denote']; intros; simp [*]
+  next => simp [denoteAnd', *]; apply Cnstr.denote_trans
+  next c cs ih => simp [denoteAnd']; intros; simp [*]
 
-def Certificate.trans' (c₁ : Cnstr) (c₂ : Certificate) : Cnstr :=
+def Cnstrs.trans' (c₁ : Cnstr) (c₂ : Cnstrs) : Cnstr :=
   match c₂ with
   | [] => c₁
   | c::c₂ => trans' (c₁.trans c) c₂
 
-@[simp] theorem Certificate.denote'_trans' (ctx : Context) (c₁ : Cnstr) (c₂ : Certificate) : denote' ctx c₁ c₂ → (trans' c₁ c₂).denote ctx := by
+@[simp] theorem Cnstrs.denote'_trans' (ctx : Context) (c₁ : Cnstr) (c₂ : Cnstrs) : denoteAnd' ctx c₁ c₂ → (trans' c₁ c₂).denote ctx := by
   induction c₂ generalizing c₁
-  next => intros; simp_all [trans', denote']
-  next c cs ih => simp [denote']; intros; simp [trans']; apply ih; apply denote'_trans <;> assumption
+  next => intros; simp_all [trans', denoteAnd']
+  next c cs ih => simp [denoteAnd']; intros; simp [trans']; apply ih; apply denote'_trans <;> assumption
 
-def Certificate.denote (ctx : Context) (c : Certificate) : Prop :=
+def Cnstrs.denoteAnd (ctx : Context) (c : Cnstrs) : Prop :=
   match c with
   | [] => True
-  | c::cs => denote' ctx c cs
+  | c::cs => denoteAnd' ctx c cs
 
-def Certificate.trans (c : Certificate) : Cnstr :=
+def Cnstrs.trans (c : Cnstrs) : Cnstr :=
   match c with
   | []    => trivialCnstr
   | c::cs => trans' c cs
 
-theorem Certificate.denote_trans {ctx : Context} {c : Certificate} : c.denote ctx → c.trans.denote ctx := by
-  cases c <;> simp [*, trans, Certificate.denote] <;> intros <;> simp [*]
+theorem Cnstrs.of_denoteAnd_trans {ctx : Context} {c : Cnstrs} : c.denoteAnd ctx → c.trans.denote ctx := by
+  cases c <;> simp [*, trans, denoteAnd] <;> intros <;> simp [*]
 
-def Certificate.isFalse (c : Certificate) : Bool :=
+def Cnstrs.isFalse (c : Cnstrs) : Bool :=
   c.trans.isFalse
 
-theorem Certificate.unsat (ctx : Context) (c : Certificate) : c.isFalse = true → ¬ c.denote ctx := by
+theorem Cnstrs.unsat' (ctx : Context) (c : Cnstrs) : c.isFalse = true → ¬ c.denoteAnd ctx := by
   simp [isFalse]; intro h₁ h₂
-  have := Certificate.denote_trans h₂
+  have := of_denoteAnd_trans h₂
   have := Cnstr.of_isFalse ctx h₁
   contradiction
 
-theorem Certificate.imp (ctx : Context) (c : Certificate) : c.denote ctx → c.trans.denote ctx := by
-  apply denote_trans
-
-end Lean.Grind
+/-- Returns `c_1 → ... → c_n → C -/
+def Cnstrs.denote (ctx : Context) (cs : Cnstrs) (C : Prop) : Prop :=
+  match cs with
+  | [] => C
+  | c::cs => c.denote ctx → denote ctx cs C
+
+theorem Cnstrs.not_denoteAnd'_eq (ctx : Context) (c : Cnstr) (cs : Cnstrs) (C : Prop) : (denoteAnd' ctx c cs → C) = denote ctx (c::cs) C := by
+  simp [denote]
+  induction cs generalizing c
+  next => simp [denoteAnd', denote]
+  next c' cs ih =>
+    simp [denoteAnd', denote, *]
+
+theorem Cnstrs.not_denoteAnd_eq (ctx : Context) (cs : Cnstrs) (C : Prop) : (denoteAnd ctx cs → C) = denote ctx cs C := by
+  cases cs
+  next => simp [denoteAnd, denote]
+  next c cs => apply not_denoteAnd'_eq
+
+def Cnstr.isImpliedBy (cs : Cnstrs) (c : Cnstr) : Bool :=
+  cs.trans == c
+
+/-! Main theorems used by `grind`. -/
+
+/-- Auxiliary theorem used by `grind` to prove that a system of offset inequalities is unsatisfiable. -/
+theorem Cnstrs.unsat (ctx : Context) (cs : Cnstrs) : cs.isFalse = true → cs.denote ctx False := by
+  intro h
+  rw [← not_denoteAnd_eq]
+  apply unsat'
+  assumption
+
+/-- Auxiliary theorem used by `grind` to prove an implied offset inequality. -/
+theorem Cnstrs.imp (ctx : Context) (cs : Cnstrs) (c : Cnstr) (h : c.isImpliedBy cs = true)  : cs.denote ctx (c.denote ctx) := by
+  rw [← eq_of_beq h]
+  rw [← not_denoteAnd_eq]
+  apply of_denoteAnd_trans
+
+end Lean.Grind.Offset
diff --git a/tests/lean/run/grind_offset_ineq_thms.lean b/tests/lean/run/grind_offset_ineq_thms.lean
new file mode 100644
index 00000000000..bb76dcfe205
--- /dev/null
+++ b/tests/lean/run/grind_offset_ineq_thms.lean
@@ -0,0 +1,31 @@
+import Lean
+
+elab tk:"#R[" ts:term,* "]" : term => do
+  let ts : Array Lean.Syntax := ts
+  let es ← ts.mapM fun stx => Lean.Elab.Term.elabTerm stx none
+  if h : 0 < es.size then
+    return (Lean.RArray.toExpr (← Lean.Meta.inferType es[0]!) id (Lean.RArray.ofArray es h))
+  else
+    throwErrorAt tk "RArray cannot be empty"
+
+open Lean.Grind.Offset
+
+macro "C[" "#" x:term:max " ≤ " "#" y:term:max "]" : term => `({ x := $x, y := $y : Cnstr })
+macro "C[" "#" x:term:max " + " k:term:max " ≤ " "#" y:term:max "]" : term => `({ x := $x, y := $y, k := $k : Cnstr })
+macro "C[" "#" x:term:max " ≤ " "#"y:term:max " + " k:term:max "]" : term => `({ x := $x, y := $y, k := $k, l := false : Cnstr })
+
+example (x y z : Nat) : x + 2 ≤ y → y ≤ z → z + 1 ≤ x → False :=
+  Cnstrs.unsat #R[x, y, z] [
+    C[ #0 + 2 ≤ #1 ],
+    C[ #1 ≤ #2 ],
+    C[ #2 + 1 ≤ #0 ]
+  ] rfl
+
+example (x y z w : Nat) : x + 2 ≤ y → y ≤ z → z ≤ w + 7 → x ≤ w + 5 :=
+  Cnstrs.imp #R[x, y, z, w] [
+    C[ #0 + 2 ≤ #1 ],
+    C[ #1 ≤ #2 ],
+    C[ #2 ≤ #3 + 7]
+  ]
+  C[ #0 ≤ #3 + 5 ]
+  rfl