feat: add decidable equality to DTreeMap/TreeMap/TreeSet and their extensional variants

src/Std/Data/DTreeMap.lean

@@ -11,3 +11,4 @@ public import Std.Data.DTreeMap.AdditionalOperations
 public import Std.Data.DTreeMap.Lemmas
 public import Std.Data.DTreeMap.Iterator
 public import Std.Data.DTreeMap.Slice
+public import Std.Data.DTreeMap.DecidableEquiv

src/Std/Data/DTreeMap/Basic.lean

@@ -870,17 +870,11 @@ end Const
 
 /-- Check if any element satisfies the predicate, short-circuiting if a predicate fails. -/
 @[inline]
-def any (t : DTreeMap α β cmp) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
-  for ⟨a, b⟩ in t do
-    if p a b then return true
-  return false
+def any (t : DTreeMap α β cmp) (p : (a : α) → β a → Bool) : Bool := t.inner.any p
 
 /-- Check if all elements satisfy the predicate, short-circuiting if a predicate fails. -/
 @[inline]
-def all (t : DTreeMap α β cmp) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
-  for ⟨a, b⟩ in t do
-    if p a b = false then return false
-  return true
+def all (t : DTreeMap α β cmp) (p : (a : α) → β a → Bool) : Bool := t.inner.all p
 
 /-- Returns a list of all keys present in the tree map in ascending  order. -/
 @[inline]
@@ -1054,6 +1048,19 @@ def inter (t₁ t₂ : DTreeMap α β cmp) : DTreeMap α β cmp :=
 
 instance : Inter (DTreeMap α β cmp) := ⟨inter⟩
 
+/--
+Compares two tree maps using Boolean equality on keys and values.
+
+Returns `true` if the maps contain the same key-value pairs, `false` otherwise.
+-/
+def beq [LawfulEqCmp cmp] [∀ k, BEq (β k)] (t₁ t₂ : DTreeMap α β cmp) : Bool :=
+  letI : Ord α := ⟨cmp⟩; t₁.inner.beq t₂.inner
+
+instance [LawfulEqCmp cmp] [∀ k, BEq (β k)] : BEq (DTreeMap α β cmp) := ⟨beq⟩
+
+@[inherit_doc DTreeMap.beq] def Const.beq {β : Type v} [BEq β] (t₁ t₂ : DTreeMap α (fun _ => β) cmp) : Bool :=
+  letI : Ord α := ⟨cmp⟩; Internal.Impl.Const.beq t₁.inner t₂.inner
+
 /--
 Computes the difference of the given tree maps.
 This function always iterates through the smaller map.

src/Std/Data/DTreeMap/DecidableEquiv.lean

@@ -0,0 +1,27 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Wojciech Różowski
+-/
+module
+
+prelude
+public import Std.Data.DTreeMap.Internal.Lemmas
+public import Std.Data.DTreeMap.Raw
+
+public section
+
+/-!
+# Decidable equivalence for `DTreeMap`
+-/
+
+open Std.DTreeMap.Internal.Impl
+
+namespace Std.DTreeMap
+
+instance {α : Type u} {β : α → Type v} {cmp : α → α → Ordering} [TransCmp cmp] [LawfulEqCmp cmp] [DecidableEq α] [∀ k, DecidableEq (β k)] {t₁ t₂ : DTreeMap α β cmp} : Decidable (t₁ ~m t₂) :=
+  let : Ord α := ⟨cmp⟩;
+  let : Decidable (t₁.inner ~m t₂.inner) := decidableEquiv t₁.1 t₂.1 t₁.2 t₂.2;
+  decidable_of_iff _ ⟨fun h => ⟨h⟩, fun h => h.1⟩
+
+end Std.DTreeMap

src/Std/Data/DTreeMap/Internal/Lemmas.lean

