feat: add BEq to DHashMap/HashMap/HashSet and their extensional variants

src/Std/Data/DHashMap/Basic.lean

@@ -360,6 +360,16 @@ instance [BEq α] [Hashable α] : Union (DHashMap α β) := ⟨union⟩
 instance [BEq α] [Hashable α] : Inter (DHashMap α β) := ⟨inter⟩
 instance [BEq α] [Hashable α] : SDiff (DHashMap α β) := ⟨diff⟩
 
+/-- Internal implementation detail of the hash map. -/
+def beq [LawfulBEq α] [∀ k, BEq (β k)] (a b : DHashMap α β) : Bool :=
+  Raw₀.beq ⟨a.1, a.2.size_buckets_pos⟩ ⟨b.1, b.2.size_buckets_pos⟩
+
+instance [LawfulBEq α] [∀ k, BEq (β k)] : BEq (DHashMap α β) := ⟨beq⟩
+
+/-- Internal implementation detail of the hash map. -/
+def Const.beq {β : Type v} [BEq α] [BEq β] (m₁ m₂ : DHashMap α (fun _ => β)) : Bool :=
+  Raw₀.Const.beq ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩
+
 section Unverified
 
 /-! We currently do not provide lemmas for the functions below. -/

src/Std/Data/DHashMap/Internal/Defs.lean

@@ -490,11 +490,14 @@ def interSmaller [BEq α] [Hashable α] (m₁ : Raw₀ α β) (m₂ : Raw α β)
 def inter [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
   if m₁.1.size ≤ m₂.1.size then m₁.filter fun k _ => m₂.contains k else interSmaller m₁ m₂
 
+/-- Internal implementation detail of the hash map -/
+def beq [BEq α] [LawfulBEq α] [Hashable α] [∀ k, BEq (β k)] (m₁ m₂ : Raw₀ α β) : Bool :=
+  if m₁.1.size ≠ m₂.1.size then false else m₁.1.all (fun k v => m₂.get? k == some v)
+
 /-- Internal implementation detail of the hash map -/
 @[inline] def diff [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
   if m₁.1.size ≤ m₂.1.size then m₁.filter (fun k _ => !m₂.contains k) else (eraseManyEntries m₁ m₂.1).1
 
-
 section
 
 variable {β : Type v}
@@ -505,6 +508,10 @@ def Const.get? [BEq α] [Hashable α] (m : Raw₀ α (fun _ => β)) (a : α) : O
   let ⟨i, h⟩ := mkIdx buckets.size h (hash a)
   buckets[i].get? a
 
+/-- Internal implementation detail of the hash map -/
+def Const.beq [BEq α] [Hashable α] [BEq β] (m₁ m₂ : Raw₀ α (fun _ => β)) : Bool :=
+  if m₁.1.size ≠ m₂.1.size then false else m₁.1.all (fun k v => Const.get? m₂ k == some v)
+
 /-- Internal implementation detail of the hash map -/
 def Const.get [BEq α] [Hashable α] (m : Raw₀ α (fun _ => β)) (a : α)
     (hma : m.contains a) : β :=

src/Std/Data/DHashMap/Internal/Raw.lean

@@ -154,6 +154,14 @@ theorem inter_eq [BEq α] [Hashable α] {m₁ m₂ : Raw α β} (h₁ : m₁.WF)
     m₁.inter m₂ = Raw₀.inter ⟨m₁, h₁.size_buckets_pos⟩ ⟨m₂, h₂.size_buckets_pos⟩ := by
   simp [Raw.inter, h₁.size_buckets_pos, h₂.size_buckets_pos]
 
+theorem beq_eq [BEq α] [Hashable α] [LawfulBEq α] [∀ k, BEq (β k)] {m₁ m₂ : Raw α β} (h₁ : m₁.WF) (h₂ : m₂.WF) :
+     m₁.beq m₂ = Raw₀.beq ⟨m₁, h₁.size_buckets_pos⟩ ⟨m₂, h₂.size_buckets_pos⟩ := by
+  simp [Raw.beq, h₁.size_buckets_pos, h₂.size_buckets_pos]
+
+theorem Const.beq_eq {β : Type v} [BEq α] [Hashable α] [BEq β] {m₁ m₂ : Raw α (fun _ => β)} (h₁ : m₁.WF) (h₂ : m₂.WF) :
+     Raw.Const.beq m₁ m₂ = Raw₀.Const.beq ⟨m₁, h₁.size_buckets_pos⟩ ⟨m₂, h₂.size_buckets_pos⟩ := by
+  simp [Raw.Const.beq, h₁.size_buckets_pos, h₂.size_buckets_pos]
+
 theorem diff_eq [BEq α] [Hashable α] {m₁ m₂ : Raw α β} (h₁ : m₁.WF) (h₂ : m₂.WF) :
     m₁.diff m₂ = Raw₀.diff ⟨m₁, h₁.size_buckets_pos⟩ ⟨m₂, h₂.size_buckets_pos⟩ := by
   simp [Raw.diff, h₁.size_buckets_pos, h₂.size_buckets_pos]

src/Std/Data/DHashMap/Internal/RawLemmas.lean

@@ -11,6 +11,7 @@ import all Std.Data.DHashMap.Internal.Defs
 public import Std.Data.DHashMap.Internal.WF
 import all Std.Data.DHashMap.Raw
 import all Std.Data.DHashMap.Basic
+import all Std.Data.DHashMap.RawDef
 meta import Std.Data.DHashMap.Basic
 
 public section
@@ -162,7 +163,10 @@ private meta def queryMap : Std.DHashMap Name (fun _ => Name × Array (MacroM (T
      ⟨`getEntry?, (``getEntry?_eq_getEntry?, #[`(getEntry?_of_perm _)])⟩,
      ⟨`getEntryD, (``getEntryD_eq_getEntryD, #[`(getEntryD_of_perm _)])⟩,
      ⟨`getEntry!, (``getEntry!_eq_getEntry!, #[`(getEntry!_of_perm _)])⟩,
