From 533af01dab41fbdd36a854edd9571e6756a3dc72 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Tue, 21 Jan 2025 19:59:45 -0800
Subject: [PATCH] feat: improve `grind` canonicalizer (#6736)

This PR ensures the canonicalizer used in `grind` does not waste time
checking whether terms with different types are definitionally equal.
---
 src/Lean/Meta/Tactic/Grind/Canon.lean | 32 +++++++++++++++------------
 src/Lean/Meta/Tactic/Grind/Types.lean |  2 +-
 tests/lean/run/grind_t1.lean          |  7 ++++++
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Canon.lean b/src/Lean/Meta/Tactic/Grind/Canon.lean
index ae60269872..524e8c6575 100644
--- a/src/Lean/Meta/Tactic/Grind/Canon.lean
+++ b/src/Lean/Meta/Tactic/Grind/Canon.lean
@@ -75,24 +75,28 @@ def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIsDefEqBou
   if let some c := s.canon.find? e then
     return c
   let key := (f, i)
+  let eType ← inferType e
   let cs := s.argMap.find? key |>.getD []
-  for c in cs do
-    if (← isDefEq e c) then
-      -- We used to check `c.fvarsSubset e` because it is not
-      -- in general safe to replace `e` with `c` if `c` has more free variables than `e`.
-      -- However, we don't revert previously canonicalized elements in the `grind` tactic.
-      -- Moreover, we store the canonicalizer state in the `Goal` because we case-split
-      -- and different locals are added in different branches.
-      modify' fun s => { s with canon := s.canon.insert e c }
-      trace[grind.debugn.canon] "found {e} ===> {c}"
-      return c
-    if useIsDefEqBounded then
-      if (← isDefEqBounded e c parent) then
+  for (c, cType) in cs do
+    -- We first check the typesr
+    if (← withDefault <| isDefEq eType cType) then
+      if (← isDefEq e c) then
+        -- We used to check `c.fvarsSubset e` because it is not
+        -- in general safe to replace `e` with `c` if `c` has more free variables than `e`.
+        -- However, we don't revert previously canonicalized elements in the `grind` tactic.
+        -- Moreover, we store the canonicalizer state in the `Goal` because we case-split
+        -- and different locals are added in different branches.
         modify' fun s => { s with canon := s.canon.insert e c }
-        trace[grind.debugn.canon] "found using `isDefEqBounded`: {e} ===> {c}"
+        trace[grind.debugn.canon] "found {e} ===> {c}"
         return c
+      if useIsDefEqBounded then
+        -- If `e` and `c` are not types, we use `isDefEqBounded`
+        if (← isDefEqBounded e c parent) then
+          modify' fun s => { s with canon := s.canon.insert e c }
+          trace[grind.debugn.canon] "found using `isDefEqBounded`: {e} ===> {c}"
+          return c
   trace[grind.debug.canon] "({f}, {i}) ↦ {e}"
-  modify' fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key (e::cs) }
+  modify' fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key ((e, eType)::cs) }
   return e
 
 abbrev canonType (parent f : Expr) (i : Nat) (e : Expr) := withDefault <| canonElemCore parent f i e (useIsDefEqBounded := false)
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 20dd6497f8..9f4ec1071a 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -336,7 +336,7 @@ structure NewFact where
 
 /-- Canonicalizer state. See `Canon.lean` for additional details. -/
 structure Canon.State where
-  argMap     : PHashMap (Expr × Nat) (List Expr) := {}
+  argMap     : PHashMap (Expr × Nat) (List (Expr × Expr)) := {}
   canon      : PHashMap Expr Expr := {}
   proofCanon : PHashMap Expr Expr := {}
   deriving Inhabited
diff --git a/tests/lean/run/grind_t1.lean b/tests/lean/run/grind_t1.lean
index 39c8e89109..79428dbd4a 100644
--- a/tests/lean/run/grind_t1.lean
+++ b/tests/lean/run/grind_t1.lean
@@ -350,3 +350,10 @@ a✝ : b = true
 #guard_msgs (error) in
 example (b : Bool) : (if b then 10 else 20) = a → b = true → False := by
   grind
+
+-- Should not generate a trace message about canonicalization issues
+#guard_msgs (info) in
+set_option trace.grind.issues true in
+example : (if n + 2 < m then a else b) = (if n + 1 < m then c else d) := by
+  fail_if_success grind (splits := 0)
+  sorry