From b7d6a4b222ec2bee13c2e629e652df908217665a Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Thu, 3 Oct 2024 16:26:21 +1000
Subject: [PATCH] feat: adding Insert/Singleton/Union instances for
 HashMap/Set.Raw (#5590)

These were missing from https://github.com/leanprover/lean4/pull/5581.
---
 src/Std/Data/DHashMap/Basic.lean     |  8 ++++++--
 src/Std/Data/DHashMap/Lemmas.lean    |  6 ++++++
 src/Std/Data/DHashMap/Raw.lean       | 15 +++++++++++++++
 src/Std/Data/DHashMap/RawLemmas.lean |  6 ++++++
 src/Std/Data/HashMap/Basic.lean      |  3 +++
 src/Std/Data/HashMap/Lemmas.lean     |  6 ++++++
 src/Std/Data/HashMap/Raw.lean        | 16 ++++++++++++++++
 src/Std/Data/HashMap/RawLemmas.lean  |  6 ++++++
 src/Std/Data/HashSet/Basic.lean      |  1 +
 src/Std/Data/HashSet/Lemmas.lean     |  4 ++++
 src/Std/Data/HashSet/Raw.lean        | 20 ++++++++++++++++++++
 src/Std/Data/HashSet/RawLemmas.lean  |  4 ++++
 tests/lean/run/hashmap.lean          |  4 ++++
 13 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/src/Std/Data/DHashMap/Basic.lean b/src/Std/Data/DHashMap/Basic.lean
index 94308f93c8..385d6611d8 100644
--- a/src/Std/Data/DHashMap/Basic.lean
+++ b/src/Std/Data/DHashMap/Basic.lean
@@ -75,8 +75,12 @@ instance [BEq α] [Hashable α] : Inhabited (DHashMap α β) where
     (b : β a) : DHashMap α β :=
   ⟨Raw₀.insert ⟨m.1, m.2.size_buckets_pos⟩ a b, .insert₀ m.2⟩
 
-instance : Singleton (Σ a, β a) (DHashMap α β) := ⟨fun ⟨a, b⟩ => DHashMap.empty.insert a b⟩
-instance : Insert (Σ a, β a) (DHashMap α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩
+instance : Singleton ((a : α) × β a) (DHashMap α β) := ⟨fun ⟨a, b⟩ => DHashMap.empty.insert a b⟩
+
+instance : Insert ((a : α) × β a) (DHashMap α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩
+
+instance : LawfulSingleton ((a : α) × β a) (DHashMap α β) :=
+  ⟨fun _ => rfl⟩
 
 @[inline, inherit_doc Raw.insertIfNew] def insertIfNew (m : DHashMap α β)
     (a : α) (b : β a) : DHashMap α β :=
diff --git a/src/Std/Data/DHashMap/Lemmas.lean b/src/Std/Data/DHashMap/Lemmas.lean
index 0862d87b29..544b5eefe5 100644
--- a/src/Std/Data/DHashMap/Lemmas.lean
+++ b/src/Std/Data/DHashMap/Lemmas.lean
@@ -87,6 +87,12 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] :
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m := by
   simpa [mem_iff_contains] using isEmpty_iff_forall_contains
 
+@[simp] theorem insert_eq_insert {p : (a : α) × β a} : Insert.insert p m = m.insert p.1 p.2 := rfl
+
+@[simp] theorem singleton_eq_insert {p : (a : α) × β a} :
+    Singleton.singleton p = (∅ : DHashMap α β).insert p.1 p.2 :=
+  rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] {k a : α} {v : β k} :
     (m.insert k v).contains a = (k == a || m.contains a) :=
diff --git a/src/Std/Data/DHashMap/Raw.lean b/src/Std/Data/DHashMap/Raw.lean
index 34a04564ce..1c01c2a4cd 100644
--- a/src/Std/Data/DHashMap/Raw.lean
+++ b/src/Std/Data/DHashMap/Raw.lean
@@ -65,6 +65,15 @@ Inserts the given mapping into the map, replacing an existing mapping for the ke
     (Raw₀.insert ⟨m, h⟩ a b).1
   else m -- will never happen for well-formed inputs
 
+instance [BEq α] [Hashable α] : Singleton ((a : α) × β a) (Raw α β) :=
+  ⟨fun ⟨a, b⟩ => Raw.empty.insert a b⟩
+
+instance [BEq α] [Hashable α] : Insert ((a : α) × β a) (Raw α β) :=
+  ⟨fun ⟨a, b⟩ s => s.insert a b⟩
+
+instance [BEq α] [Hashable α] : LawfulSingleton ((a : α) × β a) (Raw α β) :=
+  ⟨fun _ => rfl⟩
+
 /--
 If there is no mapping for the given key, inserts the given mapping into the map. Otherwise,
 returns the map unaltered.
@@ -399,6 +408,12 @@ occurrence takes precedence. -/
 @[inline] def ofList [BEq α] [Hashable α] (l : List ((a : α) × β a)) : Raw α β :=
   insertMany ∅ l
 
+/-- Computes the union of the given hash maps, by traversing `m₂` and inserting its elements into `m₁`. -/
+@[inline] def union [BEq α] [Hashable α] (m₁ m₂ : Raw α β) : Raw α β :=
+  m₂.fold (init := m₁) fun acc x => acc.insert x
+
+instance [BEq α] [Hashable α] : Union (Raw α β) := ⟨union⟩
+
 @[inline, inherit_doc Raw.ofList] def Const.ofList {β : Type v} [BEq α] [Hashable α]
     (l : List (α × β)) : Raw α (fun _ => β) :=
   Const.insertMany ∅ l
diff --git a/src/Std/Data/DHashMap/RawLemmas.lean b/src/Std/Data/DHashMap/RawLemmas.lean
index b9963ba6a1..5813e5a865 100644
--- a/src/Std/Data/DHashMap/RawLemmas.lean
+++ b/src/Std/Data/DHashMap/RawLemmas.lean
@@ -153,6 +153,12 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] (h : m.WF)
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m := by
   simpa [mem_iff_contains] using isEmpty_iff_forall_contains h
 
