feat: tree map lemmas for getThenInsertIfNew? (#7229)

This PR provides lemmas for the tree map function `getThenInsertIfNew?`.

Co-authored-by: Paul Reichert <6992158+datokrat@users.noreply.github.com>

src/Std/Data/DTreeMap/Basic.lean

@@ -629,7 +629,7 @@ variable {β : Type v}
 def getThenInsertIfNew? (t : DTreeMap α β cmp) (a : α) (b : β) :
     Option β × DTreeMap α β cmp :=
   letI : Ord α := ⟨cmp⟩
-  let p := Impl.Const.getThenInsertIfNew? a b t.inner t.wf.balanced
+  let p := Impl.Const.getThenInsertIfNew? t.inner a b t.wf.balanced
   (p.1, ⟨p.2, t.wf.constGetThenInsertIfNew?⟩)
 
 @[inline, inherit_doc DTreeMap.get?]

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

@@ -1387,4 +1387,82 @@ theorem getKeyD_insertIfNew! [TransOrd α] (h : t.WF) {k a fallback : α}
   simp only [mem_iff_contains, Bool.not_eq_true]
   simp_to_model [insertIfNew!] using List.getKeyD_insertEntryIfNew
 
+/-!
+### getThenInsertIfNew?
+-/
+
+theorem getThenInsertIfNew?_fst [TransOrd α] [LawfulEqOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v h.balanced).1 = t.get? k := by
+  rw [getThenInsertIfNew?.eq_def]
+  cases t.get? k <;> rfl
+
+theorem getThenInsertIfNew?_snd [TransOrd α] [LawfulEqOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v h.balanced).2 = (t.insertIfNew k v h.balanced).impl := by
+  rw [getThenInsertIfNew?.eq_def]
+  cases heq : t.get? k
+  · rfl
+  · rw [get?_eq_getValueCast? h.ordered] at heq
+    rw [insertIfNew, contains_eq_containsKey h.ordered, List.containsKey_eq_isSome_getValueCast?, heq]
+    rfl
+
+/-!
+### getThenInsertIfNew?!
+-/
+
+theorem getThenInsertIfNew?!_fst [TransOrd α] [LawfulEqOrd α] {k : α} {v : β k} :
+    (t.getThenInsertIfNew?! k v).1 = t.get? k := by
+  rw [getThenInsertIfNew?!.eq_def]
+  cases t.get? k <;> rfl
+
+theorem getThenInsertIfNew?!_snd [TransOrd α] [LawfulEqOrd α] (h : t.WF) {k : α} {v : β k} :
+    (t.getThenInsertIfNew?! k v).2 = (t.insertIfNew! k v) := by
+  rw [getThenInsertIfNew?!.eq_def]
+  cases heq : t.get? k
+  · rfl
+  · rw [get?_eq_getValueCast? h.ordered] at heq
+    rw [insertIfNew!, contains_eq_containsKey h.ordered, List.containsKey_eq_isSome_getValueCast?, heq]
+    rfl
+
+namespace Const
+
+variable {β : Type v} {t : Impl α β}
+
+/-!
+### getThenInsertIfNew?
+-/
+
+theorem getThenInsertIfNew?_fst [TransOrd α] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v h.balanced).1 = get? t k := by
+  rw [getThenInsertIfNew?.eq_def]
+  cases get? t k <;> rfl
+
+theorem getThenInsertIfNew?_snd [TransOrd α] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v h.balanced).2 = (t.insertIfNew k v h.balanced).impl := by
+  rw [getThenInsertIfNew?.eq_def]
+  cases heq : get? t k
+  · rfl
+  · rw [get?_eq_getValue? h.ordered] at heq
+    rw [insertIfNew, contains_eq_containsKey h.ordered, List.containsKey_eq_isSome_getValue?, heq]
+    rfl
+
+/-!
+### getThenInsertIfNew?!
+-/
+
+theorem getThenInsertIfNew?!_fst [TransOrd α] {k : α} {v : β} :
+    (getThenInsertIfNew?! t k v).1 = get? t k := by
+  rw [getThenInsertIfNew?!.eq_def]
+  cases get? t k <;> rfl
+
+theorem getThenInsertIfNew?!_snd [TransOrd α] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew?! t k v).2 = (t.insertIfNew! k v) := by
+  rw [getThenInsertIfNew?!.eq_def]
+  cases heq : get? t k
+  · rfl
+  · rw [get?_eq_getValue? h.ordered] at heq
+    rw [insertIfNew!, contains_eq_containsKey h.ordered, List.containsKey_eq_isSome_getValue?, heq]
+    rfl
+
+end Const
+
 end Std.DTreeMap.Internal.Impl

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

