Define union and provide basic proof infrastructure

Author: TwoFX

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

@@ -422,6 +422,10 @@ where
     r := ⟨r.1.insert a b, fun _ h hm => h (r.2 _ h hm)⟩
   return r
 
+/-- 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₂.insertMany m₁.1).1 else (m₁.insertMany m₂.1).1
+
 section
 
 variable {β : Type v}

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

@@ -390,6 +390,13 @@ 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 unionₘ [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) : Raw₀ α β :=
+  if m₁.1.size ≤ m₂.1.size then
+    insertListₘ m₂ (toListModel m₁.1.buckets)
+  else
+    insertListₘ m₁ (toListModel m₂.1.buckets)
+
 section
 
 variable {β : Type v}
@@ -595,7 +602,8 @@ theorem map_eq_mapₘ (m : Raw₀ α β) (f : (a : α) → β a → δ a) :
 theorem filter_eq_filterₘ (m : Raw₀ α β) (f : (a : α) → β a → Bool) :
     m.filter f = m.filterₘ f := rfl
 
-theorem insertMany_eq_insertListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α) × β a)) : insertMany m l = insertListₘ m l := by
+theorem insertMany_eq_insertListₘ [BEq α] [Hashable α] (m : Raw₀ α β) (l : List ((a : α) × β a)) :
+    insertMany m l = insertListₘ m l := by
   simp only [insertMany, Id.run_pure, Id.run_bind, pure_bind, List.forIn_pure_yield_eq_foldl]
   suffices ∀ (t : { m' // ∀ (P : Raw₀ α β → Prop),
     (∀ {m'' : Raw₀ α β} {a : α} {b : β a}, P m'' → P (m''.insert a b)) → P m → P m' }),

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

@@ -96,7 +96,7 @@ macro_rules
       | apply Raw.WF.alter₀ | apply Raw.WF.modify₀
       | apply Raw.WF.constAlter₀ | apply Raw.WF.constModify₀
       | apply Raw₀.wf_insertMany₀ | apply Raw₀.Const.wf_insertMany₀
-      | apply Raw₀.Const.wf_insertManyIfNewUnit₀
+      | apply Raw₀.Const.wf_insertManyIfNewUnit₀ | apply Raw₀.wf_union₀
       | apply Raw.WF.filter₀ | apply Raw₀.wf_map₀ | apply Raw₀.wf_filterMap₀
       | apply Raw.WF.emptyWithCapacity₀) <;> wf_trivial)
 