+     ⟨`all, (``Raw.all_eq_all_toListModel, #[])⟩,
      ⟨`toList, (``Raw.toList_eq_toListModel, #[])⟩,
+     ⟨`Const.beq, (``Raw₀.Const.beq_eq_beqModel, #[])⟩,
+     ⟨`beq, (``beq_eq_beqModel, #[])⟩,
      ⟨`keys, (``Raw.keys_eq_keys_toListModel, #[`(perm_keys_congr_left)])⟩,
      ⟨`Const.toList, (``Raw.Const.toList_eq_toListModel_map, #[`(perm_map_congr_left)])⟩,
      ⟨`foldM, (``Raw.foldM_eq_foldlM_toListModel, #[])⟩,
@@ -1453,16 +1457,7 @@ theorem any_eq_false [LawfulBEq α] {p : (a : α) → β a → Bool} (h : m.1.WF
 omit [Hashable α] [BEq α] in
 theorem all_toList {p : (a : α) → β a → Bool} :
     m.1.toList.all (fun x => p x.1 x.2) = m.1.all p := by
-  simp only [Raw.all, ForIn.forIn, Bool.not_eq_true, bind_pure_comp, map_pure, Id.run_bind]
-  rw [forIn_eq_forIn_toList, forIn_eq_forIn']
-  induction m.val.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]
+  simp_to_model
 
 omit [Hashable α] [BEq α] in
 theorem all_eq_not_any_not {p : (a : α) → β a → Bool} :
@@ -2590,8 +2585,37 @@ end Const
 
 end insertMany
 
-section Union
+section BEq
+variable {m₁ m₂ : Raw₀ α β} [LawfulBEq α] [∀ k, BEq (β k)]
 
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) : m₁.1.Equiv m₂.1 → beq m₁ m₂ := by
+  simp_to_model using List.beqModel_eq_true_of_perm
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) : beq m₁ m₂ = true → m₁.1.Equiv m₂.1 := by
+  simp_to_model using List.perm_of_beqModel
+
+theorem Equiv.beq_congr {m₃ m₄ : Raw₀ α β} (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) (h₃ : m₃.val.WF) (h₄ : m₄.val.WF) :
+    m₁.1.Equiv m₃.1 → m₂.1.Equiv m₄.1 → Raw₀.beq m₁ m₂ = Raw₀.beq m₃ m₄ := by
+  simp_to_model using List.beqModel_congr
+
+end BEq
+
+section
+
+variable {β : Type v} {m₁ m₂ : Raw₀ α (fun _ => β)}
+
+theorem Const.Equiv.beq [LawfulHashable α] [EquivBEq α] [BEq β] [ReflBEq β] (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) : m₁.1.Equiv m₂.1 → Const.beq m₁ m₂ := by
+  simp_to_model using List.Const.beqModel_eq_true_of_perm
+
+theorem Const.equiv_of_beq [LawfulBEq α] [BEq β] [LawfulBEq β] (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) : beq m₁ m₂ = true → m₁.1.Equiv m₂.1 := by
+  simp_to_model using List.Const.perm_of_beqModel
+
+theorem Const.Equiv.beq_congr [LawfulBEq α] {m₃ m₄ : Raw₀ α (fun _ => β)} [BEq β] (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) (h₃ : m₃.val.WF) (h₄ : m₄.val.WF) :
+    m₁.1.Equiv m₃.1 → m₂.1.Equiv m₄.1 → Const.beq m₁ m₂ = Const.beq m₃ m₄ := by
+  simp_to_model using List.Const.beqModel_congr
+end
+
+section Union
 variable (m₁ m₂ : Raw₀ α β)
 
 variable {m₁ m₂}

src/Std/Data/DHashMap/Internal/WF.lean

@@ -275,6 +275,19 @@ theorem forIn_eq_forIn_toListModel {δ : Type w} {l : Raw α β} {m : Type w →
       · simp
       · simpa using ih'
 
+theorem all_eq_all_toListModel {p : (a: α) → β a → Bool} {m : Raw α β} :
+    m.all p = (toListModel m.buckets).all (fun x => p x.1 x.2) := by
+  simp only [Raw.all, ForIn.forIn, Bool.not_eq_true, 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 only [all_nil, forIn'_nil, Id.run_pure]
+  | 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]
+
 end Raw
 
 namespace Raw₀
@@ -1584,4 +1597,13 @@ theorem Const.wf_insertManyIfNewUnit₀ [BEq α] [Hashable α] [EquivBEq α] [La
     {l : ρ} (h' : m.WF) : (Const.insertManyIfNewUnit ⟨m, h⟩ l).1.1.WF :=
   (Raw₀.Const.insertManyIfNewUnit ⟨m, h⟩ l).2 _ Raw.WF.insertIfNew₀ h'
 
+theorem beq_eq_beqModel [BEq α] [LawfulBEq α] [Hashable α] [∀ k, BEq (β k)] {m₁ m₂ : Raw₀ α β}  (h₁ : Raw.WFImp m₁.1) (h₂ : Raw.WFImp m₂.1) :
+    beq m₁ m₂ = beqModel (toListModel m₁.1.buckets) (toListModel m₂.1.buckets) := by
+  simp [beq, beqModel, Raw.size_eq_length h₁, Raw.size_eq_length h₂, Raw.all_eq_all_toListModel,
+    get?_eq_getValueCast? h₂]
+
+theorem Const.beq_eq_beqModel {β : Type v} [BEq α] [PartialEquivBEq α] [Hashable α] [LawfulHashable α] [BEq β] {m₁ m₂ : Raw₀ α (fun _ => β)} (h₁ : Raw.WFImp m₁.1) (h₂ : Raw.WFImp m₂.1) :
+    beq m₁ m₂ = Const.beqModel (toListModel m₁.1.buckets) (toListModel m₂.1.buckets) := by
+  simp [beq, Const.beqModel, Raw.size_eq_length h₁, Raw.size_eq_length h₂, Raw.all_eq_all_toListModel, get?_eq_getValue? h₂]
+
 end Raw₀

src/Std/Data/DHashMap/Lemmas.lean

@@ -1838,6 +1838,41 @@ theorem isEmpty_of_isEmpty_insertMany [EquivBEq α] [LawfulHashable α]
     {l : ρ} : (m.insertMany l).isEmpty → m.isEmpty :=
   Raw₀.isEmpty_of_isEmpty_insertMany ⟨m.1, _⟩ m.2
 
+section BEq
+variable {m₁ m₂ : DHashMap α β} [LawfulBEq α] [∀ k, BEq (β k)]
+
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h : m₁ ~m m₂) : m₁ == m₂ := by
+  simp [BEq.beq]
+  apply Raw₀.Equiv.beq m₁.2 m₂.2 h.1
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h : m₁ == m₂) : m₁ ~m m₂ := by
+  constructor
+  have := @Raw₀.equiv_of_beq α β _ _ ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩ _ _ _ m₁.2 m₂.2 h
+  simp only at this
+  exact this
+
+theorem Equiv.beq_congr {m₃ m₄ : DHashMap α β} : m₁ ~m m₃ → m₂ ~m m₄ → (m₁ == m₂) = (m₃ == m₄) := by
+  intro h1 h2
+  exact @Raw₀.Equiv.beq_congr α β _ _ ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩ _ _ ⟨m₃.1, m₃.2.size_buckets_pos⟩ ⟨m₄.1, m₄.2.size_buckets_pos⟩ m₁.2 m₂.2 m₃.2 m₄.2 h1.1 h2.1
+end BEq
+
+section
+variable {β : Type v} {m₁ m₂ : DHashMap α (fun _ => β)} [BEq β]
+
+theorem Const.Equiv.beq [LawfulHashable α] [EquivBEq α] [ReflBEq β] (h : m₁ ~m m₂) : DHashMap.Const.beq m₁ m₂ := by
+  apply Raw₀.Const.Equiv.beq m₁.2 m₂.2 h.1
+
+theorem Const.equiv_of_beq [LawfulBEq α] [LawfulBEq β] (h : Const.beq m₁ m₂) : m₁ ~m m₂ := by
+  constructor
+  have := @Raw₀.Const.equiv_of_beq α _ _ _ ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩ _ _ _ m₁.2 m₂.2 h
+  simp only at this
+  exact this
+
+theorem Const.Equiv.beq_congr [LawfulBEq α] {m₃ m₄ : DHashMap α (fun _ => β)} : m₁ ~m m₃ → m₂ ~m m₄ → Const.beq m₁ m₂ = Const.beq m₃ m₄ := by
+  intro h1 h2
+  exact @Raw₀.Const.Equiv.beq_congr α _ _ _ ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩ _ ⟨m₃.1, m₃.2.size_buckets_pos⟩ ⟨m₄.1, m₄.2.size_buckets_pos⟩ _ m₁.2 m₂.2 m₃.2 m₄.2 h1.1 h2.1
+end
+
 section Union
 
 variable (m₁ m₂ : DHashMap α β)

src/Std/Data/DHashMap/Raw.lean

@@ -476,17 +476,6 @@ only those mappings where the function returns `some` value.
 @[inline] def keysArray (m : Raw α β) : Array α :=
   m.fold (fun acc k _ => acc.push k) (.emptyWithCapacity m.size)
 
-/-- Checks if all elements satisfy the predicate, short-circuiting if a predicate fails. -/
-@[inline] def all (m : Raw α β) (p : (a : α) → β a → Bool) : Bool := Id.run do
-  for a in m do
-    if ¬ p a.1 a.2 then return false
-  return true
-
-/-- Checks if any element satisfies the predicate, short-circuiting if a predicate succeeds. -/
-@[inline] def any (m : Raw α β) (p : (a : α) → β a → Bool) : Bool := Id.run do
-  for a in m do
-    if p a.1 a.2 then return true
-  return false
 /--
 Computes the union of the given hash maps. If a key appears in both maps, the entry contained in
 the second argument will appear in the result.
@@ -522,6 +511,28 @@ This function always merges the smaller map into the larger map, so the expected
 
 instance [BEq α] [Hashable α] : Inter (Raw α β) := ⟨inter⟩
 
+/-- Internal implementation detail of the hash map. -/
+def beq [BEq α] [Hashable α] [LawfulBEq α] [∀ k, BEq (β k)] (m₁ m₂ : Raw α β) : Bool :=
+    if h₁ : 0 < m₁.buckets.size then
+      if h₂ : 0 < m₂.buckets.size then
+        Raw₀.beq ⟨m₁, h₁⟩ ⟨m₂, h₂⟩
+      else
+        false
+    else
+      false
+
+instance [BEq α] [Hashable α] [LawfulBEq α] [∀ k, BEq (β k)] : BEq (Raw α β) := ⟨beq⟩
+
+/-- Internal implementation detail of the hash map. -/
+def Const.beq {β : Type v} [BEq α] [Hashable α] [BEq β] (m₁ m₂ : Raw α (fun _ => β)) : Bool :=
+   if h₁ : 0 < m₁.buckets.size then
+      if h₂ : 0 < m₂.buckets.size then
+        Raw₀.Const.beq ⟨m₁, h₁⟩ ⟨m₂, h₂⟩
+      else
+        false
+    else
+      false
+
 /--
 Computes the difference of the given hash maps.
 

src/Std/Data/DHashMap/RawDef.lean

@@ -81,6 +81,18 @@ instance [Monad m] : ForM m (Raw α β) ((a : α) × β a) where
 instance [Monad m] : ForIn m (Raw α β) ((a : α) × β a) where
   forIn m init f := m.forIn (fun a b acc => f ⟨a, b⟩ acc) init
 
+/-- Checks if all elements satisfy the predicate, short-circuiting if a predicate fails. -/
+@[inline] def all (m : Raw α β) (p : (a : α) → β a → Bool) : Bool := Id.run do
+  for a in m do
+    if ¬ p a.1 a.2 then return false
+  return true
+
+/-- Checks if any element satisfies the predicate, short-circuiting if a predicate succeeds. -/
+@[inline] def any (m : Raw α β) (p : (a : α) → β a → Bool) : Bool := Id.run do
+  for a in m do
+    if p a.1 a.2 then return true
+  return false
+
 end Raw
 
 end Std.DHashMap

src/Std/Data/DHashMap/RawLemmas.lean

@@ -9,7 +9,6 @@ prelude
 public import Std.Data.DHashMap.Internal.Raw
 public import Std.Data.DHashMap.Internal.RawLemmas
 import all Std.Data.DHashMap.Raw
-
 public section
 
 /-!
@@ -48,7 +47,7 @@ private meta def baseNames : Array Name :=
     ``getKey?_eq, ``getKey_eq, ``getKey!_eq, ``getKeyD_eq,
     ``insertMany_eq, ``Const.insertMany_eq, ``Const.insertManyIfNewUnit_eq, ``ofArray_eq, ``Const.ofArray_eq,
     ``ofList_eq, ``Const.ofList_eq, ``Const.unitOfList_eq, ``Const.unitOfArray_eq,
-    ``alter_eq, ``Const.alter_eq, ``modify_eq, ``Const.modify_eq, ``union_eq, ``inter_eq, ``diff_eq]
+    ``alter_eq, ``Const.alter_eq, ``modify_eq, ``Const.modify_eq, ``union_eq, ``inter_eq, ``diff_eq, ``beq_eq, ``Const.beq_eq]
 
 
 /-- Internal implementation detail of the hash map -/
@@ -3759,6 +3758,45 @@ variable [BEq α] [Hashable α]
 
 open Internal.Raw Internal.Raw₀
 
+section BEq
+variable {m₁ m₂ : Raw α β} [LawfulBEq α] [∀ k, BEq (β k)]
+
+theorem Equiv.beq [∀ k, ReflBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) : m₁ ~m m₂ → m₁ == m₂ := by
+  simp only [BEq.beq]
+  simp_to_raw
+  intro h
+  exact Raw₀.Equiv.beq h₁ h₂ h
+
+theorem equiv_of_beq [∀ k, LawfulBEq (β k)] (h₁ : m₁.WF) (h₂ : m₂.WF) : beq m₁ m₂ = true → m₁ ~m m₂ := by
+  simp_to_raw using Raw₀.equiv_of_beq
+
+theorem Equiv.beq_congr {m₃ m₄ : Raw α β} (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₄ := by
+  simp_to_raw
+  intro w1 w2
+  exact Raw₀.Equiv.beq_congr h₁ h₂ h₃ h₄ w1 w2
+
+end BEq
+
+section
+variable {β : Type v} {m₁ m₂ : Raw α (fun _ => β)}
+
+theorem Const.Equiv.beq [LawfulHashable α] [EquivBEq α] [BEq β] [ReflBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : m₁ ~m m₂ → beq m₁ m₂ := by
+  simp_to_raw
+  intro h
+  exact Raw₀.Const.Equiv.beq h₁ h₂ h
+
+theorem Const.equiv_of_beq [LawfulBEq α] [BEq β] [LawfulBEq β] (h₁ : m₁.WF) (h₂ : m₂.WF) : beq m₁ m₂ = true → m₁ ~m m₂ := by
+  simp_to_raw using Raw₀.Const.equiv_of_beq
+
+theorem Const.Equiv.beq_congr [LawfulBEq α] [BEq β] {m₃ m₄ : Raw  α (fun _ => β)} (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₄ := by
+  simp_to_raw
+  intro w1 w2
+  exact Raw₀.Const.Equiv.beq_congr h₁ h₂ h₃ h₄ w1 w2
+
+end
+
 @[simp, grind =]
 theorem ofArray_eq_ofList (a : Array ((a : α) × (β a))) :
     ofArray a = ofList a.toList := by