@@ -100,6 +100,8 @@ private meta def queryMap : Std.DHashMap Name (fun _ => Name × Array (MacroM (T
      ⟨`getKeyD, (``getKeyD_eq_getKeyD, #[``(getKeyD_of_perm _)])⟩,
      ⟨`getKey!, (``getKey!_eq_getKey!, #[``(getKey!_of_perm _)])⟩,
      ⟨`toList, (``toList_eq_toListModel, #[])⟩,
+     ⟨`beq, (``beq_eq_beqModel, #[])⟩,
+     ⟨`Const.beq, (``Internal.Impl.Const.beq_eq_beqModel, #[])⟩,
      ⟨`keys, (``keys_eq_keys, #[])⟩,
      ⟨`Const.toList, (``Const.toList_eq_toListModel_map, #[])⟩,
      ⟨`foldlM, (``foldlM_eq_foldlM_toListModel, #[])⟩,
@@ -4851,6 +4853,37 @@ theorem get!_inter!_of_contains_eq_false_left [TransOrd α] [Inhabited β] (h₁
 
 end Const
 
+section BEq
+variable {m₁ m₂ : Impl α β} [TransOrd α] [LawfulEqOrd α] [∀ k, BEq (β k)]
+
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) : m₁.Equiv m₂ → beq m₁ m₂ := by
+  simp_to_model using List.beqModel_eq_true_of_perm
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) : beq m₁ m₂ = true → m₁.Equiv m₂ := by
+  simp_to_model using List.perm_of_beqModel
+
+theorem Equiv.beq_congr {m₃ m₄ : Impl α β} (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (h₄ : m₄.WF) :
+    m₁.Equiv m₃ → m₂.Equiv m₄ → (Impl.beq m₁ m₂) = (Impl.beq m₃ m₄) := by
+  simp_to_model using List.beqModel_congr
+
+end BEq
+
+section
+
+variable {β : Type v} [BEq β] {m₁ m₂ : Impl α (fun _ => β)}
+
+theorem Const.Equiv.beq [TransOrd α] [ReflBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : m₁.Equiv m₂ → Const.beq m₁ m₂ := by
+  simp_to_model using List.Const.beqModel_eq_true_of_perm
+
+theorem Const.equiv_of_beq [TransOrd α] [LawfulEqOrd α] [LawfulBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : Const.beq m₁ m₂ = true → m₁.Equiv m₂ := by
+  simp_to_model using List.Const.perm_of_beqModel
+
+theorem Const.Equiv.beq_congr [TransOrd α] {m₃ m₄ : Impl α (fun _ => β)} (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (h₄ : m₄.WF) :
+    m₁.Equiv m₃ → m₂.Equiv m₄ → Const.beq m₁ m₂ = Const.beq m₃ m₄ := by
+  simp_to_model using List.Const.beqModel_congr
+
+end
+
 section Diff
 
 variable {m₁ m₂ : Impl α β}
@@ -10341,4 +10374,12 @@ end Const
 
 end map
 
+section
+
+/-- Internal implementation detail -/
+def decidableEquiv {α : Type u} {β : α → Type v} [Ord α] [TransOrd α] [LawfulEqOrd α] [DecidableEq α] [∀ k, DecidableEq (β k)] (t₁ t₂ : Impl α β) (h₁ : t₁.WF) (h₂ : t₂.WF) : Decidable (t₁.Equiv t₂) :=
+  decidable_of_iff (beq t₁ t₂ = true) ⟨equiv_of_beq h₁ h₂, Equiv.beq h₁ h₂⟩
+
+end
+
 end Std.DTreeMap.Internal.Impl

src/Std/Data/DTreeMap/Internal/Operations.lean

@@ -822,6 +822,10 @@ information but still assumes the preconditions of `filter`, otherwise might pan
 def inter! [Ord α] (m₁ m₂ : Impl α β): Impl α β :=
   if m₁.size ≤ m₂.size then m₁.filter! (fun k _ => m₂.contains k) else interSmaller m₁ m₂
 
+/-- Internal implementation detail of the tree map -/
+def beq [Ord α] [LawfulEqOrd α] [∀ k, BEq (β k)] (t₁ t₂ : Impl α β) : Bool :=
+  if t₁.size ≠ t₂.size then false else t₁.all (fun k v => t₂.get? k == some v)
+
 /--
 Computes the difference of the given tree maps.
 This function always iterates through the smaller map.
@@ -950,6 +954,11 @@ def mergeWith! [Ord α] [LawfulEqOrd α] (mergeFn : (a : α) → β a → β a 
 namespace Const
 
 variable {β : Type v}
+
+/-- Internal implementation detail of the hash map -/
+def beq [Ord α] [BEq β] (t₁ t₂ : Impl α fun _ => β) : Bool :=
+  if t₁.size ≠ t₂.size then false else t₁.all (fun k v => Const.get? t₂ k == some v)
+
 local instance : Coe (Type v) (α → Type v) where coe γ := fun _ => γ
 
 /--

src/Std/Data/DTreeMap/Internal/Queries.lean

@@ -293,6 +293,20 @@ def forIn {m} [Monad m] (f : (a : α) → β a → δ → m (ForInStep δ)) (ini
 instance [Monad m] : ForIn m (Impl α β) ((a : α) × β a) where
   forIn m init f := m.forIn (fun a b acc => f ⟨a, b⟩ acc) init
 
+/-- Implementation detail. -/
+@[inline]
+def any (t : Impl α β) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
+  for ⟨a, b⟩ in t do
+    if p a b then return true
+  return false
+
+/-- Implementation detail. -/
+@[inline]
+def all (t : Impl α β) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
+  for ⟨a, b⟩ in t do
+    if p a b = false then return false
+  return true
+
 /-- Returns a `List` of the keys in order. -/
 @[inline] def keys (t : Impl α β) : List α :=
   t.foldr (init := []) fun k _ l => k :: l

src/Std/Data/DTreeMap/Internal/WF/Lemmas.lean

@@ -1900,6 +1900,29 @@ theorem WF.union! {_ : Ord α} [TransOrd α]
   . exact WF.insertManyIfNew! h₂
   . exact WF.insertMany! h₁
 
+theorem all_eq_all_toListModel {p : (a : α) → β a → Bool} {m : Impl α β} :
+    m.all p = m.toListModel.all (fun x => p x.1 x.2) := by
+  simp [all, ForIn.forIn, bind_pure_comp, map_pure, Id.run_bind]
+  rw [forIn_eq_forIn_toListModel, ← toList_eq_toListModel, forIn_eq_forIn']
+  induction m.toList with
+  | nil => simp
+  | cons hd tl ih =>
+    simp only [forIn'_eq_forIn, List.all_cons]
+    by_cases h : p hd.fst hd.snd = false
+    · simp [h]
+    · simp only [forIn'_eq_forIn] at ih
+      simp [h, ih]
+
+theorem beq_eq_beqModel {_ : Ord α} [BEq α] [TransOrd α] [LawfulBEq α] [LawfulBEqOrd α] [∀ k, BEq (β k)] {m₁ m₂ : Impl α β} (h₁ : m₁.WF) (h₂ : m₂.WF) :
+    Impl.beq m₁ m₂ = beqModel m₁.toListModel m₂.toListModel := by
+  simp [beq, beqModel, size_eq_length _ h₁.balanced, size_eq_length _ h₂.balanced, all_eq_all_toListModel,
+    get?_eq_getValueCast? h₂.ordered]
+
+theorem Const.beq_eq_beqModel {β : Type v} {_ : Ord α} [BEq α] [TransOrd α] [LawfulBEqOrd α] [BEq β] {m₁ m₂ : Impl α (fun _ => β)}  (h₁ : m₁.WF) (h₂ : m₂.WF) :
+    beq m₁ m₂ = Const.beqModel m₁.toListModel m₂.toListModel := by
+  simp [beq, Const.beqModel, size_eq_length _ h₁.balanced, size_eq_length _ h₂.balanced, all_eq_all_toListModel,
+    get?_eq_getValue? h₂.ordered]
+
 theorem WF.constInsertMany! {β : Type v} {_ : Ord α} [TransOrd α] {ρ} [ForIn Id ρ (α × β)] {l : ρ}
     {t : Impl α β} (h : t.WF) : (Const.insertMany! t l).1.WF :=
   (Const.insertMany! t l).2 h (fun _ _ _ h' => h'.insert!)

src/Std/Data/DTreeMap/Lemmas.lean

@@ -2879,6 +2879,36 @@ theorem get!_inter_of_not_mem_left [TransCmp cmp] [Inhabited β]
 
 end Const
 
+section BEq
+
+variable {m₁ m₂ : DTreeMap α β cmp} [∀ k, BEq (β k)] [LawfulEqCmp cmp] [TransCmp cmp]
+
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h : m₁ ~m m₂) : m₁ == m₂ :=
+  Impl.Equiv.beq m₁.2 m₂.2 h.1
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h : m₁ == m₂) : m₁ ~m m₂ :=
+  ⟨Impl.equiv_of_beq m₁.2 m₂.2 h⟩
+
+theorem Equiv.beq_congr {m₃ m₄ : DTreeMap α β cmp} (w₁ : m₁ ~m m₃) (w₂ : m₂ ~m m₄) : (m₁ == m₂) = (m₃ == m₄) :=
+  Impl.Equiv.beq_congr m₁.2 m₂.2 m₃.2 m₄.2 w₁.1 w₂.1
+
+end BEq
+
+section
+
+variable {β : Type v} {m₁ m₂ : DTreeMap α (fun _ => β) cmp} [BEq β]
+
+theorem Const.Equiv.beq [TransCmp cmp] [ReflBEq β] (h : m₁ ~m m₂) : DTreeMap.Const.beq m₁ m₂ :=
+  Impl.Const.Equiv.beq m₁.2 m₂.2 h.1
+
+theorem Const.equiv_of_beq [TransCmp cmp] [LawfulEqCmp cmp] [LawfulBEq β] (h : Const.beq m₁ m₂) : m₁ ~m m₂ :=
+  ⟨Impl.Const.equiv_of_beq m₁.2 m₂.2 h⟩
+
+theorem Const.Equiv.beq_congr [TransCmp cmp] {m₃ m₄ : DTreeMap α (fun _ => β) cmp} (w₁ : m₁ ~m m₃) (w₂ : m₂ ~m m₄) : Const.beq m₁ m₂ = Const.beq m₃ m₄ :=
+  Impl.Const.Equiv.beq_congr m₁.2 m₂.2 m₃.2 m₄.2 w₁.1 w₂.1
+
+end
+
 section Diff
 
 variable {t₁ t₂ : DTreeMap α β cmp}
@@ -5351,10 +5381,10 @@ theorem forM_eq [TransCmp cmp] [Monad m] [LawfulMonad m] {f : (a : α) × β a 
   h.1.forM_eq t₁.2 t₂.2
 
 theorem any_eq [TransCmp cmp] {p : (a : α) → β a → Bool} (h : t₁ ~m t₂) : t₁.any p = t₂.any p := by
-  simp only [any, h.forIn_eq]
+  simp only [any, Impl.any, ForIn.forIn, h.1.forIn_eq t₁.2 t₂.2]
 
 theorem all_eq [TransCmp cmp] {p : (a : α) → β a → Bool} (h : t₁ ~m t₂) : t₁.all p = t₂.all p := by
-  simp only [all, h.forIn_eq]
+  simp only [all, Impl.all, ForIn.forIn, h.1.forIn_eq t₁.2 t₂.2]
 
 theorem minKey?_eq [TransCmp cmp] (h : t₁ ~m t₂) : t₁.minKey? = t₂.minKey? :=
   h.1.minKey?_eq t₁.2 t₂.2

src/Std/Data/DTreeMap/Raw.lean

@@ -12,3 +12,4 @@ public import Std.Data.DTreeMap.Raw.Lemmas
 public import Std.Data.DTreeMap.Raw.WF
 public import Std.Data.DTreeMap.Raw.Iterator
 public import Std.Data.DTreeMap.Raw.Slice
+public import Std.Data.DTreeMap.Raw.DecidableEquiv

src/Std/Data/DTreeMap/Raw/Basic.lean

@@ -604,16 +604,12 @@ def forInUncurried (f : α × β → δ → m (ForInStep δ)) (init : δ) (t : R
 end Const
 
 @[inline, inherit_doc DTreeMap.any]
-def any (t : Raw α β cmp) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
-  for ⟨a, b⟩ in t do
-    if p a b then return true
-  return false
+def any (t : Raw α β cmp) (p : (a : α) → β a → Bool) : Bool :=
+  t.inner.any p
 
 @[inline, inherit_doc DTreeMap.all]
 def all (t : Raw α β cmp) (p : (a : α) → β a → Bool) : Bool := Id.run $ do
-  for ⟨a, b⟩ in t do
-    if p a b = false then return false
-  return true
+  t.inner.all p
 
 @[inline, inherit_doc DTreeMap.keys]
 def keys (t : Raw α β cmp) : List α :=
@@ -732,6 +728,15 @@ def inter (t₁ t₂ : Raw α β cmp) : Raw α β cmp :=
 
 instance : Inter (Raw α β cmp) := ⟨inter⟩
 
+
+@[inherit_doc DTreeMap.beq] def beq [LawfulEqCmp cmp] [∀ k, BEq (β k)] (t₁ t₂ : Raw α β cmp) : Bool :=
+  letI : Ord α := ⟨cmp⟩; t₁.inner.beq t₂.inner
+
+instance [LawfulEqCmp cmp] [∀ k, BEq (β k)] : BEq (Raw α β cmp) := ⟨beq⟩
+
+@[inherit_doc DTreeMap.beq] def Const.beq {β : Type v} [BEq β] (t₁ t₂ : Raw α (fun _ => β) cmp) : Bool :=
+  letI : Ord α := ⟨cmp⟩; Internal.Impl.Const.beq t₁.inner t₂.inner
+
 /--
 Computes the difference of the given tree maps.
 This function always iterates through the smaller map.

src/Std/Data/DTreeMap/Raw/DecidableEquiv.lean

@@ -0,0 +1,27 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Wojciech Różowski
+-/
+module
+
+prelude
+public import Std.Data.DTreeMap.Internal.Lemmas
+public import Std.Data.DTreeMap.Raw.Basic
+
+public section
+
+/-!
+# Decidable equivalence for `DTreeMap.Raw`
+-/
+
+open Std.DTreeMap.Internal.Impl
+
+namespace Std.DTreeMap.Raw
+
+instance instDecidableEquiv {α : Type u} {β : α → Type v} {cmp : α → α → Ordering} [TransCmp cmp] [LawfulEqCmp cmp] [DecidableEq α] [∀ k, DecidableEq (β k)] {t₁ t₂ : Raw α β cmp} (h₁ : t₁.WF) (h₂ : t₂.WF) : Decidable (t₁ ~m t₂) :=
+  let : Ord α := ⟨cmp⟩;
+  let : Decidable (t₁.inner ~m t₂.inner) := decidableEquiv t₁.1 t₂.1 h₁ h₂;
+  decidable_of_iff _ ⟨fun h => ⟨h⟩, fun h => h.1⟩
+
+end Std.DTreeMap.Raw

src/Std/Data/DTreeMap/Raw/Lemmas.lean

@@ -3044,6 +3044,36 @@ theorem get!_inter_of_not_mem_left [TransCmp cmp] [Inhabited β] (h₁ : m₁.WF
 
 end Const
 
+section BEq
+variable {m₁ m₂ : Raw α β cmp} [∀ k, BEq (β k)] [LawfulEqCmp cmp] [TransCmp cmp]
+
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) (h : m₁ ~m m₂) : beq m₁ m₂ :=
+  Impl.Equiv.beq h₁ h₂ h.1
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) : beq m₁ m₂ = true → m₁ ~m m₂ := fun hyp =>
+  let : Ord α := ⟨cmp⟩; ⟨Impl.equiv_of_beq h₁.1 h₂.1 hyp⟩
+
+theorem Equiv.beq_congr {m₃ m₄ : Raw α β cmp} (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (h₄ : m₄.WF) :
+    m₁ ~m m₃ → m₂ ~m m₄ → Raw.beq m₁ m₂ = Raw.beq m₃ m₄ :=
+  fun w1 w2 => Impl.Equiv.beq_congr h₁ h₂ h₃ h₄ w1.1 w2.1
+
+end BEq
+
+section
+variable {β : Type v} {m₁ m₂ : Raw α (fun _ => β) cmp}
+
+theorem Const.Equiv.beq [TransCmp cmp] [BEq β] [ReflBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : m₁ ~m m₂ → beq m₁ m₂ :=
+  fun h => Impl.Const.Equiv.beq h₁ h₂ h.1
+
+theorem Const.equiv_of_beq [TransCmp cmp] [LawfulEqCmp cmp] [BEq β] [LawfulBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : beq m₁ m₂ = true → m₁ ~m m₂ :=
+  fun hyp => let : Ord α := ⟨cmp⟩; ⟨Impl.Const.equiv_of_beq h₁.1 h₂.1 hyp⟩
+
+theorem Const.Equiv.beq_congr [TransCmp cmp] [BEq β] {m₃ m₄ : Raw α (fun _ => β) cmp} (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (h₄ : m₄.WF) :
+    m₁ ~m m₃ → m₂ ~m m₄ → Raw.Const.beq m₁ m₂ = Raw.Const.beq m₃ m₄ :=
+  fun w1 w2 => Impl.Const.Equiv.beq_congr h₁ h₂ h₃ h₄ w1.1 w2.1
+
+end
+
 section Diff
 
 variable {t₁ t₂ : Raw α β cmp}
@@ -5269,11 +5299,13 @@ theorem forM_eq [TransCmp cmp] [Monad m] [LawfulMonad m] {f : (a : α) × β a 
 
 theorem any_eq [TransCmp cmp] {p : (a : α) → β a → Bool} (h₁ : t₁.WF) (h₂ : t₂.WF) (h : t₁ ~m t₂) :
     t₁.any p = t₂.any p := by
-  simp only [any, h.forIn_eq h₁ h₂]
+  simp only [any, Impl.any, ForIn.forIn, bind_pure_comp, map_pure, h.1.forIn_eq h₁.1 h₂.1,
+    Id.run_bind]
 
 theorem all_eq [TransCmp cmp] {p : (a : α) → β a → Bool} (h₁ : t₁.WF) (h₂ : t₂.WF) (h : t₁ ~m t₂) :
     t₁.all p = t₂.all p := by
-  simp only [all, h.forIn_eq h₁ h₂]
+  simp only [all, Impl.all, ForIn.forIn, bind_pure_comp, map_pure, h.1.forIn_eq h₁.1 h₂.1,
+    Id.run_bind]
 
 theorem minKey?_eq [TransCmp cmp] (h₁ : t₁.WF) (h₂ : t₂.WF) (h : t₁ ~m t₂) :
     t₁.minKey? = t₂.minKey? :=