From 2a1354b3cc1eb05e1b67c656d2172ada8f54a937 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 24 May 2025 20:54:30 -0400
Subject: [PATCH] chore: add `seal` to workaround performance issue (#8469)

This PR adds `seal` commands at `grind_ite.lean` to workaround expensive
definitionally equality tests in the canonicalizer. The new module
system will automatically hide definitions such as `HashMap.insert` and
`TreeMap.insert` which are being unfolded by the canonicalizer in this
test.
This PR also adds a `profileItM` for tracking the time spent in the
`grind` canonicalizer.
---
 src/Lean/Meta/Tactic/Grind/Canon.lean | 12 ++++++++++--
 tests/lean/run/grind_ite.lean         |  5 +++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Canon.lean b/src/Lean/Meta/Tactic/Grind/Canon.lean
index d0b89fa4e0..b5abb9eba1 100644
--- a/src/Lean/Meta/Tactic/Grind/Canon.lean
+++ b/src/Lean/Meta/Tactic/Grind/Canon.lean
@@ -77,7 +77,15 @@ def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIsDefEqBou
   let eType ← inferType e
   let cs := s.argMap.find? key |>.getD []
   for (c, cType) in cs do
-    -- We first check the typesr
+    /-
+    We first check the types
+    The following checks are a performance bottleneck.
+    For example, in the test `grind_ite.lean`, there are many checks of the form:
+    ```
+    w_4 ∈ assign.insert v true → Prop =?= w_1 ∈ assign.insert v false → Prop
+    ```
+    where `grind` unfolds the definition of `DHashMap.insert` and `TreeMap.insert`.
+    -/
     if (← withDefault <| isDefEq eType cType) then
       if (← isDefEq e c) then
         -- We used to check `c.fvarsSubset e` because it is not
@@ -196,7 +204,7 @@ where
 end Canon
 
 /-- Canonicalizes nested types, type formers, and instances in `e`. -/
-def canon (e : Expr) : GoalM Expr := do
+def canon (e : Expr) : GoalM Expr := do profileitM Exception "grind canon" (← getOptions) do
   trace_goal[grind.debug.canon] "{e}"
   unsafe Canon.canonImpl e
 
diff --git a/tests/lean/run/grind_ite.lean b/tests/lean/run/grind_ite.lean
index 207e08bed8..da9dfc8573 100644
--- a/tests/lean/run/grind_ite.lean
+++ b/tests/lean/run/grind_ite.lean
@@ -126,6 +126,11 @@ we are allowed to increase the size of the branches by one, and still be smaller
   | var _ => 1
   | .ite i t e => 2 * normSize i + max (normSize t) (normSize e) + 1
 
+-- TODO: `grind` canonicalizer is spending a lot of time unfolding the following function.
+-- TODO: remove after the new module system will hide this declaration.
+seal Std.DHashMap.insert
+seal Std.TreeMap.insert
+
 def normalize (assign : Std.HashMap Nat Bool) : IfExpr → IfExpr
   | lit b => lit b
   | var v =>