From a43888f07876c759416bbf1e8cb716e62d3636c8 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Mon, 23 Dec 2024 15:48:24 -0800
Subject: [PATCH] feat: check invariant for `parents`

---
 src/Lean/Meta/Tactic/Grind/Core.lean  |  6 +++---
 src/Lean/Meta/Tactic/Grind/Inv.lean   | 15 +++++++++++++++
 src/Lean/Meta/Tactic/Grind/Types.lean | 14 ++++++++++----
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Core.lean b/src/Lean/Meta/Tactic/Grind/Core.lean
index a2eb0a7a958d..dc469d883ad3 100644
--- a/src/Lean/Meta/Tactic/Grind/Core.lean
+++ b/src/Lean/Meta/Tactic/Grind/Core.lean
@@ -185,10 +185,10 @@ Remove `root` parents from the congruence table.
 This is an auxiliary function performed while merging equivalence classes.
 -/
 private def removeParents (root : Expr) : GoalM ParentSet := do
-  let occs ← getParents root
-  for parent in occs do
+  let parents ← getParentsAndReset root
+  for parent in parents do
     modify fun s => { s with congrTable := s.congrTable.erase { e := parent } }
-  return occs
+  return parents
 
 /--
 Reinsert parents into the congruence table and detect new equalities.
diff --git a/src/Lean/Meta/Tactic/Grind/Inv.lean b/src/Lean/Meta/Tactic/Grind/Inv.lean
index 5e20ababe125..fdd59ad2a04e 100644
--- a/src/Lean/Meta/Tactic/Grind/Inv.lean
+++ b/src/Lean/Meta/Tactic/Grind/Inv.lean
@@ -40,12 +40,27 @@ private def checkEqc (root : ENode) : GoalM Unit := do
   -- The size of the equivalence class is correct.
   assert! root.size == size
 
+private def checkParents (e : Expr) : GoalM Unit := do
+  if (← isRoot e) then
+    for parent in (← getParents e) do
+      let mut found := false
+      -- There is an argument `arg` s.t. root of `arg` is `e`.
+      for arg in parent.getAppArgs do
+        if isSameExpr (← getRoot arg) e then
+          found := true
+          break
+      assert! found
+  else
+    -- All the parents are stored in the root of the equivalence class.
+    assert! (← getParents e).isEmpty
+
 /--
 Check basic invariants if `grind.debug` is enabled.
 -/
 def checkInvariants : GoalM Unit := do
   if grind.debug.get (← getOptions) then
     for (_, node) in (← get).enodes do
+      checkParents node.self
       if isSameExpr node.self node.root then
         checkEqc node
 
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 9a58f6877e00..a6754eb60646 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -302,10 +302,16 @@ Returns the set of expressions `e` is a child of, or an expression in
 The information is only up to date if `e` is the root (aka canonical representative) of the equivalence class.
 -/
 def getParents (e : Expr) : GoalM ParentSet := do
-  if let some occs := (← get).parents.find? (toENodeKey e) then
-    return occs
-  else
-    return {}
+  let some parents := (← get).parents.find? (toENodeKey e) | return {}
+  return parents
+
+/--
+Similar to `getParents`, but also removes the entry `e ↦ parents` from the parent map.
+-/
+def getParentsAndReset (e : Expr) : GoalM ParentSet := do
+  let parents ← getParents e
+  modify fun s => { s with parents := s.parents.erase (toENodeKey e) }
+  return parents
 
 /--
 Copy `parents` to the parents of `root`.