@@ -401,7 +401,7 @@ def containsThenInsertIfNew! [Ord α] (k : α) (v : β k) (t : Impl α β) :
 
 /-- Implementation detail of the tree map -/
 @[inline]
-def getThenInsertIfNew? [Ord α] [LawfulEqOrd α] (k : α) (v : β k) (t : Impl α β) (ht : t.Balanced) :
+def getThenInsertIfNew? [Ord α] [LawfulEqOrd α] (t : Impl α β) (k : α) (v : β k) (ht : t.Balanced) :
     Option (β k) × Impl α β :=
   match t.get? k with
   | none => (none, t.insertIfNew k v ht |>.impl)
@@ -412,7 +412,7 @@ Slower version of `getThenInsertIfNew?` which can be used in the absence of bala
 information but still assumes the preconditions of `getThenInsertIfNew?`, otherwise might panic.
 -/
 @[inline]
-def getThenInsertIfNew?! [Ord α] [LawfulEqOrd α] (k : α) (v : β k) (t : Impl α β) :
+def getThenInsertIfNew?! [Ord α] [LawfulEqOrd α] (t : Impl α β) (k : α) (v : β k) :
     Option (β k) × Impl α β :=
   match t.get? k with
   | none => (none, t.insertIfNew! k v)
@@ -604,7 +604,7 @@ variable {β : Type v}
 
 /-- Implementation detail of the tree map -/
 @[inline]
-def getThenInsertIfNew? [Ord α] (k : α) (v : β) (t : Impl α (fun _ => β))
+def getThenInsertIfNew? [Ord α] (t : Impl α (fun _ => β)) (k : α) (v : β)
     (ht : t.Balanced) : Option β × Impl α (fun _ => β) :=
   match get? t k with
   | none => (none, t.insertIfNew k v ht |>.impl)
@@ -615,7 +615,7 @@ Slower version of `getThenInsertIfNew?` which can be used in the absence of bala
 information but still assumes the preconditions of `getThenInsertIfNew?`, otherwise might panic.
 -/
 @[inline]
-def getThenInsertIfNew?! [Ord α] (k : α) (v : β) (t : Impl α (fun _ => β))
+def getThenInsertIfNew?! [Ord α] (t : Impl α (fun _ => β)) (k : α) (v : β)
     : Option β × Impl α (fun _ => β) :=
   match get? t k with
   | none => (none, t.insertIfNew! k v)

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

@@ -97,7 +97,7 @@ section Const
 variable {β : Type v}
 
 theorem WF.constGetThenInsertIfNew? [Ord α] {t : Impl α β} {k v} {h : t.WF} :
-    (Impl.Const.getThenInsertIfNew? k v t h.balanced).2.WF := by
+    (Impl.Const.getThenInsertIfNew? t k v h.balanced).2.WF := by
   simp only [Impl.Const.getThenInsertIfNew?]
   split
   · exact h.insertIfNew

src/Std/Data/DTreeMap/Lemmas.lean

