From 3ce7dd318d0ca86b2b499de4c8d764c3cadf9f33 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 5 Jun 2025 00:35:59 -0400
Subject: [PATCH] feat: sort equivalence classes in `grind` diagnostics (#8638)

This PR improves the diagnostic information produced by `grind`. It now
sorts the equivalence classes by generation and then `Expr. lt`.
---
 .../Meta/Tactic/Grind/Arith/Cutsat/Model.lean |  5 +-
 src/Lean/Meta/Tactic/Grind/PP.lean            |  4 +-
 src/Lean/Meta/Tactic/Grind/Types.lean         | 38 ++++---
 tests/lean/run/grind_congr.lean               | 18 +---
 tests/lean/run/grind_nested_proofs.lean       |  4 +-
 tests/lean/run/grind_pre.lean                 | 16 +--
 .../lean/run/grind_propagate_connectives.lean |  2 +-
 tests/lean/run/grind_sort_eqc.lean            | 99 +++++++++++++++++++
 tests/lean/run/grind_split_issue.lean         |  2 +-
 tests/lean/run/grind_t1.lean                  | 10 +-
 10 files changed, 154 insertions(+), 44 deletions(-)
 create mode 100644 tests/lean/run/grind_sort_eqc.lean

diff --git a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Model.lean b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Model.lean
index d319b6685f..763fc570f1 100644
--- a/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Model.lean
+++ b/src/Lean/Meta/Tactic/Grind/Arith/Cutsat/Model.lean
@@ -121,7 +121,10 @@ def mkModel (goal : Goal) : MetaM (Array (Expr × Rat)) := do
   for (e, v) in model do
     unless isInterpretedTerm e do
       r := r.push (e, v)
-  r := r.qsort fun (e₁, _) (e₂, _) => e₁.lt e₂
+  r := r.qsort fun (e₁, _) (e₂, _) =>
+    let g₁ := goal.getGeneration e₁
+    let g₂ := goal.getGeneration e₂
+    if g₁ != g₂ then g₁ < g₂ else e₁.lt e₂
   if (← isTracingEnabledFor `grind.cutsat.model) then
     for (x, v) in r do
       trace[grind.cutsat.model] "{quoteIfArithTerm x} := {v}"
diff --git a/src/Lean/Meta/Tactic/Grind/PP.lean b/src/Lean/Meta/Tactic/Grind/PP.lean
index 4628c39a46..0bfd5dca0c 100644
--- a/src/Lean/Meta/Tactic/Grind/PP.lean
+++ b/src/Lean/Meta/Tactic/Grind/PP.lean
@@ -60,7 +60,7 @@ def Goal.ppState (goal : Goal) : MetaM MessageData := do
   for e in goal.exprs do
     let node ← goal.getENode e
     r := r ++ "\n" ++ (← goal.ppENodeDecl node.self)
-  let eqcs := goal.getEqcs
+  let eqcs := goal.getEqcs (sort := true)
   for eqc in eqcs do
     if eqc.length > 1 then
       r := r ++ "\n" ++ "{" ++ (MessageData.joinSep (← eqc.mapM goal.ppENodeRef) ", ") ++  "}"
@@ -87,7 +87,7 @@ private def ppEqcs : M Unit := do
    let mut falseEqc? : Option MessageData := none
    let mut otherEqcs : Array MessageData := #[]
    let goal ← read
-   for eqc in goal.getEqcs do
+   for eqc in goal.getEqcs (sort := true) do
      if Option.isSome <| eqc.find? (·.isTrue) then
        let eqc := eqc.filter fun e => !e.isTrue
        unless eqc.isEmpty do
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 646be3e408..623856d25f 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -1359,30 +1359,46 @@ def applyFallback : GoalM Unit := do
   let fallback : GoalM Unit := (← getMethods).fallback
   fallback
 
+def Goal.getGeneration (goal : Goal) (e : Expr) : Nat :=
+  if let some n := goal.getENode? e then
+    n.generation
+  else
+    0
+
 /-- Returns expressions in the given expression equivalence class. -/
-partial def Goal.getEqc (goal : Goal) (e : Expr) : List Expr :=
-  go e e []
+partial def Goal.getEqc (goal : Goal) (e : Expr) (sort := false) : List Expr :=
+  let eqc := go e e #[]
+  if sort then
+    Array.toList <| eqc.qsort fun e₁ e₂ =>
+      let g₁ := goal.getGeneration e₁
+      let g₂ := goal.getGeneration e₂
+      if g₁ != g₂ then g₁ < g₂ else e₁.lt e₂
+  else
+    eqc.toList
 where
-  go (first : Expr) (e : Expr) (acc : List Expr) : List Expr := Id.run do
-    let some next ← goal.getNext? e | acc
-    let acc := e :: acc
+  go (first : Expr) (e : Expr) (acc : Array Expr) : Array Expr := Id.run do
+    let some next := goal.getNext? e | acc
+    let acc := acc.push e
     if isSameExpr first next then
       return acc
     else
       go first next acc
 
 @[inline, inherit_doc Goal.getEqc]
-partial def getEqc (e : Expr) : GoalM (List Expr) :=
-  return (← get).getEqc e
+partial def getEqc (e : Expr) (sort := false) : GoalM (List Expr) :=
+  return (← get).getEqc e sort
 
 /-- Returns all equivalence classes in the current goal. -/
-partial def Goal.getEqcs (goal : Goal) : List (List Expr) := Id.run do
- let mut r : List (List Expr) := []
+partial def Goal.getEqcs (goal : Goal) (sort := false) : List (List Expr) := Id.run do
+ let mut r : Array (Nat × Expr × List Expr) := #[]
  for e in goal.exprs do
     let some node := goal.getENode? e | pure ()
     if node.isRoot then
-      r := goal.getEqc node.self :: r
-  return r
+      let e :: es := goal.getEqc node.self sort | unreachable!
+      r := r.push (goal.getGeneration e, e, es)
+  if sort then
+    r := r.qsort fun (g₁, e₁, _) (g₂, e₂, _) => if g₁ != g₂ then g₁ < g₂ else e₁.lt e₂
+  return r.toList.map fun (_, e, es) => e::es
 
 @[inline, inherit_doc Goal.getEqcs]
 def getEqcs : GoalM (List (List Expr)) :=
diff --git a/tests/lean/run/grind_congr.lean b/tests/lean/run/grind_congr.lean
index 72e30aa9fd..4eee390d20 100644
--- a/tests/lean/run/grind_congr.lean
+++ b/tests/lean/run/grind_congr.lean
@@ -8,7 +8,7 @@ def g (a : Nat) := a + a
 open Lean Meta Grind in
 def fallback : Fallback := do
   let #[n, _] ← filterENodes fun e => return e.self.isApp && e.self.isAppOf ``f | unreachable!
-  let eqc ← getEqc n.self
+  let eqc ← getEqc n.self (sort := true)
   trace[Meta.debug] eqc
   (← get).mvarId.admit
 
@@ -16,30 +16,22 @@ set_option trace.Meta.debug true
 set_option grind.debug true
 set_option grind.debug.proofs true
 
-/--
-trace: [Meta.debug] [d, f b, c, f a]
--/
+/-- trace: [Meta.debug] [c, d, f a, f b] -/
 #guard_msgs (trace) in
 example (a b c d : Nat) : a = b → f a = c → f b = d → False := by
   grind on_failure fallback
 
-/--
-trace: [Meta.debug] [d, f b, c, f a]
--/
+/-- trace: [Meta.debug] [c, d, f a, f b] -/
 #guard_msgs (trace) in
 example (a b c d : Nat) : f a = c → f b = d → a = b → False := by
   grind on_failure fallback
 
-/--
-trace: [Meta.debug] [d, f (g b), c, f (g a)]
--/
+/-- trace: [Meta.debug] [c, d, f (g a), f (g b)] -/
 #guard_msgs (trace) in
 example (a b c d e : Nat) : f (g a) = c → f (g b) = d → a = e → b = e → False := by
   grind on_failure fallback
 
-/--
-trace: [Meta.debug] [d, f (g b), c, f v]
--/
+/-- trace: [Meta.debug] [c, d, f v, f (g b)] -/
 #guard_msgs (trace) in
 example (a b c d e v : Nat) : f v = c → f (g b) = d → a = e → b = e → v = g a → False := by
   grind on_failure fallback
diff --git a/tests/lean/run/grind_nested_proofs.lean b/tests/lean/run/grind_nested_proofs.lean
index d32873ac65..b1e55d3256 100644
--- a/tests/lean/run/grind_nested_proofs.lean
+++ b/tests/lean/run/grind_nested_proofs.lean
@@ -23,7 +23,7 @@ detect equalities between array access terms.
 
 /--
 trace: [Meta.debug] [‹i < a.size›, ‹j < a.size›, ‹j < b.size›]
-[Meta.debug] [a[i], b[j], a[j]]
+[Meta.debug] [a[j], b[j], a[i]]
 -/
 #guard_msgs (trace) in
 example (i j : Nat) (a b : Array Nat) (h1 : j < a.size) (h : j < b.size) (h2 : i ≤ j) : a[i] < a[j] + b[j] → i = j → a = b → False := by
@@ -31,7 +31,7 @@ example (i j : Nat) (a b : Array Nat) (h1 : j < a.size) (h : j < b.size) (h2 : i
 
 /--
 trace: [Meta.debug] [‹i < a.size›, ‹j < a.size›, ‹j < b.size›]
-[Meta.debug] [a[i], a[j]]
+[Meta.debug] [a[j], a[i]]
 -/
 #guard_msgs (trace) in
 example (i j : Nat) (a b : Array Nat) (h1 : j < a.size) (h : j < b.size) (h2 : i ≤ j) : a[i] < a[j] + b[j] → i = j → False := by
diff --git a/tests/lean/run/grind_pre.lean b/tests/lean/run/grind_pre.lean
index 4452de7c81..ecde73cc1b 100644
--- a/tests/lean/run/grind_pre.lean
+++ b/tests/lean/run/grind_pre.lean
@@ -22,18 +22,18 @@ h_1 : b = false ∨ a = false
     [prop] q
     [prop] b = false ∨ a = false
   [eqc] True propositions
-    [prop] b = true ∨ c = true
     [prop] p
     [prop] q
     [prop] b = false ∨ a = false
+    [prop] b = true ∨ c = true
     [prop] b = false
     [prop] c = true
   [eqc] False propositions
     [prop] a = false
     [prop] b = true
   [eqc] Equivalence classes
-    [eqc] {b, false}
     [eqc] {a, c, true}
+    [eqc] {b, false}
 -/
 #guard_msgs (error) in
 theorem ex (h : (f a && (b || f (f c))) = true) (h' : p ∧ q) : b && a := by
@@ -64,8 +64,8 @@ h_3 : b = false
     [prop] p
     [prop] q
   [eqc] Equivalence classes
-    [eqc] {b, false}
     [eqc] {a, c, true}
+    [eqc] {b, false}
 [grind] Diagnostics
   [cases] Cases instances
     [cases] Or ↦ 3
@@ -119,10 +119,10 @@ tail_eq_1 : as = bs
     [prop] a₃ = b₃
     [prop] as = bs
   [eqc] Equivalence classes
-    [eqc] {as, bs}
-    [eqc] {a₃, b₃}
-    [eqc] {a₂, b₂}
     [eqc] {a₁, b₁}
+    [eqc] {a₂, b₂}
+    [eqc] {a₃, b₃}
+    [eqc] {as, bs}
 -/
 #guard_msgs (error) in
 theorem ex3 (h : a₁ :: { x := a₂, y := a₃ : Point } :: as = b₁ :: { x := b₂, y := b₃} :: bs) : False := by
@@ -154,11 +154,11 @@ right : r
     [prop] p
     [prop] r
   [eqc] True propositions
-    [prop] p = r
     [prop] a
     [prop] p
     [prop] q
     [prop] r
+    [prop] p = r
   [cases] Case analyses
     [cases] [1/2]: p = r
       [cases] source: Initial goal
@@ -206,8 +206,8 @@ h_2 : ¬f a = g b
   [eqc] False propositions
     [prop] f a = g b
   [eqc] Equivalence classes
-    [eqc] {a, b}
     [eqc] {f, g}
+    [eqc] {a, b}
 [grind] Issues
   [issue] found congruence between
         g b
diff --git a/tests/lean/run/grind_propagate_connectives.lean b/tests/lean/run/grind_propagate_connectives.lean
index 87ceeaa5fd..4949410adf 100644
--- a/tests/lean/run/grind_propagate_connectives.lean
+++ b/tests/lean/run/grind_propagate_connectives.lean
@@ -70,7 +70,7 @@ example (p q r : Prop) : p ∨ ¬(s ∨ (p ↔ r)) → p = False → False := by
 /--
 trace: [Meta.debug] true:  [p]
 [Meta.debug] false: []
-[Meta.debug] [a, b]
+[Meta.debug] [b, a]
 -/
 #guard_msgs (trace) in
 example (p : Prop) (a : Vector Nat 5) (b : Vector Nat 6) : (p → HEq a b) → p → False := by
diff --git a/tests/lean/run/grind_sort_eqc.lean b/tests/lean/run/grind_sort_eqc.lean
new file mode 100644
index 0000000000..39e784dc8f
--- /dev/null
+++ b/tests/lean/run/grind_sort_eqc.lean
@@ -0,0 +1,99 @@
+set_option grind.warning false
+opaque f [Inhabited α] : α → α
+theorem feq [Inhabited α] [Add α] [One α] (x : α) : f x = f (f x) + 1 := sorry
+opaque g [Inhabited α] : α → α → α
+theorem geq [Inhabited α] (x : α) : g x x = x := sorry
+
+/--
+error: `grind` failed
+case grind
+α : Type u_1
+inst : Inhabited α
+inst_1 : Add α
+inst_2 : One α
+x z y : α
+h : g y z = z
+h_1 : g z y = g y x
+h_2 : ¬f x = g x x
+⊢ False
+[grind] Goal diagnostics
+  [facts] Asserted facts
+    [prop] g y z = z
+    [prop] g z y = g y x
+    [prop] ¬f x = g x x
+    [prop] f x = f (f x) + 1
+    [prop] g x x = x
+    [prop] f (f x) = f (f (f x)) + 1
+    [prop] f (f (f x)) = f (f (f (f x))) + 1
+  [eqc] False propositions
+    [prop] f x = g x x
+  [eqc] Equivalence classes
+    [eqc] {x, g x x}
+    [eqc] {z, g y z}
+    [eqc] {f x, f (f x) + 1}
+    [eqc] {g z y, g y x}
+    [eqc] {f (f x), f (f (f x)) + 1}
+    [eqc] {f (f (f x)), f (f (f (f x))) + 1}
+  [ematch] E-matching patterns
+    [thm] feq: [@f #4 #3 #0]
+    [thm] geq: [@g #2 #1 #0 #0]
+  [limits] Thresholds reached
+    [limit] maximum term generation has been reached, threshold: `(gen := 3)`
+[grind] Diagnostics
+  [thm] E-Matching instances
+    [thm] feq ↦ 3
+    [thm] geq ↦ 1
+-/
+#guard_msgs (error) in
+example [Inhabited α] [Add α] [One α] (x z y : α) : g y z = z → g z y = g y x → f x = g x x := by
+  grind (gen := 3) [= feq, = geq]
+
+/--
+error: `grind` failed
+case grind
+x z y : Int
+h : g y z = z
+h_1 : g z y = g y x
+h_2 : ¬f (f x) = g x x
+⊢ False
+[grind] Goal diagnostics
+  [facts] Asserted facts
+    [prop] g y z = z
+    [prop] g z y = g y x
+    [prop] ¬f (f x) = g x x
+    [prop] f x + -1 * f (f x) + -1 = 0
+    [prop] f (f x) + -1 * f (f (f x)) + -1 = 0
+    [prop] g x x = x
+    [prop] f (f (f x)) + -1 * f (f (f (f x))) + -1 = 0
+  [eqc] False propositions
+    [prop] f (f x) = g x x
+  [eqc] Equivalence classes
+    [eqc] {x, g x x}
+    [eqc] {z, g y z}
+    [eqc] {g z y, g y x}
+    [eqc] {0, f x + -1 * f (f x) + -1, f (f x) + -1 * f (f (f x)) + -1, f (f (f x)) + -1 * f (f (f (f x))) + -1}
+  [ematch] E-matching patterns
+    [thm] feq: [@f #4 #3 #0]
+    [thm] geq: [@g #2 #1 #0 #0]
+  [cutsat] Assignment satisfying linear constraints
+    [assign] x := 5
+    [assign] z := 3
+    [assign] y := 2
+    [assign] f x := 1
+    [assign] f (f x) := 0
+    [assign] g x x := 5
+    [assign] g z y := 4
+    [assign] g y x := 4
+    [assign] g y z := 3
+    [assign] f (f (f x)) := -1
+    [assign] f (f (f (f x))) := -2
+  [limits] Thresholds reached
+    [limit] maximum term generation has been reached, threshold: `(gen := 2)`
+[grind] Diagnostics
+  [thm] E-Matching instances
+    [thm] feq ↦ 3
+    [thm] geq ↦ 1
+-/
+#guard_msgs (error) in
+example (x z y : Int) : g y z = z → g z y = g y x → f (f x) = g x x := by
+  grind (gen := 2) -ring [= feq, = geq]
diff --git a/tests/lean/run/grind_split_issue.lean b/tests/lean/run/grind_split_issue.lean
index ae16f47b4b..aa765bcdaf 100644
--- a/tests/lean/run/grind_split_issue.lean
+++ b/tests/lean/run/grind_split_issue.lean
@@ -22,7 +22,7 @@ h_1 : ⋯ ≍ ⋯
     [prop] X c 0
     [prop] X c s
   [eqc] Equivalence classes
-    [eqc] {s, 0}
+    [eqc] {0, s}
   [cases] Case analyses
     [cases] [1/2]: X c 0
       [cases] source: Initial goal
diff --git a/tests/lean/run/grind_t1.lean b/tests/lean/run/grind_t1.lean
index 7c72682675..c475eb79da 100644
--- a/tests/lean/run/grind_t1.lean
+++ b/tests/lean/run/grind_t1.lean
@@ -260,9 +260,9 @@ h_1 : p
     [prop] p = q
     [prop] p
   [eqc] True propositions
-    [prop] p = q
-    [prop] q
     [prop] p
+    [prop] q
+    [prop] p = q
 -/
 #guard_msgs (error) in
 set_option trace.grind.split true in
@@ -294,9 +294,9 @@ h_1 : p
     [prop] p = ¬q
     [prop] p
   [eqc] True propositions
-    [prop] p = ¬q
-    [prop] ¬q
     [prop] p
+    [prop] ¬q
+    [prop] p = ¬q
   [eqc] False propositions
     [prop] q
 -/
@@ -378,7 +378,7 @@ h_1 : b = true
   [eqc] True propositions
     [prop] b = true
   [eqc] Equivalence classes
-    [eqc] {a, if b = true then 10 else 20, 10}
+    [eqc] {a, 10, if b = true then 10 else 20}
     [eqc] {b, true}
   [cutsat] Assignment satisfying linear constraints
     [assign] a := 10