From d8335bd661c7e8b40b994f3aa98af4f237688359 Mon Sep 17 00:00:00 2001
From: Kim Morrison <477956+kim-em@users.noreply.github.com>
Date: Tue, 28 Oct 2025 16:20:45 +1100
Subject: [PATCH] feat: add `@[grind ext]` attributes for extensional maps
 (#10993)

This PR allows `grind` to work extensionally on extensional maps/sets.
---
 src/Std/Data/ExtDHashMap/Lemmas.lean |  2 +-
 src/Std/Data/ExtDTreeMap/Lemmas.lean |  2 +-
 src/Std/Data/ExtHashMap/Lemmas.lean  |  2 +-
 src/Std/Data/ExtHashSet/Lemmas.lean  |  2 +-
 src/Std/Data/ExtTreeMap/Lemmas.lean  |  2 +-
 src/Std/Data/ExtTreeSet/Lemmas.lean  |  2 +-
 tests/lean/run/grind_ExtTreeSet.lean | 17 +++++++++++++++++
 7 files changed, 23 insertions(+), 6 deletions(-)
 create mode 100644 tests/lean/run/grind_ExtTreeSet.lean

diff --git a/src/Std/Data/ExtDHashMap/Lemmas.lean b/src/Std/Data/ExtDHashMap/Lemmas.lean
index 058f9b0f4c..8e8f4a106e 100644
--- a/src/Std/Data/ExtDHashMap/Lemmas.lean
+++ b/src/Std/Data/ExtDHashMap/Lemmas.lean
@@ -2791,7 +2791,7 @@ section Ext
 
 variable {m₁ m₂ : Std.ExtDHashMap α β}
 
-@[ext]
+@[ext, grind ext]
 theorem ext_get? [LawfulBEq α] {m₁ m₂ : Std.ExtDHashMap α β} (h : ∀ k, m₁.get? k = m₂.get? k) : m₁ = m₂ :=
   m₁.inductionOn₂ m₂ (fun _ _ h => sound (.of_forall_get?_eq h)) h
 
diff --git a/src/Std/Data/ExtDTreeMap/Lemmas.lean b/src/Std/Data/ExtDTreeMap/Lemmas.lean
index 9e4c754cc8..cadc0d1498 100644
--- a/src/Std/Data/ExtDTreeMap/Lemmas.lean
+++ b/src/Std/Data/ExtDTreeMap/Lemmas.lean
@@ -4369,7 +4369,7 @@ section Ext
 
 variable {t₁ t₂ : ExtDTreeMap α β cmp}
 
-@[ext]
+@[ext, grind ext]
 theorem ext_get? [TransCmp cmp] [LawfulEqCmp cmp] {t₁ t₂ : ExtDTreeMap α β cmp}
     (h : ∀ k, t₁.get? k = t₂.get? k) : t₁ = t₂ :=
   t₁.inductionOn₂ t₂ (fun _ _ h => sound (.of_forall_get?_eq h)) h
diff --git a/src/Std/Data/ExtHashMap/Lemmas.lean b/src/Std/Data/ExtHashMap/Lemmas.lean
index 5511d31603..99cd23adad 100644
--- a/src/Std/Data/ExtHashMap/Lemmas.lean
+++ b/src/Std/Data/ExtHashMap/Lemmas.lean
@@ -1698,7 +1698,7 @@ section Ext
 
 variable {m₁ m₂ : ExtHashMap α β}
 
-@[ext 900]
+@[ext 900, grind ext]
 theorem ext_getKey_getElem? [EquivBEq α] [LawfulHashable α]
     {m₁ m₂ : ExtHashMap α β}
     (hk : ∀ k hk hk', m₁.getKey k hk = m₂.getKey k hk')
diff --git a/src/Std/Data/ExtHashSet/Lemmas.lean b/src/Std/Data/ExtHashSet/Lemmas.lean
index 938500599d..bbe1ea6473 100644
--- a/src/Std/Data/ExtHashSet/Lemmas.lean
+++ b/src/Std/Data/ExtHashSet/Lemmas.lean
@@ -669,7 +669,7 @@ end
 
 section Ext
 
-@[ext 900]
+@[ext 900, grind ext]
 theorem ext_get? [EquivBEq α] [LawfulHashable α] {m₁ m₂ : ExtHashSet α}
     (h : ∀ k, m₁.get? k = m₂.get? k) : m₁ = m₂ :=
   ext (ExtHashMap.ext_getKey?_unit h)
diff --git a/src/Std/Data/ExtTreeMap/Lemmas.lean b/src/Std/Data/ExtTreeMap/Lemmas.lean
index 279278e45c..9c098760bb 100644
--- a/src/Std/Data/ExtTreeMap/Lemmas.lean
+++ b/src/Std/Data/ExtTreeMap/Lemmas.lean
@@ -3015,7 +3015,7 @@ section Ext
 
 variable {t₁ t₂ : ExtTreeMap α β cmp}
 
-@[ext 900]
+@[ext 900, grind ext]
 theorem ext_getKey_getElem? [TransCmp cmp] {t₁ t₂ : ExtTreeMap α β cmp}
     (hk : ∀ k hk hk', t₁.getKey k hk = t₂.getKey k hk')
     (hv : ∀ k : α, t₁[k]? = t₂[k]?) : t₁ = t₂ :=
diff --git a/src/Std/Data/ExtTreeSet/Lemmas.lean b/src/Std/Data/ExtTreeSet/Lemmas.lean
index 2d409e49fd..c104f80e7e 100644
--- a/src/Std/Data/ExtTreeSet/Lemmas.lean
+++ b/src/Std/Data/ExtTreeSet/Lemmas.lean
@@ -1544,7 +1544,7 @@ section Ext
 
 variable {t₁ t₂ : ExtTreeSet α cmp}
 
-@[ext 900]
+@[ext 900, grind ext]
 theorem ext_get? [TransCmp cmp] {t₁ t₂ : ExtTreeSet α cmp}
     (h : ∀ k, t₁.get? k = t₂.get? k) : t₁ = t₂ :=
   ext (ExtTreeMap.ext_getKey?_unit h)
diff --git a/tests/lean/run/grind_ExtTreeSet.lean b/tests/lean/run/grind_ExtTreeSet.lean
new file mode 100644
index 0000000000..3b340de940
--- /dev/null
+++ b/tests/lean/run/grind_ExtTreeSet.lean
@@ -0,0 +1,17 @@
+import Std
+
+open Std
+
+example {cmp : α → α → Ordering} [TransCmp cmp] (m : ExtTreeSet α cmp) (x : α) :
+    (m.insert x).erase x = m.erase x := by
+  grind
+
+example (m : ExtTreeSet Nat compare) (x : Nat) : (m.insert x).erase x = m.erase x := by
+  grind
+
+example [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α] (m : ExtHashSet α) (x : α) :
+    (m.insert x).erase x = m.erase x := by
+  grind
+
+example (m : ExtHashMap Int Int) (x y : Int) : (m.insert x y).erase x = m.erase x := by
+  grind