From 6259b4742c70d9621543ca4c4a086cbc6b6d55e0 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Wed, 15 Jan 2025 16:17:27 -0800
Subject: [PATCH] feat: improve case-split heuristic used in `grind` (#6658)

This PR ensures that `grind` avoids case-splitting on terms congruent to
those that have already been case-split.
---
 src/Lean/Meta/Tactic/Grind/Split.lean | 11 ++++++
 src/Lean/Meta/Tactic/Grind/Types.lean |  5 +--
 tests/lean/run/grind_split_issue.lean | 50 +++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/grind_split_issue.lean

diff --git a/src/Lean/Meta/Tactic/Grind/Split.lean b/src/Lean/Meta/Tactic/Grind/Split.lean
index 15a232bdeb1e..f0a76defeb7d 100644
--- a/src/Lean/Meta/Tactic/Grind/Split.lean
+++ b/src/Lean/Meta/Tactic/Grind/Split.lean
@@ -74,6 +74,14 @@ private def checkIffStatus (e a b : Expr) : GoalM CaseSplitStatus := do
   else
     return .notReady
 
+/-- Returns `true` is `c` is congruent to a case-split that was already performed. -/
+private def isCongrToPrevSplit (c : Expr) : GoalM Bool := do
+  (← get).resolvedSplits.foldM (init := false) fun flag { expr := c' } => do
+    if flag then
+      return true
+    else
+      return isCongruent (← get).enodes c c'
+
 private def checkCaseSplitStatus (e : Expr) : GoalM CaseSplitStatus := do
   match_expr e with
   | Or a b => checkDisjunctStatus e a b
@@ -85,6 +93,8 @@ private def checkCaseSplitStatus (e : Expr) : GoalM CaseSplitStatus := do
     if (← isResolvedCaseSplit e) then
       trace[grind.debug.split] "split resolved: {e}"
       return .resolved
+    if (← isCongrToPrevSplit e) then
+      return .resolved
     if let some info := isMatcherAppCore? (← getEnv) e then
       return .ready info.numAlts
     let .const declName .. := e.getAppFn | unreachable!
@@ -163,6 +173,7 @@ def splitNext : GrindTactic := fun goal => do
       | return none
     let gen ← getGeneration c
     let genNew := if numCases > 1 || isRec then gen+1 else gen
+    markCaseSplitAsResolved c
     trace_goal[grind.split] "{c}, generation: {gen}"
     let mvarIds ← if (← isMatcherApp c) then
       casesMatch (← get).mvarId c
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 17a00cde81d6..a4f0390e6d64 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -378,7 +378,7 @@ structure Goal where
   splitCandidates : List Expr := []
   /-- Number of splits performed to get to this goal. -/
   numSplits : Nat := 0
-  /-- Case-splits that do not have to be performed anymore. -/
+  /-- Case-splits that have already been performed, or that do not have to be performed anymore. -/
   resolvedSplits : PHashSet ENodeKey := {}
   /-- Next local E-match theorem idx. -/
   nextThmIdx : Nat := 0
@@ -879,7 +879,8 @@ def isResolvedCaseSplit (e : Expr) : GoalM Bool :=
 
 /--
 Mark `e` as a case-split that does not need to be performed anymore.
-Remark: we currently use this feature to disable `match`-case-splits
+Remark: we currently use this feature to disable `match`-case-splits.
+Remark: we also use this feature to record the case-splits that have already been performed.
 -/
 def markCaseSplitAsResolved (e : Expr) : GoalM Unit := do
   unless (← isResolvedCaseSplit e) do
diff --git a/tests/lean/run/grind_split_issue.lean b/tests/lean/run/grind_split_issue.lean
new file mode 100644
index 000000000000..604ac434415b
--- /dev/null
+++ b/tests/lean/run/grind_split_issue.lean
@@ -0,0 +1,50 @@
+variable (d : Nat) in
+inductive X : Nat → Prop
+  | f {s : Nat} : X s
+  | g {s : Nat} : X d → X s
+
+/--
+error: `grind` failed
+case grind.1
+c : Nat
+q : X c 0
+s : Nat
+h✝ : 0 = s
+h : HEq ⋯ ⋯
+⊢ False
+[grind] Diagnostics
+  [facts] Asserted facts
+    [prop] X c 0
+    [prop] 0 = s
+    [prop] HEq ⋯ ⋯
+  [eqc] True propositions
+    [prop] X c 0
+    [prop] X c s
+  [eqc] Equivalence classes
+    [eqc] {s, 0}
+case grind.2
+c : Nat
+q : X c 0
+s : Nat
+a✝¹ a✝ : X c c
+h✝ : 0 = s
+h : HEq ⋯ ⋯
+⊢ False
+[grind] Diagnostics
+  [facts] Asserted facts
+    [prop] X c 0
+    [prop] X c c
+    [prop] 0 = s
+    [prop] HEq ⋯ ⋯
+  [eqc] True propositions
+    [prop] X c 0
+    [prop] X c c
+    [prop] X c s
+  [eqc] Equivalence classes
+    [eqc] {s, 0}
+  [issues] Issues
+    [issue] this goal was not fully processed due to previous failures, threshold: `(failures := 1)`
+-/
+#guard_msgs (error) in
+example {c : Nat} (q : X c 0) : False := by
+  grind