feat: insertIfNew for the tree map

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

@@ -704,7 +704,7 @@ theorem contains_of_contains_insertIfNew [EquivBEq α] [LawfulHashable α] (h :
     {v : β k} : (m.insertIfNew k v).contains a → (k == a) = false → m.contains a := by
   simp_to_model using List.containsKey_of_containsKey_insertEntryIfNew
 
-/-- This is a restatement of `contains_insertIfNew` that is written to exactly match the proof
+/-- This is a restatement of `contains_of_contains_insertIfNew` that is written to exactly match the proof
 obligation in the statement of `get_insertIfNew`. -/
 theorem contains_of_contains_insertIfNew' [EquivBEq α] [LawfulHashable α] (h : m.1.WF) {k a : α}
     {v : β k} :

src/Std/Data/DHashMap/Lemmas.lean

@@ -815,13 +815,13 @@ theorem mem_of_mem_insertIfNew [EquivBEq α] [LawfulHashable α] {k a : α} {v :
     a ∈ m.insertIfNew k v → (k == a) = false → a ∈ m := by
   simpa [mem_iff_contains, -contains_insertIfNew] using contains_of_contains_insertIfNew
 
-/-- This is a restatement of `contains_insertIfNew` that is written to exactly match the proof
+/-- This is a restatement of `contains_of_contains_insertIfNew` that is written to exactly match the proof
 obligation in the statement of `get_insertIfNew`. -/
 theorem contains_of_contains_insertIfNew' [EquivBEq α] [LawfulHashable α] {k a : α} {v : β k} :
     (m.insertIfNew k v).contains a → ¬((k == a) ∧ m.contains k = false) → m.contains a :=
   Raw₀.contains_of_contains_insertIfNew' ⟨m.1, _⟩ m.2
 
-/-- This is a restatement of `mem_insertIfNew` that is written to exactly match the proof obligation
+/-- This is a restatement of `mem_of_mem_insertIfNew` that is written to exactly match the proof obligation
 in the statement of `get_insertIfNew`. -/
 theorem mem_of_mem_insertIfNew' [EquivBEq α] [LawfulHashable α] {k a : α} {v : β k} :
     a ∈ m.insertIfNew k v → ¬((k == a) ∧ ¬k ∈ m) → a ∈ m := by

src/Std/Data/DHashMap/RawLemmas.lean

@@ -895,14 +895,14 @@ theorem mem_of_mem_insertIfNew [EquivBEq α] [LawfulHashable α] (h : m.WF) {k a
     a ∈ m.insertIfNew k v → (k == a) = false → a ∈ m := by
   simpa [mem_iff_contains, -contains_insertIfNew] using contains_of_contains_insertIfNew h
 
-/-- This is a restatement of `contains_insertIfNew` that is written to exactly match the proof
+/-- This is a restatement of `contains_of_contains_insertIfNew` that is written to exactly match the proof
 obligation in the statement of `get_insertIfNew`. -/
 theorem contains_of_contains_insertIfNew' [EquivBEq α] [LawfulHashable α] (h : m.WF) {k a : α}
     {v : β k} :
     (m.insertIfNew k v).contains a → ¬((k == a) ∧ m.contains k = false) → m.contains a := by
   simp_to_raw using Raw₀.contains_of_contains_insertIfNew'
 
-/-- This is a restatement of `mem_insertIfNew` that is written to exactly match the proof obligation
+/-- This is a restatement of `mem_of_mem_insertIfNew` that is written to exactly match the proof obligation
 in the statement of `get_insertIfNew`. -/
 theorem mem_of_mem_insertIfNew' [EquivBEq α] [LawfulHashable α] (h : m.WF) {k a : α} {v : β k} :
     a ∈ m.insertIfNew k v → ¬((k == a) ∧ ¬k ∈ m) → a ∈ m := by

src/Std/Data/DTreeMap/Basic.lean

@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Markus Himmel
 -/
 prelude
-import Std.Data.DTreeMap.Internal.Impl
+import Std.Data.DTreeMap.Internal.WF
 import Std.Data.DTreeMap.Raw
 
 /-!
@@ -97,9 +97,16 @@ def containsThenInsert (t : DTreeMap α β cmp) (a : α) (b : β a) : Bool × DT
   let p := t.inner.containsThenInsert a b t.wf.balanced
   (p.1, ⟨p.2.impl, t.wf.containsThenInsert⟩)
 
+@[inline, inherit_doc Raw.insertIfNew]
+def insertIfNew (t : DTreeMap α β cmp) (a : α) (b : β a) : DTreeMap α β cmp :=
+    letI : Ord α := ⟨cmp⟩; ⟨(t.inner.insertIfNew a b t.wf.balanced).impl, t.wf.insertIfNew⟩
+
 instance : Membership α (DTreeMap α β cmp) where
   mem m a := m.contains a
 
+instance {m : DTreeMap α β cmp} {a : α} : Decidable (a ∈ m) :=
+  show Decidable (m.contains a) from inferInstance
+
 end DTreeMap
 
 end Std

src/Std/Data/DTreeMap/Internal/Impl/Query.lean

@@ -45,6 +45,9 @@ def contains [Ord α] (k : α) (t : Impl α β) : Bool :=
 instance [Ord α] : Membership α (Impl α β) where
   mem t a := t.contains a
 
+instance [Ord α] {m : Impl α β} {a : α} : Decidable (a ∈ m) :=
+  show Decidable (m.contains a) from inferInstance
+
 /-- Returns `true` if the tree is empty. -/
 @[inline]
 def isEmpty (t : Impl α β) : Bool :=

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

@@ -28,6 +28,7 @@ scoped macro "wf_trivial" : tactic => `(tactic|
     | apply WF.ordered | apply WF.balanced | apply WF.insert | apply WF.insertSlow
     | apply WF.erase | apply WF.eraseSlow
     | apply WF.containsThenInsert | apply WF.containsThenInsertSlow
+    | apply WF.insertIfNew | apply WF.insertIfNewSlow
     | apply Ordered.distinctKeys
     | assumption
     ))
@@ -42,7 +43,8 @@ private def queryNames : Array Name :=
 
 private def modifyNames : Array Name :=
   #[``toListModel_insert, ``toListModel_insertSlow, ``toListModel_erase, ``toListModel_eraseSlow,
-    ``toListModel_containsThenInsert, ``toListModel_containsThenInsertSlow]
+    ``toListModel_containsThenInsert, ``toListModel_containsThenInsertSlow,
+    ``toListModel_insertIfNew, ``toListModel_insertIfNewSlow]
 
 private def congrNames : MacroM (Array (TSyntax `term)) := do
   return #[← `(_root_.List.Perm.isEmpty_eq), ← `(containsKey_of_perm),
@@ -195,6 +197,118 @@ theorem containsThenInsertSlow_snd [TransOrd α] (h : t.WF) {k : α} {v : β k}
   rw [snd_containsThenInsertSlow_eq_containsThenInsert _ h.balanced, containsThenInsert_snd h,
     insert_eq_insertSlow]
 
+@[simp]
+theorem isEmpty_insertIfNew [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.isEmpty = false := by
+  simp_to_model using List.isEmpty_insertEntryIfNew
+
+@[simp]
+theorem isEmpty_insertIfNewSlow [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNewSlow k v).isEmpty = false := by
+  simp_to_model using List.isEmpty_insertEntryIfNew
+
+@[simp]
+theorem contains_insertIfNew [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.contains a = (k == a || t.contains a) := by
+  simp_to_model using List.containsKey_insertEntryIfNew
+
+@[simp]
+theorem contains_insertIfNewSlow [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNewSlow k v).contains a = (k == a || t.contains a) := by
+  simp_to_model using List.containsKey_insertEntryIfNew
+
+@[simp]
+theorem mem_insertIfNew [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ (t.insertIfNew k v h.balanced).impl ↔ k == a ∨ a ∈ t := by
+  simp [mem_iff_contains, contains_insertIfNew, h]
+
+@[simp]
+theorem mem_insertIfNewSlow [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ t.insertIfNewSlow k v ↔ k == a ∨ a ∈ t := by
+  simp [mem_iff_contains, contains_insertIfNew, h]
+
+theorem contains_insertIfNew_self [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.contains k := by
+  simp [contains_insertIfNew, h]
+
+theorem contains_insertIfNewSlow_self [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNewSlow k v).contains k := by
+  simp [contains_insertIfNewSlow, h]
+
+theorem mem_insertIfNew_self [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    k ∈ (t.insertIfNew k v h.balanced).impl := by
+  simp [contains_insertIfNew, h]
+
+theorem mem_insertIfNewSlow_self [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    k ∈ t.insertIfNewSlow k v := by
+  simp [contains_insertIfNewSlow, h]
+
+theorem contains_of_contains_insertIfNew [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.contains a → (k == a) = false → t.contains a := by
+  simp_to_model using List.containsKey_of_containsKey_insertEntryIfNew
+
+theorem contains_of_contains_insertIfNewSlow [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNewSlow k v).contains a → (k == a) = false → t.contains a := by
+  simp_to_model using List.containsKey_of_containsKey_insertEntryIfNew
+
+theorem mem_of_mem_insertIfNew [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ (t.insertIfNew k v h.balanced).impl → (k == a) = false → a ∈ t := by
+  simpa [mem_iff_contains] using contains_of_contains_insertIfNew h
+
+theorem mem_of_mem_insertIfNewSlow [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ t.insertIfNewSlow k v → (k == a) = false → a ∈ t := by
+  simpa [mem_iff_contains] using contains_of_contains_insertIfNewSlow h
+
+/-- This is a restatement of `contains_of_contains_insertIfNew` that is written to exactly match the proof
+obligation in the statement of `get_insertIfNew`. -/
+theorem contains_of_contains_insertIfNew' [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.contains a → ¬((k == a) ∧ t.contains k = false) → t.contains a := by
+  simp_to_model using List.containsKey_of_containsKey_insertEntryIfNew'
+
+/-- This is a restatement of `contains_of_contains_insertIfNewSlow` that is written to exactly match the proof
+obligation in the statement of `get_insertIfNewSlow`. -/
+theorem contains_of_contains_insertIfNewSlow' [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    (t.insertIfNewSlow k v).contains a → ¬((k == a) ∧ t.contains k = false) → t.contains a := by
+  simp_to_model using List.containsKey_of_containsKey_insertEntryIfNew'
+
+/-- This is a restatement of `mem_of_mem_insertIfNew` that is written to exactly match the proof obligation
+in the statement of `get_insertIfNew`. -/
+theorem mem_of_mem_insertIfNew' [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ (t.insertIfNew k v h.balanced).impl → ¬((k == a) ∧ ¬k ∈ t) → a ∈ t := by
+  simpa [mem_iff_contains] using contains_of_contains_insertIfNew' h
+
+/-- This is a restatement of `mem_of_mem_insertIfNew` that is written to exactly match the proof obligation
+in the statement of `get_insertIfNew`. -/
+theorem mem_of_mem_insertIfNewSlow' [TransOrd α] (h : t.WF) {k a : α} {v : β k} :
+    a ∈ t.insertIfNewSlow k v → ¬((k == a) ∧ ¬k ∈ t) → a ∈ t := by
+  simpa [mem_iff_contains] using contains_of_contains_insertIfNewSlow' h
+
+theorem size_insertIfNew [TransOrd α] {k : α} (h : t.WF) {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.size = if k ∈ t then t.size else t.size + 1 := by
+  simp only [mem_iff_contains]
+  simp_to_model using List.length_insertEntryIfNew
+
+theorem size_insertIfNewSlow [TransOrd α] {k : α} (h : t.WF) {v : β k} :
+    (t.insertIfNewSlow k v).size = if k ∈ t then t.size else t.size + 1 := by
+  simp only [mem_iff_contains]
+  simp_to_model using List.length_insertEntryIfNew
+
+theorem size_le_size_insertIfNew [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    t.size ≤ (t.insertIfNew k v h.balanced).impl.size := by
+  simp_to_model using List.length_le_length_insertEntryIfNew
+
+theorem size_le_size_insertIfNewSlow [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    t.size ≤ (t.insertIfNewSlow k v).size := by
+  simp_to_model using List.length_le_length_insertEntryIfNew
+
+theorem size_insertIfNew_le [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNew k v h.balanced).impl.size ≤ t.size + 1 := by
+  simp_to_model using List.length_insertEntryIfNew_le
+
+theorem size_insertIfNewSlow_le [TransOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.insertIfNewSlow k v).size ≤ t.size + 1 := by
+  simp_to_model using List.length_insertEntryIfNew_le
+
 end Std.DTreeMap.Internal.Impl
 
 /-- Used as instance by the dependents of this file -/

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

@@ -216,6 +216,15 @@ Internal implementation detail of the tree map
 def eraseₘ [Ord α] (k : α) (t : Impl α β) (h : t.Balanced) : Impl α β :=
   updateCell k (fun _ => .empty) t h |>.impl
 
+/--
+Model implementation of the `insertIfNew` function.
+Internal implementation detail of the tree map
+-/
+def insertIfNewₘ [Ord α] (k : α) (v : β k) (l : Impl α β) (h : l.Balanced) : Impl α β :=
+  updateCell k (fun
+    | ⟨.none, _⟩ => .of k v
+    | c => c) l h |>.impl
+
 /-!
 ## Helper theorems for reasoning with key-value pairs
 -/
@@ -380,6 +389,25 @@ theorem snd_containsThenInsertIfNewSlow_eq_containsThenInsert [Ord α] (t : Impl
   · rfl
   · simp [insertSlow_eq_insertₘ, insert_eq_insertₘ, htb]
 
+theorem insertIfNew_eq_insertIfNewSlow [Ord α] {k : α} {v : β k} {l : Impl α β} {h} :
+    (insertIfNew k v l h).impl = insertIfNewSlow k v l := by
+  simp [insertIfNew, insertIfNewSlow]
+  split
+  · rfl
+  · simp [insert_eq_insertSlow]
+
+-- theorem insertIfNew_eq_insertₘ [Ord α] {k : α} {v : β k} {l : Impl α β} {h} :
+--     (insertIfNew k v l h).impl = insertₘ k v l h := by
+--   simp only [insertIfNewₘ]
+--   induction l
+--   · simp only [insertIfNew, updateCell]
+--     split <;> split <;> simp_all [balanceL_eq_balance, balanceR_eq_balance]
+--   · simp [insertIfNew, insertₘ, updateCell]
+
+-- theorem insertIfNewSlow_eq_insertₘ [Ord α] {k : α} {v : β k} {l : Impl α β} (h : l.Balanced) :
+--     insertIfNewSlow k v l = insertₘ k v l h := by
+--   rw [← insertIfNew_eq_insertSlow (h := h), insert_eq_insertₘ]
+
 end Impl
 
 end Std.DTreeMap.Internal

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

@@ -717,6 +717,58 @@ theorem toListModel_containsThenInsertSlow [Ord α] [TransOrd α] {k : α} {v :
   rw [containsThenInsertSlow_eq_insertₘ]
   exact toListModel_insertₘ htb hto
 
+
+-- /-!
+-- ### `insertIfNewₘ`
+-- -/
+
+-- theorem WF.insert
+
+-- theorem ordered_insertIfNewₘ [Ord α] [TransOrd α] {k : α} {v : β k} {l : Impl α β} (hlb : l.Balanced)
+--     (hlo : l.Ordered) : (l.insertIfNewₘ k v hlb).Ordered :=
+--   ordered_updateAtKey _ hlo
+
+-- theorem toListModel_insertIfNewₘ [Ord α] [TransOrd α] {k : α} {v : β k} {l : Impl α β} (hlb : l.Balanced)
+--     (hlo : l.Ordered) : (l.insertIfNewₘ k v hlb).toListModel.Perm (insertEntry k v l.toListModel) := by
+--   refine toListModel_updateAtKey_perm _ hlo ?_ insertIfNewEntry_of_perm
+--     insertIfNewEntry_append_of_not_contains_right
+--   rintro ⟨(_|l), hl⟩
+--   · simp
+--   · simp only [Option.toList_some, Cell.of_inner]
+--     have h : l.fst == k := by simpa using OrientedCmp.eq_symm (hl l rfl)
+--     rw [insertIfNewEntry_of_containsKey (containsKey_cons_of_beq h), replaceEntry_cons_of_true h]
+
+/-!
+### `insertIfNew`
+-/
+
+theorem WF.insertIfNew [Ord α] {k : α} {v : β k} {l : Impl α β}
+    (h : l.WF) : (l.insertIfNew k v h.balanced).impl.WF := by
+  simp [Impl.insertIfNew]
+  split <;> simp only [h, h.insert]
+
+theorem toListModel_insertIfNew [Ord α] [TransOrd α] {k : α} {v : β k} {l : Impl α β}
+    (hlb : l.Balanced) (hlo : l.Ordered) :
+    (l.insertIfNew k v hlb).impl.toListModel.Perm (insertEntryIfNew k v l.toListModel) := by
+  simp only [Impl.insertIfNew, insertEntryIfNew, cond_eq_if, apply_contains hlo]
+  split
+  · rfl
+  · refine (toListModel_insert hlb hlo).trans ?_
+    simp [insertEntry_of_containsKey_eq_false, *]
+
+/-!
+### `insertIfNewSlow`
+-/
+
+theorem WF.insertIfNewSlow [Ord α] {k : α} {v : β k} {l : Impl α β}
+    (h : l.WF) : (l.insertIfNewSlow k v).WF := by
+  simpa [insertIfNew_eq_insertIfNewSlow] using WF.insertIfNew h
+
+theorem toListModel_insertIfNewSlow [Ord α] [TransOrd α] {k : α} {v : β k} {l : Impl α β}
+    (hlb : l.Balanced) (hlo : l.Ordered) :
+    (l.insertIfNewSlow k v).toListModel.Perm (insertEntryIfNew k v l.toListModel) := by
+  simpa [insertIfNew_eq_insertIfNewSlow] using toListModel_insertIfNew hlb hlo
+
 /-!
 ## Deducing that well-formed trees are ordered
 -/

src/Std/Data/DTreeMap/Lemmas.lean

@@ -98,4 +98,59 @@ theorem containsThenInsert_snd [TransCmp cmp] {k : α} {v : β k} :
     (t.containsThenInsert k v).2 = t.insert k v :=
   ext <| congrArg Impl.TreeB.impl <| Impl.containsThenInsert_snd t.wf
 
+@[simp]
+theorem isEmpty_insertIfNew [TransCmp cmp] {k : α} {v : β k} :
+    (t.insertIfNew k v).isEmpty = false :=
+  Impl.isEmpty_insertIfNew t.wf
+
+@[simp]
+theorem contains_insertIfNew [TransCmp cmp] {k a : α} {v : β k} :
+    (t.insertIfNew k v).contains a = (cmp k a == .eq || t.contains a) :=
+  Impl.contains_insertIfNew t.wf
+
+@[simp]
+theorem mem_insertIfNew [TransCmp cmp] {k a : α} {v : β k} :
+    a ∈ t.insertIfNew k v ↔ cmp k a == .eq ∨ a ∈ t :=
+  Impl.mem_insertIfNew t.wf
+
+theorem contains_insertIfNew_self [TransCmp cmp] {k : α} {v : β k} :
+    (t.insertIfNew k v).contains k :=
+  Impl.contains_insertIfNew_self t.wf
+
+theorem mem_insertIfNew_self [TransCmp cmp] {k : α} {v : β k} :
+    k ∈ t.insertIfNew k v :=
+  Impl.mem_insertIfNew_self t.wf
+
+theorem contains_of_contains_insertIfNew [TransCmp cmp] {k a : α} {v : β k} :
+    (t.insertIfNew k v).contains a → (cmp k a == .eq) = false → t.contains a :=
+  Impl.contains_of_contains_insertIfNew t.wf
+
+theorem mem_of_mem_insertIfNew [TransCmp cmp] {k a : α} {v : β k} :
+    a ∈ t.insertIfNew k v → (cmp k a == .eq) = false → a ∈ t :=
+  Impl.contains_of_contains_insertIfNew t.wf
+
+/-- This is a restatement of `contains_of_contains_insertIfNew` that is written to exactly match the proof
+obligation in the statement of `get_insertIfNew`. -/
+theorem contains_of_contains_insertIfNew' [TransCmp cmp] {k a : α} {v : β k} :
+    (t.insertIfNew k v).contains a → ¬((cmp k a == .eq) ∧ t.contains k = false) → t.contains a :=
+  Impl.contains_of_contains_insertIfNew' t.wf
+
+/-- This is a restatement of `mem_of_mem_insertIfNew` that is written to exactly match the proof obligation
+in the statement of `get_insertIfNew`. -/
+theorem mem_of_mem_insertIfNew' [TransCmp cmp] {k a : α} {v : β k} :
+    a ∈ t.insertIfNew k v → ¬((cmp k a == .eq) ∧ ¬k ∈ t) → a ∈ t :=
+  Impl.mem_of_mem_insertIfNew' t.wf
+
+theorem size_insertIfNew [TransCmp cmp] {k : α} {v : β k} :
+    (t.insertIfNew k v).size = if k ∈ t then t.size else t.size + 1 :=
+  Impl.size_insertIfNew t.wf
+
+theorem size_le_size_insertIfNew [TransCmp cmp] {k : α} {v : β k} :
+    t.size ≤ (t.insertIfNew k v).size :=
+  Impl.size_le_size_insertIfNew t.wf
+
+theorem size_insertIfNew_le [TransCmp cmp] {k : α} {v : β k} :
+    (t.insertIfNew k v).size ≤ t.size + 1 :=
+  Impl.size_insertIfNew_le t.wf
+
 end Std.DTreeMap

src/Std/Data/DTreeMap/Raw.lean

@@ -155,9 +155,19 @@ def containsThenInsert (t : Raw α β cmp) (a : α) (b : β a) : Bool × Raw α
   let p := t.inner.containsThenInsertSlow a b
   (p.1, ⟨p.2⟩)
 
+/--
+If there is no mapping for the given key, inserts the given mapping into the map. Otherwise,
+returns the map unaltered.
+-/
+@[inline] def insertIfNew (t : Raw α β cmp) (a : α) (b : β a) : Raw α β cmp :=
+    letI : Ord α := ⟨cmp⟩; ⟨t.inner.insertIfNewSlow a b⟩
+
 instance : Membership α (Raw α β cmp) where
   mem m a := m.contains a
 
+instance {m : Raw α β cmp} {a : α} : Decidable (a ∈ m) :=
+  show Decidable (m.contains a) from inferInstance
+
 end Raw
 
 end DTreeMap