@@ -896,4 +896,30 @@ theorem getKeyD_insertIfNew [TransCmp cmp] {k a fallback : α} {v : β k} :
       if cmp k a = .eq ∧ ¬ k ∈ t then k else t.getKeyD a fallback :=
   Impl.getKeyD_insertIfNew t.wf
 
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] [LawfulEqCmp cmp] {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v).1 = t.get? k :=
+  Impl.getThenInsertIfNew?_fst t.wf
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] [LawfulEqCmp cmp] {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v).2 = t.insertIfNew k v :=
+  ext <| Impl.getThenInsertIfNew?_snd t.wf
+
+namespace Const
+
+variable {β : Type v} {t : DTreeMap α β cmp}
+
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).1 = get? t k :=
+  Impl.Const.getThenInsertIfNew?_fst t.wf
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).2 = t.insertIfNew k v :=
+  ext <| Impl.Const.getThenInsertIfNew?_snd t.wf
+
+end Const
+
 end Std.DTreeMap

src/Std/Data/DTreeMap/Raw.lean

@@ -406,7 +406,7 @@ variable {β : Type v}
 @[inline, inherit_doc DTreeMap.Const.getThenInsertIfNew?]
 def getThenInsertIfNew? (t : Raw α β cmp) (a : α) (b : β) : Option β × Raw α β cmp :=
   letI : Ord α := ⟨cmp⟩
-  let p := Impl.Const.getThenInsertIfNew?! a b t.inner
+  let p := Impl.Const.getThenInsertIfNew?! t.inner a b
   (p.1, ⟨p.2⟩)
 
 @[inline, inherit_doc DTreeMap.Const.get?]

src/Std/Data/DTreeMap/RawLemmas.lean

@@ -903,4 +903,30 @@ theorem getKeyD_insertIfNew [TransCmp cmp] (h : t.WF) {k a fallback : α}
       if cmp k a = .eq ∧ ¬ k ∈ t then k else t.getKeyD a fallback :=
   Impl.getKeyD_insertIfNew! h
 
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] [LawfulEqCmp cmp] (_ : t.WF) {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v).1 = t.get? k :=
+  Impl.getThenInsertIfNew?!_fst
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] [LawfulEqCmp cmp] (h : t.WF) {k : α} {v : β k} :
+    (t.getThenInsertIfNew? k v).2 = t.insertIfNew k v :=
+  ext <| Impl.getThenInsertIfNew?!_snd h
+
+namespace Const
+
+variable {β : Type v} {t : Raw α β cmp}
+
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] (_ : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).1 = get? t k :=
+  Impl.Const.getThenInsertIfNew?!_fst
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).2 = t.insertIfNew k v :=
+  ext <| Impl.Const.getThenInsertIfNew?!_snd h
+
+end Const
+
 end Std.DTreeMap.Raw

src/Std/Data/TreeMap/Lemmas.lean

@@ -645,4 +645,14 @@ theorem getKeyD_insertIfNew [TransCmp cmp] {k a fallback : α}
       if cmp k a = .eq ∧ ¬ k ∈ t then k else t.getKeyD a fallback :=
   DTreeMap.getKeyD_insertIfNew
 
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).1 = get? t k :=
+  DTreeMap.Const.getThenInsertIfNew?_fst
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).2 = t.insertIfNew k v :=
+  ext <| DTreeMap.Const.getThenInsertIfNew?_snd
+
 end Std.TreeMap

src/Std/Data/TreeMap/RawLemmas.lean

@@ -651,4 +651,14 @@ theorem getKeyD_insertIfNew [TransCmp cmp] (h : t.WF) {k a fallback : α}
       if cmp k a = .eq ∧ ¬ k ∈ t then k else t.getKeyD a fallback :=
   DTreeMap.Raw.getKeyD_insertIfNew h
 
+@[simp]
+theorem getThenInsertIfNew?_fst [TransCmp cmp] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).1 = get? t k :=
+  DTreeMap.Raw.Const.getThenInsertIfNew?_fst h
+
+@[simp]
+theorem getThenInsertIfNew?_snd [TransCmp cmp] (h : t.WF) {k : α} {v : β} :
+    (getThenInsertIfNew? t k v).2 = t.insertIfNew k v :=
+  ext <| DTreeMap.Raw.Const.getThenInsertIfNew?_snd h
+
 end Std.TreeMap.Raw