feat: add intersection on DTreeMap/TreeMap/TreeSet (#11165)

This PR provides intersection on `DTreeMap`/`TreeMap`/`TreeSet`and
provides several lemmas about it.

---------

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

Author: wkrozowski

src/Std/Data/DTreeMap/Basic.lean

@@ -1044,6 +1044,15 @@ def union (t₁ t₂ : DTreeMap α β cmp) : DTreeMap α β cmp :=
 
 instance : Union (DTreeMap α β cmp) := ⟨union⟩
 
+/--
+Computes the intersection of the given tree maps. The result will only contain entries from the first map.
+
+This function always merges the smaller map into the larger map.
+-/
+def inter (t₁ t₂ : DTreeMap α β cmp) : DTreeMap α β cmp :=
+  letI : Ord α := ⟨cmp⟩; ⟨t₁.inner.inter t₂.inner t₁.wf.balanced,  @Impl.WF.inter _ _ _ _ t₂.inner t₁.wf.balanced t₁.wf⟩
+
+instance : Inter (DTreeMap α β cmp) := ⟨inter⟩
 /--
 Erases multiple mappings from the tree map by iterating over the given collection and calling
 `erase`.

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

@@ -772,6 +772,26 @@ def filter! [Ord α] (f : (a : α) → β a → Bool) (t : Impl α β) : Impl α
     | false => link2! (filter! f l) (filter! f r)
     | true => link! k v (filter! f l) (filter! f r)
 
+/-- Internal implementation detail of the tree map -/
+@[inline]
+def interSmallerFn [Ord α] (m : Impl α β) (sofar : { t : Impl α β // t.Balanced } ) (k : α) : { res : Impl α β // res.Balanced } :=
+  match m.getEntry? k with
+  | some kv' => let ⟨val, prop, _, _⟩ := (sofar.val.insert kv'.1 kv'.2 sofar.2); ⟨val, prop⟩
+  | none => sofar
+
+/-- Internal implementation detail of the tree map -/
+def interSmaller [Ord α] (m₁ : Impl α β) (m₂ : Impl α β) : Impl α β :=
+  (m₂.foldl (fun sofar k _ => interSmallerFn m₁ sofar k) ⟨empty, balanced_empty⟩).1
+
+/-- Internal implementation detail of the hash map -/
+def inter [Ord α] (m₁ m₂ : Impl α β) (h₁ : Balanced m₁) : Impl α β :=
+  if m₁.size ≤ m₂.size then (m₁.filter (fun k _ => m₂.contains k) h₁).impl else interSmaller m₁ m₂
+
+/-- Slower version of `inter!` which can be used in the absence of balance
+information but still assumes the preconditions of `filter`, otherwise might panic. -/
+def inter! [Ord α] (m₁ m₂ : Impl α β): Impl α β :=
+  if m₁.size ≤ m₂.size then m₁.filter! (fun k _ => m₂.contains k) else interSmaller m₁ m₂
+
 /--
 Changes the mapping of the key `k` by applying the function `f` to the current mapped value
 (if any). This function can be used to insert a new mapping, modify an existing one or delete it.
@@ -848,6 +868,15 @@ theorem balanced_modify [Ord α] [LawfulEqOrd α] {k f} {t : Impl α β} (ht : t
     have ihr := ihr ht.right
     tree_tac
 
+theorem balanced_inter [Ord α] {t₁ t₂ : Impl α β} (ht : t₁.Balanced) : (t₁.inter t₂ ht).Balanced := by
+  rw [inter]
+  split
+  · generalize (filter (fun k x => contains k t₂) t₁ ht) = m
+    exact m.balanced_impl
+  · rw [interSmaller]
+    generalize (foldl (fun sofar k x => t₁.interSmallerFn sofar k) ⟨empty, _⟩ t₂) = m
+    exact m.2
+
 /--
 Returns a map that contains all mappings of `t₁` and `t₂`. In case that both maps contain the
 same key `k` with respect to `cmp`, the provided function is used to determine the new value from

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

@@ -61,6 +61,8 @@ inductive WF [Ord α] : {β : α → Type v} → Impl α β → Prop where
   | mergeWith {t₁ t₂ f h} [LawfulEqOrd α] : WF t₁ → WF (t₁.mergeWith f t₂ h).impl
   /-- `mergeWith` preserves well-formedness. Later shown to be subsumed by `.wf`. -/
   | constMergeWith {t₁ t₂ f h} : WF t₁ → WF (Impl.Const.mergeWith f t₁ t₂ h).impl
+  /-- `inter` preserves well-formedness. Later shown to be subsumed by `.wf`. -/
+  | inter {t₁ t₂ h} : WF t₁ → WF (t₁.inter t₂ h)
 
 /--
 A well-formed tree is balanced. This is needed here already because we need to know that the
@@ -72,6 +74,7 @@ theorem WF.balanced [Ord α] {t : Impl α β} (h : WF t) : t.Balanced := by
   case empty => exact balanced_empty
   case modify ih => exact balanced_modify ih
   case constModify ih => exact Const.balanced_modify ih
+  case inter ih => exact balanced_inter ih
 
 theorem WF.eraseMany [Ord α] {ρ} [ForIn Id ρ α] {t : Impl α β} {l : ρ} {h} (hwf : WF t) :
     WF (t.eraseMany l h).val :=

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

@@ -1529,6 +1529,40 @@ theorem toArray_eq_toArray_map {t : Impl α β} :
 
 end Const
 
+/-!
+### interSmallerFn
+-/
+
+theorem WF.interSmallerFn {_ : Ord α} [TransOrd α] [BEq α] (m₁ : Impl α β) (m₂ : Impl α β)
+    (hm₂ :m₂.WF) (k : α) : (m₁.interSmallerFn ⟨m₂,hm₂.balanced⟩ k).1.WF := by
+  rw [Impl.interSmallerFn]
+  split
+  · exact WF.insert hm₂
+  · exact hm₂
+
+theorem ordered_inter [Ord α] [TransOrd α] {l₁ l₂ : Impl α β} (hlb : l₁.Balanced)
+    (hlo : l₁.Ordered) : (l₁.inter l₂ hlb).Ordered := by
+  rw [inter]
+  split
+  · exact ordered_filter hlo
+  · rw [interSmaller]
+    rw [foldl_eq_foldl]
+    generalize l₂.toListModel = l₂
+    suffices ∀ {start}, start.val.Ordered → (List.foldl (fun acc p => interSmallerFn l₁ acc p.fst) start l₂).val.Ordered by
+      apply this
+      apply ordered_empty
+    intro s swf
+    induction l₂ generalizing s
+    case nil => simp [swf]
+    case cons h t ih =>
+      simp only [List.foldl_cons]
+      apply ih
+      simp only [Impl.interSmallerFn]
+      split
+      · apply ordered_insert
+        · exact swf
+      · exact swf
+
 /-!
 ## Deducing that well-formed trees are ordered
 -/
@@ -1549,6 +1583,7 @@ theorem WF.ordered [Ord α] [TransOrd α] {l : Impl α β} (h : WF l) : l.Ordere
   · exact ordered_filter ‹_›
   · exact ordered_mergeWith ‹_› ‹_›
   · exact Const.ordered_mergeWith ‹_› ‹_›
+  · exact ordered_inter ‹_› ‹_›
 
 /-!
 ## Deducing that additional operations are well-formed
@@ -1995,6 +2030,91 @@ theorem WF.filter! {_ : Ord α} {t : Impl α β} {f : (a : α) → β a → Bool
   rw [← filter_eq_filter! (h := h.balanced)]
   exact h.filter
 
+theorem toListModel_interSmallerFn {_ : Ord α} [TransOrd α] [BEq α] [LawfulBEqOrd α] (m sofar : Impl α β) (h₁ : m.WF) (h₂ : sofar.WF)
+    (l : List ((a : α) × β a))
+    (k : α) (hml : List.Perm (sofar.toListModel) l) :
+    List.Perm (toListModel ((interSmallerFn m ⟨sofar,h₂.balanced⟩ k)).1)
+      (List.interSmallerFn (toListModel m) l k) := by
+  rw [interSmallerFn, List.interSmallerFn]
+  split
+  case h_1 _ val heq =>
+    simp only
+    rw [getEntry?_eq_getEntry?] at heq
+    simp only [heq]
+    apply List.Perm.trans
+    · apply toListModel_insert
+      · exact h₂.ordered
+    · apply insertEntry_of_perm
+      · apply Ordered.distinctKeys h₂.ordered
+      · exact hml
+    · exact h₁.ordered
+  case h_2 heq =>
+    simp only
+    rw [getEntry?_eq_getEntry?] at heq
+    simp only [heq, hml]
+    exact h₁.ordered
+
+/-!
+### interSmaller
+-/
+
+theorem toListModel_interSmaller {_ : Ord α} [TransOrd α] [BEq α] [LawfulBEqOrd α]
+    (m₁ : Impl α β) (m₂ : Impl α β) (hm₁ : m₁.WF) :
+    List.Perm (toListModel (m₁.interSmaller m₂))
+      (List.interSmaller (toListModel m₁) (toListModel m₂)) := by
+  rw [interSmaller, foldl_eq_foldl, List.interSmaller]
+  generalize toListModel m₂ = l
+  suffices ∀ m l', (hm : m.WF) → List.Perm (toListModel m) l' →
+      List.Perm (toListModel (List.foldl (fun a b => interSmallerFn m₁ a b.fst) ⟨m, hm.balanced⟩ l).val)
+        (List.foldl (fun sofar kv => List.interSmallerFn (toListModel m₁) sofar kv.fst) l' l) by
+    simpa using this empty [] WF.empty (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 _ _ (by apply WF.interSmallerFn _ _ hm) (toListModel_interSmallerFn _ _ hm₁ hm _ _ hml')
+
+/-!
+### inter
+-/
+
+theorem toListModel_inter {_ : Ord α} [TransOrd α] [BEq α] [LawfulBEqOrd α]
+    (m₁ : Impl α β) (m₂ : Impl α β) (hm₁ : m₁.WF) (hm₂ : m₂.WF) :
+    List.Perm (toListModel (m₁.inter m₂ hm₁.balanced)) ((toListModel m₁).filter fun p => containsKey p.1 (toListModel m₂)) := by
+  rw [inter]
+  split
+  · rw [toListModel_filter]
+    conv =>
+      lhs
+      lhs
+      ext e
+      rw [@contains_eq_containsKey α β _ _ _ _ e.fst m₂ hm₂.ordered]
+  · apply List.Perm.trans (toListModel_interSmaller _ _ hm₁) (List.interSmaller_perm_filter _ _ hm₁.ordered.distinctKeys)
+
+/-!
+### inter!
+-/
+
+theorem inter_eq_inter! [Ord α] {l₁ l₂: Impl α β} {h} :
+    (inter l₁ l₂ h) = inter! l₁ l₂ := by
+  rw [inter, inter!]
+  split
+  · rw [filter_eq_filter!]
+  · rfl
+
+theorem toListModel_inter! {_ : Ord α} [TransOrd α] [BEq α] [LawfulBEqOrd α]
+    (m₁ : Impl α β) (m₂ : Impl α β) (hm₁ : m₁.WF) (hm₂ : m₂.WF) :
+    List.Perm (toListModel (m₁.inter! m₂)) ((toListModel m₁).filter fun p => containsKey p.1 (toListModel m₂)) := by
+  rw [← @inter_eq_inter! _ _ _ _ _ hm₁.balanced]
+  exact toListModel_inter _ _ hm₁ hm₂
+
+theorem WF.inter! {_ : Ord α} [TransOrd α]
+    {m₁ m₂ : Impl α β} (wh₁ : m₁.WF) :
+    (inter! m₁ m₂).WF := by
+  rw [← @inter_eq_inter! _ _ _ _ _ wh₁.balanced]
+  exact WF.inter wh₁
+
 /-!
 ### map
 -/