@@ -111,6 +111,7 @@ private def modifyMap : Std.DHashMap Name (fun _ => Name) :=
      ⟨`erase, ``toListModel_erase⟩,
      ⟨`insertIfNew, ``toListModel_insertIfNew⟩,
      ⟨`insertMany, ``toListModel_insertMany_list⟩,
+     ⟨`union, ``toListModel_union⟩,
      ⟨`Const.insertMany, ``Const.toListModel_insertMany_list⟩,
      ⟨`Const.insertManyIfNewUnit, ``Const.toListModel_insertManyIfNewUnit_list⟩,
      ⟨`alter, ``toListModel_alter⟩,
@@ -2440,14 +2441,11 @@ section Union
 
 variable (m₁ m₂ : Raw₀ α β)
 
---Temporary: presumably we'll want to have this somewhere else (and redefine it in terms of `insertMany`?)
-def union : Raw₀ α β := ⟨m₁.val.union m₂.val, sorry⟩
-
 variable {m₁ m₂}
 
 @[simp]
 theorem union_insert_emptyWithCapacity {k : α} {v : β k} [EquivBEq α] [LawfulHashable α] (h : m.val.WF) :
-    m.union (emptyWithCapacity.insert k v) = m.insert k v  := by
+    m.union (emptyWithCapacity.insert k v) = m.insert k v := by
   sorry
 
 theorem contains_union_of_left  [EquivBEq α] [LawfulHashable α] (h₁ : m₁.val.WF)
@@ -2464,7 +2462,7 @@ theorem contains_union_of_right  [EquivBEq α] [LawfulHashable α] (h₁ : m₁.
 theorem contains_union [EquivBEq α] [LawfulHashable α] (h₁ : m₁.val.WF)
     (h₂ : m₂.val.WF) {k : α} :
     (m₁.union m₂).contains k = (m₁.contains k || m₂.contains k) := by
-  sorry
+  simp_to_model [contains, union] using List.containsKey_insertSmallerList
 
 @[simp]
 theorem contains_union_iff [EquivBEq α] [LawfulHashable α] (h₁ : m₁.val.WF)

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

@@ -983,7 +983,7 @@ theorem wfImp_filter [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] {m
 
 /-! # `insertListₘ` -/
 
-theorem toListModel_insertListₘ [BEq α] [Hashable α] [EquivBEq α][LawfulHashable α]
+theorem toListModel_insertListₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α]
     {m : Raw₀ α β} {l : List ((a : α) × β a)} (h : Raw.WFImp m.1) :
     Perm (toListModel (insertListₘ m l).1.buckets)
       (List.insertList (toListModel m.1.buckets) l) := by
@@ -995,6 +995,37 @@ theorem toListModel_insertListₘ [BEq α] [Hashable α] [EquivBEq α][LawfulHas
     apply Perm.trans (ih (wfImp_insert h))
     apply List.insertList_perm_of_perm_first (toListModel_insert h) (wfImp_insert h).distinct
 
+/-! # `unionₘ` -/
+
+theorem insertMany_eq_insertListₘ_toListModel [BEq α] [Hashable α] (m m₂ : Raw₀ α β) :
+    insertMany m m₂.1 = insertListₘ m (toListModel m₂.1.buckets) := by
+  simp only [insertMany, bind_pure_comp, map_pure, bind_pure]
+  simp only [ForIn.forIn]
+  simp only [Raw.forIn_eq_forIn_toListModel, forIn_pure_yield_eq_foldl, Id.run_pure]
+  generalize toListModel m₂.val.buckets = l
+  suffices ∀ (t : { m' // ∀ (P : Raw₀ α β → Prop),
+    (∀ {m'' : Raw₀ α β} {a : α} {b : β a}, P m'' → P (m''.insert a b)) → P m → P m' }),
+      (List.foldl (fun m' p => ⟨m'.val.insert p.1 p.2, fun P h₁ h₂ => h₁ (m'.2 _ h₁ h₂)⟩) t l).val =
+    t.val.insertListₘ l from this _
+  intro t
+  induction l generalizing m with
+  | nil => simp [insertListₘ]
+  | cons hd tl ih =>
+    simp only [List.foldl_cons, insertListₘ]
+    apply ih
+
+theorem union_eq_unionₘ [BEq α] [Hashable α] (m₁ m₂ : Raw₀ α β) :
+    union m₁ m₂ = unionₘ m₁ m₂ := by
+  rw [union, unionₘ]
+  split <;> rw [insertMany_eq_insertListₘ_toListModel]
+
+theorem toListModel_unionₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α]
+    {m₁ m₂ : Raw₀ α β} (h₁ : Raw.WFImp m₁.1) (h₂ : Raw.WFImp m₂.1) :
+    Perm (toListModel (unionₘ m₁ m₂).1.buckets)
+      (List.insertSmallerList (toListModel m₁.1.buckets) (toListModel m₂.1.buckets)) := by
+  rw [unionₘ, insertSmallerList, h₁.size_eq, h₂.size_eq]
+  split <;> exact toListModel_insertListₘ ‹_›
+
 end Raw₀
 
 namespace Raw
@@ -1134,6 +1165,22 @@ theorem toListModel_insertMany_list [BEq α] [Hashable α] [EquivBEq α] [Lawful
   apply toListModel_insertListₘ
   exact h
 
+/-! # `union` -/
+
+theorem wf_union₀ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α]
+    {m₁ m₂ : Raw α β} {h₁ : 0 < m₁.buckets.size} {h₂ : 0 < m₂.buckets.size} (h'₁ : m₁.WF)
+    (h'₂ : m₂.WF) :
+    (Raw₀.union ⟨m₁, h₁⟩ ⟨m₂, h₂⟩).1.WF := by
+  rw [union]
+  split <;> 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.insertSmallerList (toListModel m₁.1.buckets) (toListModel m₂.1.buckets)) := by
+  rw [union_eq_unionₘ]
+  exact toListModel_unionₘ h₁ h₂
+
 /-! # `Const.insertListₘ` -/
 
 theorem Const.toListModel_insertListₘ {β : Type v} [BEq α] [Hashable α] [EquivBEq α]

src/Std/Data/DHashMap/Raw.lean

@@ -397,20 +397,6 @@ by `foldM`. -/
 def foldRev (f : δ → (a : α) → β a → δ) (init : δ) (b : Raw α β) : δ :=
   Id.run (Internal.foldRevM (pure <| f · · ·) init b)
 
-/-- Carries out a monadic action on each mapping in the hash map in some order. -/
-@[inline] def forM (f : (a : α) → β a → m PUnit) (b : Raw α β) : m PUnit :=
-  b.buckets.forM (AssocList.forM f)
-
-/-- Support for the `for` loop construct in `do` blocks. -/
-@[inline] def forIn (f : (a : α) → β a → δ → m (ForInStep δ)) (init : δ) (b : Raw α β) : m δ :=
-  ForIn.forIn b.buckets init (fun bucket acc => bucket.forInStep acc f)
-
-instance : ForM m (Raw α β) ((a : α) × β a) where
-  forM m f := m.forM (fun a b => f ⟨a, b⟩)
-
-instance : ForIn m (Raw α β) ((a : α) × β a) where
-  forIn m init f := m.forIn (fun a b acc => f ⟨a, b⟩ acc) init
-
 namespace Const
 
 variable {β : Type v}

src/Std/Data/DHashMap/RawDef.lean

@@ -16,10 +16,12 @@ This file defines the type `Std.Data.DHashMap.Raw`. All of its functions are def
 set_option linter.missingDocs true
 set_option autoImplicit false
 
-universe u v
+universe u v w
 
 namespace Std.DHashMap
 
+open Internal
+
 /--
 Dependent hash maps without a bundled well-formedness invariant, suitable for use in nested
 inductive types. The well-formedness invariant is called `Raw.WF`. When in doubt, prefer `DHashMap`
@@ -46,4 +48,24 @@ structure Raw (α : Type u) (β : α → Type v) where
   /-- Internal implementation detail of the hash map -/
   buckets : Array (DHashMap.Internal.AssocList α β)
 
+namespace Raw
+
+variable {α : Type u} {β : α → Type v} {δ : Type w} {m : Type w → Type w}
+
+/-- Carries out a monadic action on each mapping in the hash map in some order. -/
+@[inline] def forM [Monad m] (f : (a : α) → β a → m PUnit) (b : Raw α β) : m PUnit :=
+  b.buckets.forM (AssocList.forM f)
+
+/-- Support for the `for` loop construct in `do` blocks. -/
+@[inline] def forIn [Monad m] (f : (a : α) → β a → δ → m (ForInStep δ)) (init : δ) (b : Raw α β) : m δ :=
+  ForIn.forIn b.buckets init (fun bucket acc => bucket.forInStep acc f)
+
+instance x : ForM m (Raw α β) ((a : α) × β a) where
+  forM m f := m.forM (fun a b => f ⟨a, b⟩)
+
+instance : ForIn m (Raw α β) ((a : α) × β a) where
+  forIn m init f := m.forIn (fun a b acc => f ⟨a, b⟩ acc) init
+
+end Raw
+
 end Std.DHashMap

src/Std/Data/Internal/List/Associative.lean

@@ -2972,6 +2972,22 @@ theorem isEmpty_insertList [BEq α]
     rw [insertList, List.isEmpty_cons, ih, isEmpty_insertEntry]
     simp
 
+/-- Internal implementation detail of the hash map -/
+def insertSmallerList [BEq α] (l₁ l₂ : List ((a : α) × β a)) : List ((a : α) × β a) :=
+  if l₁.length ≤ l₂.length then insertList l₂ l₁ else insertList l₁ l₂
+
+theorem DistinctKeys.insertSmallerList [BEq α] [PartialEquivBEq α] {l₁ l₂ : List ((a : α) × β a)}
+    (h₁ : DistinctKeys l₁) (h₂ : DistinctKeys l₂) : DistinctKeys (insertSmallerList l₁ l₂) := by
+  rw [List.insertSmallerList]
+  split <;> exact DistinctKeys.insertList ‹_›
+
+theorem containsKey_insertSmallerList [BEq α] [PartialEquivBEq α] {l₁ l₂ : List ((a : α) × β a)}
+    {k : α} : containsKey k (List.insertSmallerList l₁ l₂) = (containsKey k l₁ || containsKey k l₂) := by
+  rw [List.insertSmallerList]
+  split
+  · rw [containsKey_insertList, ← containsKey_eq_contains_map_fst, Bool.or_comm]
+  · rw [containsKey_insertList, ← containsKey_eq_contains_map_fst]
+
 section
 
 variable {β : Type v}