+@[simp] theorem insert_eq_insert {p : (a : α) × β a} : Insert.insert p m = m.insert p.1 p.2 := rfl
+
+@[simp] theorem singleton_eq_insert {p : (a : α) × β a} :
+    Singleton.singleton p = (∅ : Raw α β).insert p.1 p.2 :=
+  rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] (h : m.WF) {a k : α} {v : β k} :
     (m.insert k v).contains a = (k == a || m.contains a) := by
diff --git a/src/Std/Data/HashMap/Basic.lean b/src/Std/Data/HashMap/Basic.lean
index d676bdd696..4a9afb3335 100644
--- a/src/Std/Data/HashMap/Basic.lean
+++ b/src/Std/Data/HashMap/Basic.lean
@@ -77,8 +77,11 @@ instance [BEq α] [Hashable α] : Inhabited (HashMap α β) where
   ⟨m.inner.insert a b⟩
 
 instance : Singleton (α × β) (HashMap α β) := ⟨fun ⟨a, b⟩ => HashMap.empty.insert a b⟩
+
 instance : Insert (α × β) (HashMap α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩
 
+instance : LawfulSingleton (α × β) (HashMap α β) := ⟨fun _ => rfl⟩
+
 @[inline, inherit_doc DHashMap.insertIfNew] def insertIfNew (m : HashMap α β)
     (a : α) (b : β) : HashMap α β :=
   ⟨m.inner.insertIfNew a b⟩
diff --git a/src/Std/Data/HashMap/Lemmas.lean b/src/Std/Data/HashMap/Lemmas.lean
index b14e18b4f7..e116921d22 100644
--- a/src/Std/Data/HashMap/Lemmas.lean
+++ b/src/Std/Data/HashMap/Lemmas.lean
@@ -95,6 +95,12 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] :
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m :=
   DHashMap.isEmpty_iff_forall_not_mem
 
+@[simp] theorem insert_eq_insert {p : α × β} : Insert.insert p m = m.insert p.1 p.2 := rfl
+
+@[simp] theorem singleton_eq_insert {p : α × β} :
+    Singleton.singleton p = (∅ : HashMap α β).insert p.1 p.2 :=
+  rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] {k a : α} {v : β} :
     (m.insert k v).contains a = (k == a || m.contains a) :=
