feat: add difference on DHashMap/HashMap/HashSet (#11212)

This PR adds support for difference operation for
`DHashMap`/`HashMap`/`HashSet` and proves several lemmas about it.

---------

Co-authored-by: Markus Himmel <[email protected]>

Author: wkrozowski

src/Std/Data/DHashMap/Basic.lean

@@ -346,8 +346,19 @@ This function always iterates through the smaller map, so the expected runtime i
   inner := Raw₀.inter ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩
   wf := Std.DHashMap.Raw.WF.inter₀ m₁.2 m₂.2
 
+/--
+Computes the difference of the given hash maps.
+
+This function always iterates through the smaller map, so the expected runtime is
+`O(min(m₁.size, m₂.size))`.
+-/
+@[inline] def diff [BEq α] [Hashable α] (m₁ m₂ : DHashMap α β) : DHashMap α β where
+  inner := Raw₀.diff ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩
+  wf := Std.DHashMap.Raw.WF.diff₀ m₁.2 m₂.2
+
 instance [BEq α] [Hashable α] : Union (DHashMap α β) := ⟨union⟩
 instance [BEq α] [Hashable α] : Inter (DHashMap α β) := ⟨inter⟩
+instance [BEq α] [Hashable α] : SDiff (DHashMap α β) := ⟨diff⟩
 
 section Unverified
 

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

@@ -451,6 +451,16 @@ def insertMany {ρ : Type w} [ForIn Id ρ ((a : α) × β a)] [BEq α] [Hashable
     r := ⟨r.1.insert a b, fun _ h hm => h (r.2 _ h hm)⟩
   return r
 
+/-- Internal implementation detail of the hash map -/
+def eraseManyEntries {ρ : Type w} [ForIn Id ρ ((a : α) × β a)] [BEq α] [Hashable α]
+    (m : Raw₀ α β) (l : ρ) : { m' : Raw₀ α β // ∀ (P : Raw₀ α β → Prop),
+      (∀ {m'' a}, P m'' → P (m''.erase a)) → P m → P m' } := Id.run do
+  let mut r : { m' : Raw₀ α β // ∀ (P : Raw₀ α β → Prop),
+      (∀ {m'' a}, P m'' → P (m''.erase a)) → P m → P m' } := ⟨m, fun _ _ => id⟩
+  for ⟨a, _⟩ in l do
+    r := ⟨r.1.erase a, fun _ h hm => h (r.2 _ h hm)⟩
+  return r
+
 /-- Internal implementation detail of the hash map -/
 @[inline] def insertManyIfNew {ρ : Type w} [ForIn Id ρ ((a : α) × β a)] [BEq α] [Hashable α]
     (m : Raw₀ α β) (l : ρ) : { m' : Raw₀ α β // ∀ (P : Raw₀ α β → Prop),
@@ -480,6 +490,11 @@ 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 -/
+@[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}

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

@@ -407,6 +407,19 @@ def insertListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α)
   | .nil => m
   | .cons hd tl => insertListₘ (m.insert hd.1 hd.2) tl
 
+/-- Internal implementation detail of the hash map -/
+def eraseListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List α) : Raw₀ α β :=
+  match l with
+  | .nil => m
+  | .cons hd tl => eraseListₘ (m.erase hd) tl
+
+/-- Internal implementation detail of the hash map -/
+def diffₘ [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
+  if m₁.1.size ≤ m₂.1.size then
+    filterₘ m₁ (fun k _ => !containsₘ m₂ k)
+  else
+    eraseManyEntries m₁ (toListModel m₂.1.buckets)
+
 /-- Internal implementation detail of the hash map -/
 def insertListIfNewₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α) × β a)) : Raw₀ α β :=
   match l with
@@ -656,6 +669,20 @@ theorem insertMany_eq_insertListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l
     simp only [List.foldl_cons, insertListₘ]
     apply ih
 
+theorem eraseManyEntries_eq_eraseListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α) × β a)) :
+    eraseManyEntries m l = eraseListₘ m (l.map (·.1)) := by
+  simp only [eraseManyEntries, Id.run_pure, pure_bind, List.forIn_pure_yield_eq_foldl]
+  suffices ∀ (t : { m' // ∀ (P : Raw₀ α β → Prop),
+      (∀ {m'' : Raw₀ α β} {a : α}, P m'' → P (m''.erase a)) → P m → P m' }),
+        (List.foldl (fun m' p => ⟨m'.val.erase p.1, fun P h₁ h₂ => h₁ (m'.2 _ h₁ h₂)⟩) t l).val =
+      t.val.eraseListₘ (l.map (·.1)) from this _
+  intro t
+  induction l generalizing m with
+  | nil => simp [eraseListₘ]
+  | cons hd tl ih =>
+    simp only [List.foldl_cons]
+    apply ih
+
 theorem insertManyIfNew_eq_insertListIfNewₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α) × β a)) :
     insertManyIfNew m l = insertListIfNewₘ m l := by
   simp only [insertManyIfNew, Id.run_pure, pure_bind, List.forIn_pure_yield_eq_foldl]

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

@@ -154,6 +154,10 @@ 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 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]
+
 section
 
 variable {β : Type v}