fix: congruence closure used in the grind tactic

This PR fixes a bug in the congruence closure data structure used in
the `grind` tactic. The new test includes an example that previously caused a
panic. A similar panic was also occurring in the test `grind_nested_proofs.lean`.

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

@@ -41,8 +41,9 @@ This is an auxiliary function performed while merging equivalence classes.
 private def removeParents (root : Expr) : GoalM ParentSet := do
   let parents ← getParentsAndReset root
   for parent in parents do
-    trace[grind.debug.parent] "remove: {parent}"
-    modify fun s => { s with congrTable := s.congrTable.erase { e := parent } }
+    if (← isCongrRoot parent) then
+      trace[grind.debug.parent] "remove: {parent}"
+      modify fun s => { s with congrTable := s.congrTable.erase { e := parent } }
   return parents
 
 /--
@@ -51,8 +52,9 @@ This is an auxiliary function performed while merging equivalence classes.
 -/
 private def reinsertParents (parents : ParentSet) : GoalM Unit := do
   for parent in parents do
-    trace[grind.debug.parent] "reinsert: {parent}"
-    addCongrTable parent
+    if (← isCongrRoot parent) then
+      trace[grind.debug.parent] "reinsert: {parent}"
+      addCongrTable parent
 
 /-- Closes the goal when `True` and `False` are in the same equivalence class. -/
 private def closeGoalWithTrueEqFalse : GoalM Unit := do

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

@@ -122,7 +122,7 @@ private partial def processMatch (c : Choice) (p : Expr) (e : Expr) : M Unit :=
     let n ← getENode curr
     if n.generation <= maxGeneration
        -- uses heterogeneous equality or is the root of its congruence class
-       && (n.heqProofs || isSameExpr curr n.cgRoot)
+       && (n.heqProofs || n.isCongrRoot)
        && eqvFunctions pFn curr.getAppFn
        && curr.getAppNumArgs == numArgs then
       if let some c ← matchArgs? c p curr |>.run then
@@ -140,7 +140,7 @@ private def processContinue (c : Choice) (p : Expr) : M Unit := do
   for app in apps do
     let n ← getENode app
     if n.generation <= maxGeneration
-       && (n.heqProofs || isSameExpr n.cgRoot app) then
+       && (n.heqProofs || n.isCongrRoot) then
       if let some c ← matchArgs? c p app |>.run then
         let gen := n.generation
         let c := { c with gen := Nat.max gen c.gen }
@@ -225,7 +225,7 @@ private def main (p : Expr) (cnstrs : List Cnstr) : M Unit := do
   for app in apps do
     if (← checkMaxInstancesExceeded) then return ()
     let n ← getENode app
-    if (n.heqProofs || isSameExpr n.cgRoot app) &&
+    if (n.heqProofs || n.isCongrRoot) &&
        (!useMT || n.mt == gmt) then
       if let some c ← matchArgs? { cnstrs, assignment, gen := n.generation } p app |>.run then
         modify fun s => { s with choiceStack := [c] }

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

@@ -25,7 +25,7 @@ def addCongrTable (e : Expr) : GoalM Unit := do
     trace[grind.debug.congr] "{e} = {e'}"
     pushEqHEq e e' congrPlaceholderProof
     let node ← getENode e
-    setENode e { node with cgRoot := e' }
+    setENode e { node with congr := e' }
   else
     modify fun s => { s with congrTable := s.congrTable.insert { e } }
 

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

@@ -24,9 +24,9 @@ private def checkEqc (root : ENode) : GoalM Unit := do
     if curr.isApp then
       if let some { e } := (← get).congrTable.find? { e := curr } then
         if (← hasSameType e.getAppFn curr.getAppFn) then
-          assert! isSameExpr e (← getENode curr).cgRoot
+          assert! isSameExpr e (← getCongrRoot curr)
       else
-        assert! isSameExpr curr (← getENode curr).cgRoot
+        assert! (← isCongrRoot curr)
     -- If the equivalence class does not have HEq proofs, then the types must be definitionally equal.
     unless root.heqProofs do
       assert! (← hasSameType curr root.self)

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

@@ -173,8 +173,12 @@ structure ENode where
   next : Expr
   /-- Root (aka canonical representative) of the equivalence class -/
   root : Expr
-  /-- Root of the congruence class. This is field is a don't care if `e` is not an application. -/
-  cgRoot : Expr
+  /--
+  `congr` is the term `self` is congruent to.
+  We say `self` is the congruence class root if `isSameExpr congr self`.
+  This field is initialized to `self` even if `e` is not an application.
+  -/
+  congr : Expr
   /--
   When `e` was added to this equivalence class because of an equality `h : e = target`,
   then we store `target` here, and `h` at `proof?`.
@@ -206,6 +210,9 @@ structure ENode where
   -- TODO: see Lean 3 implementation
   deriving Inhabited, Repr
 
+def ENode.isCongrRoot (n : ENode) :=
+  isSameExpr n.self n.congr
+
 /-- New equality to be processed. -/
 structure NewEq where
   lhs   : Expr
@@ -540,7 +547,7 @@ def setENode (e : Expr) (n : ENode) : GoalM Unit :=
 
 def mkENodeCore (e : Expr) (interpreted ctor : Bool) (generation : Nat) : GoalM Unit := do
   setENode e {
-    self := e, next := e, root := e, cgRoot := e, size := 1
+    self := e, next := e, root := e, congr := e, size := 1
     flipped := false
     heqProofs := false
     hasLambdas := e.isLambda
@@ -562,7 +569,13 @@ def mkENode (e : Expr) (generation : Nat) : GoalM Unit := do
 
 /-- Returns `true` is `e` is the root of its congruence class. -/
 def isCongrRoot (e : Expr) : GoalM Bool := do
-  return isSameExpr e (← getENode e).cgRoot
+  return (← getENode e).isCongrRoot
+
+/-- Returns the root of the congruence class containing `e`. -/
+partial def getCongrRoot (e : Expr) : GoalM Expr := do
+  let n ← getENode e
+  if isSameExpr n.congr e then return e
+  getCongrRoot n.congr
 
 /-- Return `true` if the goal is inconsistent. -/
 def isInconsistent : GoalM Bool :=

tests/lean/run/grind_ematch2.lean

@@ -17,3 +17,15 @@ example (as bs cs : Array α) (v₁ v₂ : α)
         (h₆ : j < as.size)
         : cs[j] = as[j] := by
   grind
+
+example (as bs cs : Array α) (v₁ v₂ : α)
+        (i₁ i₂ j : Nat)
+        (h₁ : i₁ < as.size)
+        (h₂ : as.set i₁ v₁ = bs)
+        (h₃ : i₂ < bs.size)
+        (h₃ : bs.set i₂ v₂ = cs)
+        (h₄ : i₁ ≠ j ∧ j ≠ i₂)
+        (h₅ : j < cs.size)
+        (h₆ : j < as.size)
+        : cs[j] = as[j] := by
+  grind

tests/lean/run/grind_nested_proofs.lean

@@ -12,9 +12,8 @@ elab "grind_test" : tactic => withMainContext do
     let [_, n, _] := nodes.toList | unreachable!
     logInfo (← getEqc n.self)
 
--- TODO: fix
--- set_option grind.debug true
--- set_option grind.debug.proofs true
+set_option grind.debug true
+set_option grind.debug.proofs true
 
 /-
 Recall that array access terms, such as `a[i]`, have nested proofs.