diff --git a/src/Std/Data/HashMap/Raw.lean b/src/Std/Data/HashMap/Raw.lean
index afaf810041..825ce0a314 100644
--- a/src/Std/Data/HashMap/Raw.lean
+++ b/src/Std/Data/HashMap/Raw.lean
@@ -74,6 +74,12 @@ set_option linter.unusedVariables false in
     (a : α) (b : β) : Raw α β :=
   ⟨m.inner.insert a b⟩
 
+instance [BEq α] [Hashable α] : Singleton (α × β) (Raw α β) := ⟨fun ⟨a, b⟩ => Raw.empty.insert a b⟩
+
+instance [BEq α] [Hashable α] : Insert (α × β) (Raw α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩
+
+instance [BEq α] [Hashable α] : LawfulSingleton (α × β) (Raw α β) := ⟨fun _ => rfl⟩
+
 @[inline, inherit_doc DHashMap.Raw.insertIfNew] def insertIfNew [BEq α] [Hashable α] (m : Raw α β)
     (a : α) (b : β) : Raw α β :=
   ⟨m.inner.insertIfNew a b⟩
@@ -231,10 +237,20 @@ m.inner.values
     (l : List (α × β)) : Raw α β :=
   ⟨DHashMap.Raw.Const.ofList l⟩
 
+/-- Computes the union of the given hash maps, by traversing `m₂` and inserting its elements into `m₁`. -/
+@[inline] def union [BEq α] [Hashable α] (m₁ m₂ : Raw α β) : Raw α β :=
+  m₂.fold (init := m₁) fun acc x => acc.insert x
+
+instance [BEq α] [Hashable α] : Union (Raw α β) := ⟨union⟩
+
 @[inline, inherit_doc DHashMap.Raw.Const.unitOfList] def unitOfList [BEq α] [Hashable α]
     (l : List α) : Raw α Unit :=
   ⟨DHashMap.Raw.Const.unitOfList l⟩
 
+@[inline, inherit_doc DHashMap.Raw.Const.unitOfArray] def unitOfArray [BEq α] [Hashable α]
+    (l : Array α) : Raw α Unit :=
+  ⟨DHashMap.Raw.Const.unitOfArray l⟩
+
 @[inherit_doc DHashMap.Raw.Internal.numBuckets] def Internal.numBuckets (m : Raw α β) : Nat :=
   DHashMap.Raw.Internal.numBuckets m.inner
 
diff --git a/src/Std/Data/HashMap/RawLemmas.lean b/src/Std/Data/HashMap/RawLemmas.lean
index 2bd7c89c08..6b08495191 100644
--- a/src/Std/Data/HashMap/RawLemmas.lean
+++ b/src/Std/Data/HashMap/RawLemmas.lean
@@ -108,6 +108,12 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] (h : m.WF)
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m :=
   DHashMap.Raw.isEmpty_iff_forall_not_mem h.out
 
+@[simp] theorem insert_eq_insert {p : α × β} : Insert.insert p m = m.insert p.1 p.2 := rfl
+
+@[simp] theorem singleton_eq_insert {p : α × β} :
+    Singleton.singleton p = (∅ : Raw α β).insert p.1 p.2 :=
+  rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] (h : m.WF) {k a : α} {v : β} :
     (m.insert k v).contains a = (k == a || m.contains a) :=
diff --git a/src/Std/Data/HashSet/Basic.lean b/src/Std/Data/HashSet/Basic.lean
index dedad90724..762ea2f113 100644
--- a/src/Std/Data/HashSet/Basic.lean
+++ b/src/Std/Data/HashSet/Basic.lean
@@ -78,6 +78,7 @@ equal (with regard to `==`) to the given element, then the hash set is returned
   ⟨m.inner.insertIfNew a ()⟩
 
 instance : Singleton α (HashSet α) := ⟨fun a => HashSet.empty.insert a⟩
+
 instance : Insert α (HashSet α) := ⟨fun a s => s.insert a⟩
 
 /--
diff --git a/src/Std/Data/HashSet/Lemmas.lean b/src/Std/Data/HashSet/Lemmas.lean
index e4b256a94f..972990cbef 100644
--- a/src/Std/Data/HashSet/Lemmas.lean
+++ b/src/Std/Data/HashSet/Lemmas.lean
@@ -89,6 +89,10 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] :
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m :=
   HashMap.isEmpty_iff_forall_not_mem
 
