fix: proof generation for grind tactic

This PR fixes a bug in the proof assembly procedure utilized by the `grind` tactic.

src/Lean/Meta/Tactic/Grind/Inv.lean

@@ -80,7 +80,7 @@ private def checkProofs : GoalM Unit := do
     for a in eqc do
       for b in eqc do
         unless isSameExpr a b do
-          let p ← mkEqProof a b
+          let p ← mkEqHEqProof a b
           trace[grind.debug.proofs] "{a} = {b}"
           check p
           trace[grind.debug.proofs] "checked: {← inferType p}"

src/Lean/Meta/Tactic/Grind/Proof.lean

@@ -198,42 +198,41 @@ mutual
     -- `h' : lhs = target`
     mkTrans' h' h heq
 
+  /--
+  Returns a proof of `lhs = rhs` (`HEq lhs rhs`) if `heq = false` (`heq = true`).
+  If `heq = false`, this function assumes that `lhs` and `rhs` have the same type.
+  -/
   private partial def mkEqProofCore (lhs rhs : Expr) (heq : Bool) : GoalM Expr := do
     if isSameExpr lhs rhs then
       return (← mkRefl lhs heq)
+    -- The equivalence class contains `HEq` proofs. So, we build a proof using HEq. Otherwise, we use `Eq`.
+    let heqProofs := (← getRootENode lhs).heqProofs
     let n₁ ← getENode lhs
     let n₂ ← getENode rhs
     assert! isSameExpr n₁.root n₂.root
     let common ← findCommon lhs rhs
-    let lhsEqCommon? ← mkProofTo lhs common none heq
-    let some lhsEqRhs ← mkProofFrom rhs common lhsEqCommon? heq | unreachable!
-    return lhsEqRhs
+    let lhsEqCommon? ← mkProofTo lhs common none heqProofs
+    let some lhsEqRhs ← mkProofFrom rhs common lhsEqCommon? heqProofs | unreachable!
+    if heq == heqProofs then
+      return lhsEqRhs
+    else if heq then
+      mkHEqOfEq lhsEqRhs
+    else
+      mkEqOfHEq lhsEqRhs
+
 end
 
 /--
-Returns a proof that `a = b` (or `HEq a b`).
+Returns a proof that `a = b`.
 It assumes `a` and `b` are in the same equivalence class.
 -/
 @[export lean_grind_mk_eq_proof]
 def mkEqProofImpl (a b : Expr) : GoalM Expr := do
-  let p ← go
-  trace[grind.debug.proof] "{p}"
-  return p
-where
-  go : GoalM Expr := do
-    let n ← getRootENode a
-    if !n.heqProofs then
-      trace[grind.debug.proof] "{a} = {b}"
-      mkEqProofCore a b (heq := false)
-    else
-      if (← hasSameType a b) then
-        trace[grind.debug.proof] "{a} = {b}"
-        mkEqOfHEq (← mkEqProofCore a b (heq := true))
-      else
-        trace[grind.debug.proof] "{a} ≡ {b}"
-        mkEqProofCore a b (heq := true)
+  assert! (← hasSameType a b)
+  mkEqProofCore a b (heq := false)
 
-def mkHEqProof (a b : Expr) : GoalM Expr :=
+@[export lean_grind_mk_heq_proof]
+def mkHEqProofImpl (a b : Expr) : GoalM Expr :=
   mkEqProofCore a b (heq := true)
 
 end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -563,13 +563,31 @@ def isInconsistent : GoalM Bool :=
   return (← get).inconsistent
 
 /--
-Returns a proof that `a = b` (or `HEq a b`).
-It assumes `a` and `b` are in the same equivalence class.
+Returns a proof that `a = b`.
+It assumes `a` and `b` are in the same equivalence class, and have the same type.
 -/
 -- Forward definition
 @[extern "lean_grind_mk_eq_proof"]
 opaque mkEqProof (a b : Expr) : GoalM Expr
 
+/--
+Returns a proof that `HEq a b`.
+It assumes `a` and `b` are in the same equivalence class.
+-/
+-- Forward definition
+@[extern "lean_grind_mk_heq_proof"]
+opaque mkHEqProof (a b : Expr) : GoalM Expr
+
+/--
+Returns a proof that `a = b` if they have the same type. Otherwise, returns a proof of `HEq a b`.
+It assumes `a` and `b` are in the same equivalence class.
+-/
+def mkEqHEqProof (a b : Expr) : GoalM Expr := do
+  if (← hasSameType a b) then
+    mkEqProof a b
+  else
+    mkHEqProof a b
+
 /--
 Returns a proof that `a = True`.
 It assumes `a` and `True` are in the same equivalence class.

tests/lean/run/grind_heq_proof_issue.lean

@@ -0,0 +1,7 @@
+def f (a : α) := a
+
+example (a b : α) (x y : β) : HEq a x → x = y → HEq y b → f a = f b := by
+  grind
+
+example (a b : α) (x y : β) : x = y → HEq a x → HEq y b → f a = f b := by
+  grind