From 4316629119b710c0518e04d4f154db3db7df0d3a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 29 May 2025 19:47:40 -0400
Subject: [PATCH] fix: `BEq` support in `grind` (#8536)

This PR fixes the support for `LawfulBEq` and `BEq` in `grind`.
---
 src/Init/Data/BEq.lean                    |  2 +-
 src/Init/Grind/Lemmas.lean                |  6 ++++
 src/Lean/Meta/Tactic/Grind/Propagate.lean | 41 +++++++++++++++++++----
 tests/lean/grind/list_count.lean          | 32 ------------------
 tests/lean/run/grind_bool_prop.lean       |  5 +++
 tests/lean/run/grind_fastEraseDups.lean   |  6 +++-
 tests/lean/run/grind_list_count.lean      | 22 ++++++++++++
 7 files changed, 74 insertions(+), 40 deletions(-)
 delete mode 100644 tests/lean/grind/list_count.lean

diff --git a/src/Init/Data/BEq.lean b/src/Init/Data/BEq.lean
index 2500efaed4..2cb394352a 100644
--- a/src/Init/Data/BEq.lean
+++ b/src/Init/Data/BEq.lean
@@ -27,7 +27,7 @@ class EquivBEq (α) [BEq α] : Prop extends PartialEquivBEq α, ReflBEq α
 theorem BEq.symm [BEq α] [PartialEquivBEq α] {a b : α} : a == b → b == a :=
   PartialEquivBEq.symm
 
-@[grind] theorem BEq.comm [BEq α] [PartialEquivBEq α] {a b : α} : (a == b) = (b == a) :=
+theorem BEq.comm [BEq α] [PartialEquivBEq α] {a b : α} : (a == b) = (b == a) :=
   Bool.eq_iff_iff.2 ⟨BEq.symm, BEq.symm⟩
 
 theorem bne_comm [BEq α] [PartialEquivBEq α] {a b : α} : (a != b) = (b != a) := by
diff --git a/src/Init/Grind/Lemmas.lean b/src/Init/Grind/Lemmas.lean
index 31ce7a33a3..930e02bda3 100644
--- a/src/Init/Grind/Lemmas.lean
+++ b/src/Init/Grind/Lemmas.lean
@@ -89,6 +89,12 @@ theorem beq_eq_true_of_eq {α : Type u} {_ : BEq α} {_ : LawfulBEq α} {a b : 
 theorem beq_eq_false_of_diseq {α : Type u} {_ : BEq α} {_ : LawfulBEq α} {a b : α} (h : ¬ a = b) : (a == b) = false := by
   simp[*]
 
+theorem eq_of_beq_eq_true {α : Type u} {_ : BEq α} {_ : LawfulBEq α} {a b : α} (h : (a == b) = true) : a = b := by
+  simp [beq_iff_eq.mp h]
+
+theorem ne_of_beq_eq_false {α : Type u} {_ : BEq α} {_ : LawfulBEq α} {a b : α} (h : (a == b) = false) : (a = b) = False := by
+  simp [beq_eq_false_iff_ne.mp h]
+
 /-! Bool.and -/
 
 theorem Bool.and_eq_of_eq_true_left {a b : Bool} (h : a = true) : (a && b) = b := by simp [h]
diff --git a/src/Lean/Meta/Tactic/Grind/Propagate.lean b/src/Lean/Meta/Tactic/Grind/Propagate.lean
index 636b2242e2..855fbe3cd0 100644
--- a/src/Lean/Meta/Tactic/Grind/Propagate.lean
+++ b/src/Lean/Meta/Tactic/Grind/Propagate.lean
@@ -173,6 +173,25 @@ builtin_grind_propagator propagateEqDown ↓Eq := fun e => do
       for thm in (← getExtTheorems α) do
         instantiateExtTheorem thm e
 
+private def getLawfulBEqInst? (u : List Level) (α : Expr) (binst : Expr) : MetaM (Option Expr) := do
+  let lawfulBEq := mkApp2 (mkConst ``LawfulBEq u) α binst
+  let .some linst ← trySynthInstance lawfulBEq | return none
+  return some linst
+
+/-
+Note about `BEq.beq`
+Given `a b : α` in a context where we have `[BEq α] [LawfulBEq α]`
+The normalizer (aka `simp`) fails to normalize `if a == b then ... else ...` to `if a = b then ... else ...` using
+```
+theorem beq_iff_eq [BEq α] [LawfulBEq α] {a b : α} : a == b ↔ a = b :=
+  ⟨eq_of_beq, beq_of_eq⟩
+```
+The main issue is that `ite_congr` requires that the resulting proposition to be decidable,
+and we don't have `[DecidableEq α]`. Thus, the normalization step fails.
+The following propagators for `BEq.beq` ensure `grind` does not assume this normalization
+rule has been applied.
+-/
+
 builtin_grind_propagator propagateBEqUp ↑BEq.beq := fun e => do
   /-
   `grind` uses the normalization rule `Bool.beq_eq_decide_eq`, but it is only applicable if
@@ -181,17 +200,27 @@ builtin_grind_propagator propagateBEqUp ↑BEq.beq := fun e => do
   Thus, we have added this propagator as a backup.
   -/
   let_expr f@BEq.beq α binst a b := e | return ()
+  let u := f.constLevels!
   if (← isEqv a b) then
-    let u := f.constLevels!
-    let lawfulBEq := mkApp2 (mkConst ``LawfulBEq u) α binst
-    let .some linst ← trySynthInstance lawfulBEq | return ()
+    let some linst ← getLawfulBEqInst? u α binst | return ()
     pushEqBoolTrue e <| mkApp6 (mkConst ``Grind.beq_eq_true_of_eq u) α binst linst a b (← mkEqProof a b)
   else if let some h ← mkDiseqProof? a b then
-    let u := f.constLevels!
-    let lawfulBEq := mkApp2 (mkConst ``LawfulBEq u) α binst
-    let .some linst ← trySynthInstance lawfulBEq | return ()
+    let some linst ← getLawfulBEqInst? u α binst | return ()
     pushEqBoolFalse e <| mkApp6 (mkConst ``Grind.beq_eq_false_of_diseq u) α binst linst a b h
 
+builtin_grind_propagator propagateBEqDown ↓BEq.beq := fun e => do
+  /- See comment above -/
+  let_expr f@BEq.beq α binst a b := e | return ()
+  let u := f.constLevels!
+  if (← isEqBoolTrue e) then
+    let some linst ← getLawfulBEqInst? u α binst | return ()
+    pushEq a b <| mkApp6 (mkConst ``Grind.eq_of_beq_eq_true u) α binst linst a b (← mkEqProof e (← getBoolTrueExpr))
+  else if (← isEqBoolFalse e) then
+    let some linst ← getLawfulBEqInst? u α binst | return ()
+    let eq ← shareCommon (mkApp3 (mkConst ``Eq [u.head!.succ]) α a b)
+    internalize eq (← getGeneration a)
+    pushEqFalse eq <| mkApp6 (mkConst ``Grind.ne_of_beq_eq_false u) α binst linst a b (← mkEqProof e (← getBoolFalseExpr))
+
 /-- Propagates `EqMatch` downwards -/
 builtin_grind_propagator propagateEqMatchDown ↓Grind.EqMatch := fun e => do
   if (← isEqTrue e) then
diff --git a/tests/lean/grind/list_count.lean b/tests/lean/grind/list_count.lean
deleted file mode 100644
index a94b2107c3..0000000000
--- a/tests/lean/grind/list_count.lean
+++ /dev/null
@@ -1,32 +0,0 @@
-open List
-
-set_option grind.warning false
-
-variable [BEq α] [LawfulBEq α]
-
--- These tests should move back to `tests/lean/run/grind_list_count.lean` once fixed.
-
-theorem count_pos_iff {a : α} {l : List α} : 0 < count a l ↔ a ∈ l := by
-  induction l with grind -- fails, having proved `head = a` is false and `head == a` is true.
-
-theorem one_le_count_iff {a : α} {l : List α} : 1 ≤ count a l ↔ a ∈ l := by
-  induction l with grind -- fails, similarly
-
-theorem count_eq_zero_of_not_mem {a : α} {l : List α} (h : a ∉ l) : count a l = 0 := by
-  induction l with grind -- fails
-
-theorem count_eq_zero {l : List α} : count a l = 0 ↔ a ∉ l := by
-  induction l with grind -- fails
-
-theorem count_filter {l : List α} (h : p a) : count a (filter p l) = count a l := by
-  induction l with grind -- similarly
-
-theorem count_le_count_map {β} [BEq β] [LawfulBEq β] {l : List α} {f : α → β} {x : α} :
-    count x l ≤ count (f x) (map f l) := by
-  induction l with grind
-
-theorem count_erase {a b : α} {l : List α} : count a (l.erase b) = count a l - if b == a then 1 else 0 := by
-  induction l with grind [-List.count_erase]
-  -- fails with inconsistent equivalence clases:
-  -- [] {head == a, false}
-  -- [] {b == a, head == b, true}
diff --git a/tests/lean/run/grind_bool_prop.lean b/tests/lean/run/grind_bool_prop.lean
index 95ffef4820..b4c2e977e9 100644
--- a/tests/lean/run/grind_bool_prop.lean
+++ b/tests/lean/run/grind_bool_prop.lean
@@ -1,3 +1,4 @@
+set_option grind.warning false
 example (f : Bool → Nat) : f (a && b) = 0 → a = false → f false = 0 := by grind (splits := 0)
 example (f : Bool → Nat) : f (a && b) = 0 → b = false → f false = 0 := by grind (splits := 0)
 example (f : Bool → Nat) : f (a && b) = 0 → a = true → f b = 0 := by grind (splits := 0)
@@ -31,3 +32,7 @@ example (a b : Bool) : (a ^^ b, c) = d → d = (false, true) → a = b := by gri
 example (a b : Bool) : (a == b, c) = d → d = (true, true) → a = true → true = b := by grind (splits := 0)
 
 example (h : α = β) (a : α) (b : β) : h ▸ a = b → HEq a b := by grind
+
+example {α : Type u} [BEq α] [LawfulBEq α] (x : Nat) (a b : α)
+    : x = (if a == b then 2 else 1) → x = (if (b == a) then 1 else 2) → False := by
+  grind
diff --git a/tests/lean/run/grind_fastEraseDups.lean b/tests/lean/run/grind_fastEraseDups.lean
index 4b05a5b135..c73d2f504f 100644
--- a/tests/lean/run/grind_fastEraseDups.lean
+++ b/tests/lean/run/grind_fastEraseDups.lean
@@ -24,6 +24,10 @@ where
       eraseDupsBy.loop (· == ·) l seenl = fastEraseDups.go l seenl seen := by
     induction l generalizing seenl seen with
     | nil => grind [eraseDupsBy.loop, fastEraseDups.go]
-    | cons x => cases h : seenl.contains x <;> grind [eraseDupsBy.loop, fastEraseDups.go]
+    | cons x =>
+      -- In the following example `BEq` is not lawful. To complete the proof we need to add `BEq.comm`
+      -- TODO: add support for arbitrary partial equivalence and equivalence relations.
+      -- Remark: `BEq.comm` is noise when `BEq` is lawful.
+      cases h : seenl.contains x <;> grind [eraseDupsBy.loop, fastEraseDups.go, BEq.comm]
 
 end List
diff --git a/tests/lean/run/grind_list_count.lean b/tests/lean/run/grind_list_count.lean
index 556a1d3d9f..d53ec9ff94 100644
--- a/tests/lean/run/grind_list_count.lean
+++ b/tests/lean/run/grind_list_count.lean
@@ -186,4 +186,26 @@ theorem count_erase_self {a : α} {l : List α} :
 theorem count_erase_of_ne (ab : a ≠ b) {l : List α} : count a (l.erase b) = count a l := by
   grind
 
+theorem count_pos_iff {a : α} {l : List α} : 0 < count a l ↔ a ∈ l := by
+  induction l with grind
+
+theorem one_le_count_iff {a : α} {l : List α} : 1 ≤ count a l ↔ a ∈ l := by
+  induction l with grind
+
+theorem count_eq_zero_of_not_mem {a : α} {l : List α} (h : a ∉ l) : count a l = 0 := by
+  induction l with grind
+
+theorem count_eq_zero {l : List α} : count a l = 0 ↔ a ∉ l := by
+  induction l with grind
+
+theorem count_filter {l : List α} (h : p a) : count a (filter p l) = count a l := by
+  induction l with grind
+
+theorem count_le_count_map {β} [BEq β] [LawfulBEq β] {l : List α} {f : α → β} {x : α} :
+    count x l ≤ count (f x) (map f l) := by
+  induction l with grind
+
+theorem count_erase {a b : α} {l : List α} : count a (l.erase b) = count a l - if b == a then 1 else 0 := by
+  induction l <;> grind [-List.count_erase]
+
 end count