From 84cdab1da42bba27ad1634e47f7c753840927cf7 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Wed, 8 Jan 2025 17:55:20 -0800
Subject: [PATCH] fix: `grind` canonicalizer

This PR fixes a bug in the `grind` canonicalizer.
---
 src/Lean/Meta/Tactic/Grind.lean       |  2 +-
 src/Lean/Meta/Tactic/Grind/Canon.lean | 17 ++++++++++++++---
 tests/lean/run/grind_canon_bug.lean   | 13 +++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)
 create mode 100644 tests/lean/run/grind_canon_bug.lean

diff --git a/src/Lean/Meta/Tactic/Grind.lean b/src/Lean/Meta/Tactic/Grind.lean
index 80edeeef6ca..47f4b491181 100644
--- a/src/Lean/Meta/Tactic/Grind.lean
+++ b/src/Lean/Meta/Tactic/Grind.lean
@@ -53,5 +53,5 @@ builtin_initialize registerTraceClass `grind.debug.parent
 builtin_initialize registerTraceClass `grind.debug.final
 builtin_initialize registerTraceClass `grind.debug.forallPropagator
 builtin_initialize registerTraceClass `grind.debug.split
-
+builtin_initialize registerTraceClass `grind.debug.canon
 end Lean
diff --git a/src/Lean/Meta/Tactic/Grind/Canon.lean b/src/Lean/Meta/Tactic/Grind/Canon.lean
index 6ac02c336eb..17c3f43fc6d 100644
--- a/src/Lean/Meta/Tactic/Grind/Canon.lean
+++ b/src/Lean/Meta/Tactic/Grind/Canon.lean
@@ -66,11 +66,13 @@ def canonElemCore (f : Expr) (i : Nat) (e : Expr) (isInst : Bool) : StateT State
       -- 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.
       modify fun s => { s with canon := s.canon.insert e c }
+      trace[grind.debug.canon] "found {e} ===> {c}"
       return c
     if isInst then
-      if (← isDiagnosticsEnabled <&&> pure (c.fvarsSubset e) <&&> (withDefault <| isDefEq e c)) then
+      if (← isDiagnosticsEnabled <&&> (withDefault <| isDefEq e c)) then
         -- TODO: consider storing this information in some structure that can be browsed later.
         trace[grind.issues] "the following `grind` static elements are definitionally equal with `default` transparency, but not with `instances` transparency{indentExpr e}\nand{indentExpr 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) }
   return e
 
@@ -92,6 +94,12 @@ private inductive ShouldCanonResult where
     visit
   deriving Inhabited
 
+instance : Repr ShouldCanonResult where
+  reprPrec r _ := match r with
+    | .canonType => "canonType"
+    | .canonInst => "canonInst"
+    | .visit => "visit"
+
 /--
 See comments at `ShouldCanonResult`.
 -/
@@ -102,7 +110,9 @@ def shouldCanon (pinfos : Array ParamInfo) (i : Nat) (arg : Expr) : MetaM Should
       return .canonInst
     else if pinfo.isProp then
       return .visit
-  if (← isTypeFormer arg) then
+  if (← isProp arg) then
+    return .visit
+  else if (← isTypeFormer arg) then
     return .canonType
   else
     return .visit
@@ -151,7 +161,8 @@ where
     return e'
 
 /-- Canonicalizes nested types, type formers, and instances in `e`. -/
-def canon (e : Expr) : StateT State MetaM Expr :=
+def canon (e : Expr) : StateT State MetaM Expr := do
+  trace[grind.debug.canon] "{e}"
   unsafe canonImpl e
 
 end Canon
diff --git a/tests/lean/run/grind_canon_bug.lean b/tests/lean/run/grind_canon_bug.lean
new file mode 100644
index 00000000000..1605f7a8921
--- /dev/null
+++ b/tests/lean/run/grind_canon_bug.lean
@@ -0,0 +1,13 @@
+open List
+
+attribute [grind =] getElem_cons_zero
+attribute [grind =] getElem?_cons_zero
+
+example (h : (a :: t)[0]? = some b) : (a :: t)[0] = b := by
+  grind -- ✓
+
+example [Inhabited α] : ∀ {l : List α} {i : Nat}, l[i]? = some b → l[i]! = b
+  | a::t, 0, _ => by
+    rw [getElem!_pos _ _ sorry]
+    grind
+  | _::l, _+1, e => sorry