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

@@ -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)

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

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