fix: bug at mkEqProof in grind

src/Lean/Meta/Tactic/Grind.lean

@@ -31,5 +31,6 @@ builtin_initialize registerTraceClass `grind.debug
 builtin_initialize registerTraceClass `grind.debug.proofs
 builtin_initialize registerTraceClass `grind.simp
 builtin_initialize registerTraceClass `grind.congr
+builtin_initialize registerTraceClass `grind.proof
 
 end Lean

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

@@ -142,7 +142,6 @@ where
       flipped
     }
     let parents ← removeParents lhsRoot.self
-    -- TODO: set propagateBool
     updateRoots lhs rhsNode.root
     trace[grind.debug] "{← ppENodeRef lhs} new root {← ppENodeRef rhsNode.root}, {← ppENodeRef (← getRoot lhs)}"
     reinsertParents parents
@@ -162,7 +161,6 @@ where
 
   updateRoots (lhs : Expr) (rootNew : Expr) : GoalM Unit := do
     let rec loop (e : Expr) : GoalM Unit := do
-      -- TODO: propagateBool
       let n ← getENode e
       setENode e { n with root := rootNew }
       unless (← isInconsistent) do

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

@@ -20,6 +20,9 @@ private def checkEqc (root : ENode) : GoalM Unit := do
     size := size + 1
     -- The root of `curr` must be `root`
     assert! isSameExpr (← getRoot curr) root.self
+    -- If the equivalence class does not have HEq proofs, then the types must be definitionally equal.
+    unless root.heqProofs do
+      assert! (← withDefault <| isDefEq (← inferType curr) (← inferType root.self))
     -- Starting at `curr`, following the `target?` field leads to `root`.
     let mut n := curr
     repeat

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

@@ -124,7 +124,8 @@ Returns a proof that `a = b` (or `HEq a b`).
 It assumes `a` and `b` are in the same equivalence class.
 -/
 def mkEqProof (a b : Expr) : GoalM Expr := do
-  let n ← getENode a
+  let n ← getRootENode a
+  trace[grind.proof] "{a} {if n.heqProofs then "≡" else "="} {b}"
   if !n.heqProofs then
     mkEqProofCore a b (heq := false)
   else if (← withDefault <| isDefEq (← inferType a) (← inferType b)) then

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

@@ -330,6 +330,10 @@ def getRoot? (e : Expr) : GoalM (Option Expr) := do
 def getRoot (e : Expr) : GoalM Expr :=
   return (← getENode e).root
 
+/-- Returns the root enode in the equivalence class of `e`. -/
+def getRootENode (e : Expr) : GoalM ENode := do
+  getENode (← getRoot e)
+
 /-- Returns the next element in the equivalence class of `e`. -/
 def getNext (e : Expr) : GoalM Expr :=
   return (← getENode e).next
@@ -350,7 +354,7 @@ def pushEqCore (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit :=
   modify fun s => { s with newEqs := s.newEqs.push { lhs, rhs, proof, isHEq } }
 
 @[inline] def pushEqHEq (lhs rhs proof : Expr) : GoalM Unit := do
-  if (← isDefEq (← inferType lhs) (← inferType rhs)) then
+  if (← withDefault <| isDefEq (← inferType lhs) (← inferType rhs)) then
     pushEqCore lhs rhs proof (isHEq := false)
   else
     pushEqCore lhs rhs proof (isHEq := true)

tests/lean/run/grind_nested_proofs.lean

@@ -13,8 +13,7 @@ elab "grind_test" : tactic => withMainContext do
     logInfo (← getEqc n.self)
 
 set_option grind.debug true
--- TODO: fix nested proof support
--- set_option grind.debug.proofs true
+set_option grind.debug.proofs true
 
 /-
 Recall that array access terms, such as `a[i]`, have nested proofs.