+@[simp] theorem insert_eq_insert {a : α} : Insert.insert a m = m.insert a := rfl
+
+@[simp] theorem singleton_eq_insert {a : α} : Singleton.singleton a = (∅ : HashSet α).insert a := rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] {k a : α} :
     (m.insert k).contains a = (k == a || m.contains a) :=
diff --git a/src/Std/Data/HashSet/Raw.lean b/src/Std/Data/HashSet/Raw.lean
index 83f2cf30ee..6becfc534f 100644
--- a/src/Std/Data/HashSet/Raw.lean
+++ b/src/Std/Data/HashSet/Raw.lean
@@ -78,6 +78,12 @@ equal (with regard to `==`) to the given element, then the hash set is returned
 @[inline] def insert [BEq α] [Hashable α] (m : Raw α) (a : α) : Raw α :=
   ⟨m.inner.insertIfNew a ()⟩
 
+instance [BEq α] [Hashable α] : Singleton α (Raw α) := ⟨fun a => Raw.empty.insert a⟩
+
+instance [BEq α] [Hashable α] : Insert α (Raw α) := ⟨fun a s => s.insert a⟩
+
+instance [BEq α] [Hashable α] : LawfulSingleton α (Raw α) := ⟨fun _ => rfl⟩
+
 /--
 Checks whether an element is present in a set and inserts the element if it was not found.
 If the hash set already contains an element that is equal (with regard to `==`) to the given
@@ -225,6 +231,20 @@ in the collection will be present in the returned hash set.
 @[inline] def ofList [BEq α] [Hashable α] (l : List α) : Raw α :=
   ⟨HashMap.Raw.unitOfList l⟩
 
+/--
+Creates a hash set from an array of elements. Note that unlike repeatedly calling `insert`, if the
+collection contains multiple elements that are equal (with regard to `==`), then the last element
+in the collection will be present in the returned hash set.
+-/
+@[inline] def ofArray [BEq α] [Hashable α] (l : Array α) : Raw α :=
+  ⟨HashMap.Raw.unitOfArray l⟩
+
+/-- Computes the union of the given hash sets, by traversing `m₂` and inserting its elements into `m₁`. -/
+@[inline] def union [BEq α] [Hashable α] (m₁ m₂ : Raw α) : Raw α :=
+  m₂.fold (init := m₁) fun acc x => acc.insert x
+
+instance [BEq α] [Hashable α] : Union (Raw α) := ⟨union⟩
+
 /--
 Returns the number of buckets in the internal representation of the hash set. This function may
 be useful for things like monitoring system health, but it should be considered an internal
diff --git a/src/Std/Data/HashSet/RawLemmas.lean b/src/Std/Data/HashSet/RawLemmas.lean
index 82b251158e..43edb645fa 100644
--- a/src/Std/Data/HashSet/RawLemmas.lean
+++ b/src/Std/Data/HashSet/RawLemmas.lean
@@ -104,6 +104,10 @@ theorem isEmpty_iff_forall_not_mem [EquivBEq α] [LawfulHashable α] (h : m.WF)
     m.isEmpty = true ↔ ∀ a, ¬a ∈ m :=
   HashMap.Raw.isEmpty_iff_forall_not_mem h.out
 
+@[simp] theorem insert_eq_insert {a : α} : Insert.insert a m = m.insert a := rfl
+
+@[simp] theorem singleton_eq_insert {a : α} : Singleton.singleton a = (∅ : Raw α).insert a := rfl
+
 @[simp]
 theorem contains_insert [EquivBEq α] [LawfulHashable α] (h : m.WF) {k a : α} :
     (m.insert k).contains a = (k == a || m.contains a) :=
diff --git a/tests/lean/run/hashmap.lean b/tests/lean/run/hashmap.lean
index 54fea48e57..1a7db220bb 100644
--- a/tests/lean/run/hashmap.lean
+++ b/tests/lean/run/hashmap.lean
@@ -398,6 +398,10 @@ def addKeyToState (k : Nat) : StateM Nat PUnit := do
 #guard_msgs in
 #eval m.toArray
 
+/-- info: Std.HashSet.Raw.ofList [1000000000, 2, 1, 16] -/
+#guard_msgs in
+#eval m ∪ {16, 16}
+
 /-- info: [1000000000, 2, 1, 16] -/
 #guard_msgs in
 #eval (m.insertMany [16, 16]).toList