fix: check function types when detecting congruences in grind (#6466)

This PR completes the implementation of `addCongrTable` in the (WIP)
`grind` tactic. It also adds a new test to demonstrate why the extra
check is needed. It also updates the field `cgRoot` (congruence root).

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

@@ -13,10 +13,18 @@ namespace Lean.Meta.Grind
 /-- Adds `e` to congruence table. -/
 def addCongrTable (e : Expr) : GoalM Unit := do
   if let some { e := e' } := (← get).congrTable.find? { e } then
+    -- `f` and `g` must have the same type.
+    -- See paper: Congruence Closure in Intensional Type Theory
+    let f := e.getAppFn
+    let g := e'.getAppFn
+    unless isSameExpr f g do
+      unless (← hasSameType f g) do
+        trace[grind.issues] "found congruence between{indentExpr e}\nand{indentExpr e'}\nbut functions have different types"
+        return ()
     trace[grind.congr] "{e} = {e'}"
     pushEqHEq e e' congrPlaceholderProof
-    -- TODO: we must check whether the types of the functions are the same
-    -- TODO: update cgRoot for `e`
+    let node ← getENode e
+    setENode e { node with cgRoot := e' }
   else
     modify fun s => { s with congrTable := s.congrTable.insert { e } }
 

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

@@ -20,9 +20,16 @@ private def checkEqc (root : ENode) : GoalM Unit := do
     size := size + 1
     -- The root of `curr` must be `root`
     assert! isSameExpr (← getRoot curr) root.self
+    -- Check congruence root
+    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
+      else
+        assert! isSameExpr curr (← getENode curr).cgRoot
     -- 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))
+      assert! (← hasSameType curr root.self)
     -- Starting at `curr`, following the `target?` field leads to `root`.
     let mut n := curr
     repeat

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

@@ -226,7 +226,7 @@ where
       trace[grind.proof] "{a} = {b}"
       mkEqProofCore a b (heq := false)
     else
-      if (← withDefault <| isDefEq (← inferType a) (← inferType b)) then
+      if (← hasSameType a b) then
         trace[grind.proof] "{a} = {b}"
         mkEqOfHEq (← mkEqProofCore a b (heq := true))
       else

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

@@ -356,8 +356,12 @@ Otherwise, it pushes `HEq lhs rhs`.
 def pushEqCore (lhs rhs proof : Expr) (isHEq : Bool) : GoalM Unit :=
   modify fun s => { s with newEqs := s.newEqs.push { lhs, rhs, proof, isHEq } }
 
+/-- Return `true` if `a` and `b` have the same type. -/
+def hasSameType (a b : Expr) : MetaM Bool :=
+  withDefault do isDefEq (← inferType a) (← inferType b)
+
 @[inline] def pushEqHEq (lhs rhs proof : Expr) : GoalM Unit := do
-  if (← withDefault <| isDefEq (← inferType lhs) (← inferType rhs)) then
+  if (← hasSameType lhs rhs) then
     pushEqCore lhs rhs proof (isHEq := false)
   else
     pushEqCore lhs rhs proof (isHEq := true)

tests/lean/run/grind_pre.lean

@@ -100,3 +100,21 @@ example (a b : Nat) (f : Nat → Nat) : (h₁ : a = b) → (h₂ : f a ≠ f b)
 
 example (a : α) (p q r : Prop) : (h₁ : HEq p a) → (h₂ : HEq q a) → (h₃ : p = r) → q = r := by
   grind
+
+/--
+warning: declaration uses 'sorry'
+---
+info: [grind.issues] found congruence between
+      g b
+    and
+      f a
+    but functions have different types
+-/
+#guard_msgs in
+set_option trace.grind.issues true in
+set_option trace.grind.proof.detail false in
+set_option trace.grind.proof false in
+set_option trace.grind.pre false in
+example (f : Nat → Bool) (g : Int → Bool) (a : Nat) (b : Int) : HEq f g → HEq a b → f a = g b := by
+  fail_if_success grind
+  sorry