feat: add intersection on DHashMap (#11112)

This PR adds intersection operation on `DHashMap`/`HashMap`/`HashSet`
and provides several lemmas about its behaviour.

---------

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

Author: wkrozowski

src/Std/Data/DHashMap/Basic.lean

@@ -336,7 +336,18 @@ This function always merges the smaller map into the larger map, so the expected
   inner := Raw₀.union ⟨m₁.1, m₁.2.size_buckets_pos⟩ ⟨m₂.1, m₂.2.size_buckets_pos⟩
   wf := Std.DHashMap.Raw.WF.union₀ m₁.2 m₂.2
 
+/--
+Computes the intersection of the given hash maps. The result will only contain entries from the first map.
+
+This function always iterates through the smaller map, so the expected runtime is
+`O(min(m₁.size, m₂.size))`.
+-/
+@[inline] def inter [BEq α] [Hashable α] (m₁ m₂ : DHashMap α β) : DHashMap α β where
+  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
+
 instance [BEq α] [Hashable α] : Union (DHashMap α β) := ⟨union⟩
+instance [BEq α] [Hashable α] : Inter (DHashMap α β) := ⟨inter⟩
 
 section Unverified
 

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

@@ -461,10 +461,25 @@ def insertMany {ρ : Type w} [ForIn Id ρ ((a : α) × β a)] [BEq α] [Hashable
     r := ⟨r.1.insertIfNew a b, fun _ h hm => h (r.2 _ h hm)⟩
   return r
 
+/-- Internal implementation detail of the hash map -/
+@[inline]
+def interSmallerFn [BEq α] [Hashable α] (m sofar : Raw₀ α β) (k : α) : Raw₀ α β :=
+  match m.getEntry? k with
+  | some kv' => sofar.insert kv'.1 kv'.2
+  | none => sofar
+
+/-- Internal implementation detail of the hash map -/
+def interSmaller [BEq α] [Hashable α] (m₁ : Raw₀ α β) (m₂ : Raw α β) : Raw₀ α β :=
+  (m₂.fold (fun sofar k _ => interSmallerFn m₁ sofar k) emptyWithCapacity)
+
 /-- Internal implementation detail of the hash map -/
 @[inline] def union [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
   if m₁.1.size ≤ m₂.1.size then (m₂.insertManyIfNew m₁.1).1 else (m₁.insertMany m₂.1).1
 
+/-- Internal implementation detail of the hash map -/
+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₂
+
 section
 
 variable {β : Type v}

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

@@ -420,6 +420,12 @@ def unionₘ [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
   else
     insertListₘ m₁ (toListModel m₂.1.buckets)
 
+/-- Internal implementation detail of the hash map -/
+def interSmallerFnₘ [BEq α] [Hashable α] (m sofar : Raw₀ α β) (k : α) : Raw₀ α β :=
+  match m.getEntry?ₘ k with
+  | some kv' => sofar.insertₘ kv'.1 kv'.2
+  | none => sofar
+
 section
 
 variable {β : Type v}
@@ -664,6 +670,14 @@ theorem insertManyIfNew_eq_insertListIfNewₘ [BEq α] [Hashable α] (m : Raw₀
     simp only [List.foldl_cons, insertListIfNewₘ]
     apply ih
 
+theorem interSmallerFn_eq_interSmallerFnₘ [BEq α] [Hashable α] (m sofar : Raw₀ α β) (k : α) :
+    interSmallerFn m sofar k = interSmallerFnₘ m sofar k := by
+  rw [interSmallerFn, interSmallerFnₘ]
+  rw [getEntry?_eq_getEntry?ₘ]
+  congr
+  ext
+  rw [insert_eq_insertₘ]
+
 section
 
 variable {β : Type v}

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

@@ -145,6 +145,10 @@ theorem union_eq [BEq α] [Hashable α] {m₁ m₂ : Raw α β} (h₁ : m₁.WF)
     m₁.union m₂ = Raw₀.union ⟨m₁, h₁.size_buckets_pos⟩ ⟨m₂, h₂.size_buckets_pos⟩ := by
   simp [Raw.union, h₁.size_buckets_pos, h₂.size_buckets_pos]
 
+theorem inter_eq [BEq α] [Hashable α] {m₁ m₂ : Raw α β} (h₁ : m₁.WF) (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]
+
 section
 
 variable {β : Type v}

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

@@ -167,6 +167,23 @@ theorem foldM_eq_foldlM_toListModel {δ : Type w} {m : Type w → Type w'} [Mona
       funext init'
       rw [ih]
 
+theorem fold_induction {δ : Type w}
+    {f : δ → (a : α) → β a → δ} {init : δ} {b : Raw α β} {P : δ → Prop}
+    (base : P init) (step : ∀ acc a b , P acc → P (f acc a b)) :
+    P (b.fold f init) := by
+  simp [Raw.fold, Raw.foldM, ← Array.foldlM_toList]
+  induction b.buckets.toList generalizing init with
+  | nil => simp [base]
+  | cons hd tl ih =>
+    apply ih
+    induction hd generalizing init with
+    | nil => simp [AssocList.foldlM, pure, base]
+    | cons hda hdb tl ih =>
+      simp only [AssocList.foldlM, pure_bind]
+      apply ih
+      apply step
+      exact base
+
 theorem fold_eq_foldl_toListModel {l : Raw α β} {f : γ → (a : α) → β a → γ} {init : γ} :
     l.fold f init = (toListModel l.buckets).foldl (fun a b => f a b.1 b.2) init := by
   simp [Raw.fold, foldM_eq_foldlM_toListModel]
@@ -1107,7 +1124,6 @@ theorem insertMany_eq_insertListₘ_toListModel [BEq α] [Hashable α] (m m₂ :
     simp only [List.foldl_cons, insertListₘ]
     apply ih
 
-
 theorem insertManyIfNew_eq_insertListIfNewₘ_toListModel [BEq α] [Hashable α] (m m₂ : Raw₀ α β) :
     insertManyIfNew m m₂.1 = insertListIfNewₘ m (toListModel m₂.1.buckets) := by
   simp only [insertManyIfNew, bind_pure_comp, map_pure, bind_pure]
@@ -1142,6 +1158,20 @@ theorem toListModel_unionₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashabl
   · exact toListModel_insertListIfNewₘ ‹_›
   · exact toListModel_insertListₘ ‹_›
 
+theorem wfImp_inter [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α]
+    {m₁ m₂ : Raw α β} {h₁ : 0 < m₁.buckets.size} {h₂ : 0 < m₂.buckets.size} (wh₁ : Raw.WFImp m₁) :
+    Raw.WFImp (Raw₀.inter ⟨m₁, h₁⟩ ⟨m₂, h₂⟩).val := by
+  rw [inter]
+  split
+  · apply wfImp_filter wh₁
+  · rw [interSmaller]
+    apply @Raw.fold_induction _ β _ (fun sofar k x => interSmallerFn ⟨m₁, h₁⟩ sofar k) emptyWithCapacity m₂ (Raw.WFImp ·.val) wfImp_emptyWithCapacity
+    intro acc a b wf
+    rw [interSmallerFn]
+    split
+    · apply wfImp_insert wf
+    · apply wf
+
 end Raw₀
 
 namespace Raw
@@ -1163,6 +1193,7 @@ theorem WF.out [BEq α] [Hashable α] [i₁ : EquivBEq α] [i₂ : LawfulHashabl
   | constModify₀ _ h => exact Raw₀.Const.wfImp_modify (by apply h)
   | alter₀ _ h => exact Raw₀.wfImp_alter (by apply h)
   | constAlter₀ _ h => exact Raw₀.Const.wfImp_alter (by apply h)
+  | inter₀ _ _ h _  => exact Raw₀.wfImp_inter (by apply h)
 
 end Raw
 
@@ -1312,13 +1343,84 @@ theorem wf_union₀ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α]
   · exact wf_insertManyIfNew₀ ‹_›
   · exact wf_insertMany₀ ‹_›
 
+
 theorem toListModel_union [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] {m₁ m₂ : Raw₀ α β}
     (h₁ : Raw.WFImp m₁.1) (h₂ : Raw.WFImp m₂.1) :
     Perm (toListModel (m₁.union m₂).1.buckets)
       (List.insertList (toListModel m₁.1.buckets) (toListModel m₂.1.buckets)) := by
   rw [union_eq_unionₘ]
   exact toListModel_unionₘ h₁ h₂
 
+
+/-! # `inter` -/
+
+theorem wfImp_interSmallerFnₘ [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α] (m₁ : Raw₀ α β) (m₂ : Raw₀ α β)
+    (hm₂ : Raw.WFImp m₂.1) (k : α) : Raw.WFImp (m₁.interSmallerFnₘ m₂ k).1 := by
+  rw [interSmallerFnₘ]
+  split
+  · exact wfImp_insertₘ hm₂
+  · exact hm₂
+
+/-- Internal implementation detail of the hash map -/
+def interSmallerₘ [BEq α] [Hashable α] (m₁ : Raw₀ α β) (m₂ : Raw α β) : Raw₀ α β :=
+  m₂.fold (fun sofar k _ => interSmallerFnₘ m₁ sofar k) emptyWithCapacity
+
+theorem interSmaller_eq_interSmallerₘ [BEq α] [Hashable α] (m₁ : Raw₀ α β) (m₂ : Raw α β) :
+    m₁.interSmaller m₂ = m₁.interSmallerₘ m₂ := by
+  rw [interSmaller, interSmallerₘ]
+  simp only [interSmallerFn_eq_interSmallerFnₘ]
+
+theorem foldl_perm_cong  [BEq α] {init₁ init₂ : List ((a : α) × β a)} {l : List ((a : α) × β a)}
+    {f : List ((a : α) × β a) → ((a : α) × β a) → List ((a : α) × β a)} (h₁ : Perm init₁ init₂)
+    (h₂ : ∀ h l₁ l₂, (w : DistinctKeys l₁) → Perm l₁ l₂ → Perm (f l₁ h) (f l₂ h) ∧ DistinctKeys (f l₁ h))
+    (h₃ : DistinctKeys init₁)
+     : Perm (List.foldl f init₁ l) (List.foldl f init₂ l) := by
+  induction l generalizing init₁ init₂
+  case nil =>
+    simp only [foldl_nil, h₁]
+  case cons h t ih =>
+    simp only [foldl_cons]
+    apply ih
+    · exact (h₂ h init₁ init₂ h₃ h₁).1
+    · exact (h₂ h init₁ init₂ h₃ h₁).2
+
+theorem toListModel_interSmallerFnₘ [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α] (m sofar : Raw₀ α β)
+    (l : List ((a : α) × β a))
+    (hm : Raw.WFImp m.1) (hs : Raw.WFImp sofar.1) (k : α) (hml : toListModel sofar.1.buckets ~ l) :
+    Perm (toListModel ((interSmallerFnₘ m sofar k).1.buckets))
+      (List.interSmallerFn (toListModel m.1.buckets) l k) := by
+  rw [interSmallerFnₘ, getEntry?ₘ_eq_getEntry? hm, List.interSmallerFn]
+  split
+  · simpa [*] using (toListModel_insertₘ hs).trans (List.insertEntry_of_perm hs.distinct hml)
+  · simp [*]
+
+theorem toListModel_interSmallerₘ [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α]
+    (m₁ : Raw₀ α β) (m₂ : Raw α β) (hm₁ : Raw.WFImp m₁.1) :
+    toListModel (m₁.interSmallerₘ m₂).1.buckets ~
+      List.interSmaller (toListModel m₁.1.buckets) (toListModel m₂.buckets) := by
+  rw [interSmallerₘ, Raw.fold_eq_foldl_toListModel, List.interSmaller]
+  generalize toListModel m₂.buckets = l
+  suffices ∀ m l', Raw.WFImp m.1 → toListModel m.1.buckets ~ l' → toListModel (foldl (fun a b => m₁.interSmallerFnₘ a b.fst) m l).val.buckets ~
+      foldl (fun sofar kv => List.interSmallerFn (toListModel m₁.val.buckets) sofar kv.fst) l' l by
+    simpa using this emptyWithCapacity [] wfImp_emptyWithCapacity (by simp)
+  intro m l' hm hml'
+  induction l generalizing m l' with
+  | nil => simpa
+  | cons ht tl ih =>
+    rw [List.foldl_cons, List.foldl_cons]
+    exact ih _ _ (wfImp_interSmallerFnₘ _ _ hm _) (toListModel_interSmallerFnₘ _ _ _ hm₁ hm _ hml')
+
+theorem toListModel_inter [BEq α] [EquivBEq α] [Hashable α] [LawfulHashable α] (m₁ m₂ : Raw₀ α β) (hm₁ : Raw.WFImp m₁.1) (hm₂ : Raw.WFImp m₂.1) :
+    Perm (toListModel (m₁.inter m₂).1.buckets) ((toListModel m₁.1.buckets).filter fun p => containsKey p.1 (toListModel m₂.1.buckets) ) := by
+  simp [inter]
+  split
+  · rw [filter_eq_filterₘ]
+    simp only [contains_eq_containsKey hm₂]
+    exact toListModel_filterₘ
+  · rw [interSmaller_eq_interSmallerₘ]
+    exact Perm.trans (toListModel_interSmallerₘ _ _ hm₁)
+      (interSmaller_perm_filter _ _ hm₁.distinct)
+
 /-! # `Const.insertListₘ` -/
 
 theorem Const.toListModel_insertListₘ {β : Type v} [BEq α] [Hashable α] [EquivBEq α]