From 6c62f720c8b6f6a172df41693719a4aca0110faa Mon Sep 17 00:00:00 2001
From: Paul Reichert <datokrat@users.noreply.github.com>
Date: Wed, 26 Feb 2025 11:29:51 +0100
Subject: [PATCH] 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              |  2 +-
 src/Std/Data/DTreeMap/Internal/Lemmas.lean    | 78 +++++++++++++++++++
 .../Data/DTreeMap/Internal/Operations.lean    |  8 +-
 src/Std/Data/DTreeMap/Internal/WF/Defs.lean   |  2 +-
 src/Std/Data/DTreeMap/Lemmas.lean             | 26 +++++++
 src/Std/Data/DTreeMap/Raw.lean                |  2 +-
 src/Std/Data/DTreeMap/RawLemmas.lean          | 26 +++++++
 src/Std/Data/TreeMap/Lemmas.lean              | 10 +++
 src/Std/Data/TreeMap/RawLemmas.lean           | 10 +++
 9 files changed, 157 insertions(+), 7 deletions(-)

diff --git a/src/Std/Data/DTreeMap/Basic.lean b/src/Std/Data/DTreeMap/Basic.lean
index f6254e41a6..12decd73ac 100644
--- a/src/Std/Data/DTreeMap/Basic.lean
+++ b/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?]
diff --git a/src/Std/Data/DTreeMap/Internal/Lemmas.lean b/src/Std/Data/DTreeMap/Internal/Lemmas.lean
index 8c4bb37bea..1eb6909601 100644
--- a/src/Std/Data/DTreeMap/Internal/Lemmas.lean
+++ b/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
diff --git a/src/Std/Data/DTreeMap/Internal/Operations.lean b/src/Std/Data/DTreeMap/Internal/Operations.lean
index 3810c2c46e..a4a0326aab 100644
--- a/src/Std/Data/DTreeMap/Internal/Operations.lean
+++ b/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)
diff --git a/src/Std/Data/DTreeMap/Internal/WF/Defs.lean b/src/Std/Data/DTreeMap/Internal/WF/Defs.lean
index cda3d42cab..a438e720d9 100644
--- a/src/Std/Data/DTreeMap/Internal/WF/Defs.lean
+++ b/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
diff --git a/src/Std/Data/DTreeMap/Lemmas.lean b/src/Std/Data/DTreeMap/Lemmas.lean
index 53e2d3b79e..172b98fc5a 100644
--- a/src/Std/Data/DTreeMap/Lemmas.lean
+++ b/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
diff --git a/src/Std/Data/DTreeMap/Raw.lean b/src/Std/Data/DTreeMap/Raw.lean
index 35cd6b9594..7c6f952939 100644
--- a/src/Std/Data/DTreeMap/Raw.lean
+++ b/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?]
diff --git a/src/Std/Data/DTreeMap/RawLemmas.lean b/src/Std/Data/DTreeMap/RawLemmas.lean
index 9f67aeb022..90a40cf64f 100644
--- a/src/Std/Data/DTreeMap/RawLemmas.lean
+++ b/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
diff --git a/src/Std/Data/TreeMap/Lemmas.lean b/src/Std/Data/TreeMap/Lemmas.lean
index 8d7a40077d..ea405a7d85 100644
--- a/src/Std/Data/TreeMap/Lemmas.lean
+++ b/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
diff --git a/src/Std/Data/TreeMap/RawLemmas.lean b/src/Std/Data/TreeMap/RawLemmas.lean
index b8951eb19c..9c6cf36e89 100644
--- a/src/Std/Data/TreeMap/RawLemmas.lean
+++ b/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