feat: add minimal support for getEntry/getEntry?/getEntry!/getEntryD for DTreeMap (#11161)

This PR adds getEntry/getEntry?/getEntry!/getEntryD operation on
DTreeMap.

src/Std/Data/DTreeMap/Basic.lean

@@ -272,6 +272,38 @@ If a mapping exists the result is guaranteed to be pointer equal to the key in t
 def getKeyD (t : DTreeMap α β cmp) (a : α) (fallback : α) : α :=
   letI : Ord α := ⟨cmp⟩; t.inner.getKeyD a fallback
 
+
+/--
+Checks if a mapping for the given key exists and returns the key-value pair if it does, otherwise `none`.
+The key in the returned pair will be `BEq` to the input `a`.
+-/
+@[inline]
+def getEntry? (t : DTreeMap α β cmp) (a : α) : Option ((a : α) × β a) :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry? a
+
+/--
+Retrieves the key-value pair, whose key matches `a`. Ensures that such a mapping exists by requiring a proof of `a ∈ m`. The key in the returned pair will be `BEq` to the input `a`.
+-/
+@[inline]
+def getEntry (t : DTreeMap α β cmp) (a : α) (h : a ∈ t) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry a h
+
+/--
+Checks if a mapping for the given key exists and returns the key-value pair if it does, otherwise `fallback`.
+The key in the returned pair will compare equal to the input `a`.
+-/
+@[inline]
+def getEntryD (t : DTreeMap α β cmp) (a : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntryD a fallback
+
+/--
+Checks if a mapping for the given key exists and returns the key-value pair if it does, otherwise panics.
+The key in the returned pair will be `BEq` to the input `a`.
+-/
+@[inline]
+def getEntry! [Inhabited ((a : α) × β a)] (t : DTreeMap α β cmp) (a : α) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry! a
+
 /--
 Tries to retrieve the key-value pair with the smallest key in the tree map, returning `none` if the
 map is empty.

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

@@ -94,6 +94,12 @@ def get? [Ord α] [OrientedOrd α] [LawfulEqOrd α] {k : α} (c : Cell α β (co
   | none => none
   | some p => some (cast (congrArg β (compare_eq_iff_eq.mp (c.property _ h)).symm) p.2)
 
+/-- Internal implementation detail of the tree map -/
+def getEntry? [Ord α] {k : α} (c : Cell α β (compare k)) : Option ((a : α) × β a) :=
+  match c.inner with
+  | none => none
+  | some p => some ⟨p.1, p.2⟩
+
 @[simp]
 theorem get?_empty [Ord α] [OrientedOrd α] [LawfulEqOrd α] {k : α} :
     (Cell.empty : Cell α β (compare k)).get? = none :=
@@ -105,6 +111,11 @@ def getKey? [Ord α] {k : α} (c : Cell α β (compare k)) : Option α :=
   | none => none
   | some p => some p.1
 
+@[simp]
+theorem getEntry?_empty [Ord α] [OrientedOrd α] [LawfulEqOrd α] {k : α} :
+    (Cell.empty : Cell α β (compare k)).getEntry? = none :=
+  rfl
+
 @[simp]
 theorem getKey?_empty [Ord α] {k : α} : (Cell.empty : Cell α β (compare k)).getKey? = none :=
   rfl

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

@@ -87,6 +87,10 @@ private meta def queryMap : Std.DHashMap Name (fun _ => Name × Array (MacroM (T
      ⟨`get, (``get_eq_getValueCast, #[``(getValueCast_of_perm _)])⟩,
      ⟨`get!, (``get!_eq_getValueCast!, #[``(getValueCast!_of_perm _)])⟩,
      ⟨`getD, (``getD_eq_getValueCastD, #[``(getValueCastD_of_perm _)])⟩,
+     ⟨`getEntry?, (``getEntry?_eq_getEntry?, #[``(List.getEntry?_of_perm _)])⟩,
+     ⟨`getEntry, (``getEntry_eq_getEntry, #[``(List.getEntry_of_perm _)])⟩,
+     ⟨`getEntry!, (``getEntry!_eq_getEntry!, #[``(List.getEntry!_of_perm _)])⟩,
+     ⟨`getEntryD, (``getEntryD_eq_getEntryD, #[``(List.getEntryD_of_perm _)])⟩,
      ⟨`Const.get!, (``Const.get!_eq_getValue!, #[``(getValue!_of_perm _)])⟩,
      ⟨`Const.getD, (``Const.getD_eq_getValueD, #[``(getValueD_of_perm _)])⟩,
      ⟨`getKey?, (``getKey?_eq_getKey?, #[``(getKey?_of_perm _)])⟩,
@@ -3352,10 +3356,10 @@ theorem contains_of_contains_union!_of_contains_eq_false_left [TransOrd α]
 theorem union_insert_right_equiv_insert_union [TransOrd α] {p : (a : α) × β a}
     (h₁ : m₁.WF) (h₂ : m₂.WF) :
     Equiv (m₁.union (m₂.insert p.fst p.snd h₂.balanced).impl h₁.balanced h₂.insert.balanced) ((m₁.union m₂ h₁.balanced h₂.balanced).insert p.fst p.snd (@WF.union _ _ _ m₁ h₁ m₂ h₂).balanced).1 := by
-  simp_to_model [union, insert]
+  simp_to_model [Equiv, union, insert]
   apply List.Perm.trans
   . apply insertList_perm_of_perm_second
-    simp_to_model [insert]
+    simp_to_model [Equiv, insert]
     . apply insertEntry_of_perm
       . wf_trivial
       . apply List.Perm.refl
@@ -3400,7 +3404,7 @@ theorem union_equiv_congr_left {m₃ : Impl α β} [TransOrd α]
     (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (equiv : m₁.Equiv m₂) :
     (m₁.union m₃ h₁.balanced h₃.balanced).Equiv (m₂.union m₃ h₂.balanced h₃.balanced) := by
   revert equiv
-  simp_to_model [union]
+  simp_to_model [Equiv, union]
   intro equiv
   apply List.insertList_perm_of_perm_first equiv
   wf_trivial
@@ -3416,7 +3420,7 @@ theorem union_equiv_congr_right {m₃ : Impl α β} [TransOrd α]
     (h₁ : m₁.WF) (h₂ : m₂.WF) (h₃ : m₃.WF) (equiv : m₂.Equiv m₃) :
     (m₁.union m₂ h₁.balanced h₂.balanced).Equiv (m₁.union m₃ h₁.balanced h₃.balanced) := by
   revert equiv
-  simp_to_model [union]
+  simp_to_model [Equiv, union]
   intro equiv
   apply List.insertList_perm_of_perm_second equiv
   all_goals wf_trivial

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

@@ -330,6 +330,35 @@ Internal implementation detail of the tree map
 def getDₘ [Ord α] [OrientedOrd α] [LawfulEqOrd α] (k : α) (l : Impl α β) (fallback : β k) : β k :=
   get?ₘ l k |>.getD fallback
 
+/--
+Model implementation of the `getEntry?` function.
+Internal implementation detail of the tree map
+-/
+def getEntry?ₘ [Ord α] (l : Impl α β) (k : α) : Option ((a : α) × β a) :=
+  applyCell k l fun c _ => c.getEntry?
+
+/--
+Model implementation of the `getEntry` function.
+Internal implementation detail of the tree map
+-/
+def getEntryₘ [Ord α] (l : Impl α β) (k : α) (h : (getEntry?ₘ l k).isSome) :
+    (a : α) × β a :=
+  getEntry?ₘ l k |>.get h
+
+/--
+Model implementation of the `getEntry!` function.
+Internal implementation detail of the tree map
+-/
+def getEntry!ₘ [Ord α] [Inhabited ((a : α) × β a)] (l : Impl α β) (k : α) : (a : α) × β a :=
+  getEntry?ₘ l k |>.get!
+
+/--
+Model implementation of the `getEntryD` function.
+Internal implementation detail of the tree map
+-/
+def getEntryDₘ [Ord α] (k : α) (l : Impl α β) (fallback : (a : α) × β a) : (a : α) × β a :=
+  getEntry?ₘ l k |>.getD fallback
+
 /--
 Model implementation of the `getKey?` function.
 Internal implementation detail of the tree map
@@ -520,6 +549,42 @@ theorem getKey?_eq_getKey?ₘ [Ord α] (k : α) (l : Impl α β) :
     split <;> simp_all [Cell.getKey?, Cell.ofEq]
   · simp [getKey?, applyCell]
 
+theorem getEntry?_eq_getEntry?ₘ [Ord α] (k : α) (l : Impl α β) :
+    l.getEntry? k = l.getEntry?ₘ k := by
+  simp only [getEntry?ₘ]
+  induction l
+  · simp only [applyCell, getEntry?]
+    split <;> simp_all [Cell.getEntry?, Cell.ofEq]
+  · simp only [getEntry?, applyCell]; rfl
+
+theorem getEntry_eq_getEntry? [Ord α] (k : α) (l : Impl α β) {h} :
+    some (l.getEntry k h) = l.getEntry? k := by
+  induction l
+  · simp only [getEntry, getEntry?]
+    split <;> simp_all
+  · contradiction
+
+theorem getEntry_eq_getEntryₘ [Ord α] (k : α) (l : Impl α β) {h} (h') :
+    l.getEntry k h = l.getEntryₘ k h' := by
+  apply Option.some.inj
+  simp [getEntry_eq_getEntry?, getEntry?_eq_getEntry?ₘ, getEntryₘ]
+
+theorem getEntry!_eq_getEntry!ₘ [Ord α] [Inhabited ((a : α) × (β a))] (k : α) (l : Impl α β) :
+    l.getEntry! k = l.getEntry!ₘ k := by
+  simp only [getEntry!ₘ, getEntry?ₘ]
+  induction l
+  · simp only [applyCell, getEntry!]
+    split <;> simp_all [Cell.getEntry?, Cell.ofEq]
+  · simp only [getEntry!, applyCell]; rfl
+
+theorem getEntryD_eq_getEntryDₘ [Ord α] (k : α) (l : Impl α β)
+    (fallback : (a : α) × β a) : l.getEntryD k fallback = l.getEntryDₘ k fallback := by
+  simp only [getEntryDₘ, getEntry?ₘ]
+  induction l
+  · simp only [applyCell, getEntryD]
+    split <;> simp_all [Cell.getEntry?, Cell.ofEq]
+  · simp only [getEntryD, applyCell]; rfl
+
 theorem getKey_eq_getKey? [Ord α] (k : α) (l : Impl α β) {h} :
     some (l.getKey k h) = l.getKey? k := by
   induction l

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

@@ -116,6 +116,45 @@ def getD [Ord α] [LawfulEqOrd α] (t : Impl α β) (k : α) (fallback : β k) :
     | .gt => getD r k fallback
     | .eq => cast (congrArg β (compare_eq_iff_eq.mp h).symm) v'
 
+/-- Returns the entry (key-value pair) for the key `k`, or `none` if such a key does not exist. -/
+def getEntry? [Ord α] (t : Impl α β) (k : α) : Option ((a : α) × β a) :=
+  match t with
+  | .leaf => none
+  | .inner _ k' v' l r =>
+    match compare k k' with
+    | .lt => getEntry? l k
+    | .gt => getEntry? r k
+    | .eq => some ⟨k', v'⟩
+
+/-- Returns the entry (key-value pair) for the key `k`. -/
+def getEntry [Ord α] (t : Impl α β) (k : α) (hlk : k ∈ t) : (a : α) × β a :=
+  match t with
+  | .inner _ k' v' l r =>
+    match h : compare k k' with
+    | .lt => getEntry l k (by simpa [mem_iff_contains, contains, h] using hlk)
+    | .gt => getEntry r k (by simpa [mem_iff_contains, contains, h] using hlk)
+    | .eq => ⟨k', v'⟩
+
+/-- Returns the entry (key-value pair) for the key `k`, or panics if such a key does not exist. -/
+def getEntry! [Ord α] [Inhabited ((a : α) × β a)] (t : Impl α β) (k : α) : (a : α) × β a :=
+  match t with
+  | .leaf => panic! "Key is not present in map"
+  | .inner _ k' v' l r =>
+    match compare k k' with
+    | .lt => getEntry! l k
+    | .gt => getEntry! r k
+    | .eq => ⟨k', v'⟩
+
+/-- Returns the entry (key-value pair) for the key `k`, or `fallback` if such a key does not exist. -/
+def getEntryD [Ord α] (t : Impl α β) (k : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  match t with
+  | .leaf => fallback
+  | .inner _ k' v' l r =>
+    match compare k k' with
+    | .lt => getEntryD l k fallback
+    | .gt => getEntryD r k fallback
+    | .eq => ⟨k', v'⟩
+
 /-- Implementation detail of the tree map -/
 def getKey? [Ord α] (t : Impl α β) (k : α) : Option α :=
   match t with

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

@@ -613,6 +613,7 @@ theorem get?_eq_getValueCast? [instBEq : BEq α] [Ord α] [i : LawfulBEqOrd α]
     (hto : t.Ordered) : t.get? k = getValueCast? k t.toListModel := by
   rw [get?_eq_get?ₘ, get?ₘ_eq_getValueCast? hto]
 
+
 /-!
 ### `get`
 -/
@@ -724,6 +725,70 @@ theorem getKeyD_eq_getKeyD [Ord α] [TransOrd α] [instBEq : BEq α] [LawfulBEqO
     t.getKeyD k fallback = List.getKeyD k t.toListModel fallback := by
   rw [getKeyD_eq_getKeyDₘ, getKeyDₘ_eq_getKeyD hto]
 
+/-!
+### `getEntry?`
+-/
+
+theorem getEntry?ₘ_eq_getEntry? [Ord α] [TransOrd α] [BEq α] [LawfulBEqOrd α] {k : α} {t : Impl α β}
+    (hto : t.Ordered) : t.getEntry?ₘ k = List.getEntry? k t.toListModel := by
+  rw [getEntry?ₘ, applyCell_eq_apply_toListModel hto (fun l _ => List.getEntry? k l)]
+  · rintro ⟨(_|p), hp⟩ -
+    · simp [Cell.getEntry?]
+    · simp only [Cell.getEntry?, Option.toList_some, List.getEntry?]
+      simp [BEq.symm <| compare_eq_iff_beq.mp (hp p rfl)]
+  · exact fun l₁ l₂ h => List.getEntry?_of_perm
+  · exact fun l₁ l₂ h => List.getEntry?_append_of_containsKey_eq_false
+
+theorem getEntry?_eq_getEntry? [Ord α] [TransOrd α] [instBEq : BEq α] [LawfulBEqOrd α] {k : α} {t : Impl α β}
+    (hto : t.Ordered) : t.getEntry? k = List.getEntry? k t.toListModel := by
+  rw [getEntry?_eq_getEntry?ₘ, getEntry?ₘ_eq_getEntry? hto]
+
+/-!
+### `getEntry`
+-/
+
+theorem contains_eq_isSome_getEntry?ₘ [Ord α] [TransOrd α] [BEq α] [LawfulBEqOrd α] {k : α}
+    {t : Impl α β} (hto : t.Ordered) : contains k t = (t.getEntry?ₘ k).isSome := by
+  rw [getEntry?ₘ_eq_getEntry? hto, contains_eq_containsKey hto, containsKey_eq_isSome_getEntry?]
+
+theorem getEntryₘ_eq_getEntry [Ord α] [TransOrd α] [BEq α] [LawfulBEqOrd α] {k : α} {t : Impl α β} (h)
+    {h'} (hto : t.Ordered) : t.getEntryₘ k h' = List.getEntry k t.toListModel h := by
+  simp only [getEntryₘ]
+  revert h'
+  rw [getEntry?ₘ_eq_getEntry? hto]
+  simp [getEntry?_eq_some_getEntry ‹_›]
+
+theorem getEntry_eq_getEntry [Ord α] [TransOrd α] [instBEq : BEq α] [LawfulBEqOrd α] {k : α} {t : Impl α β} {h}
+    (hto : t.Ordered): t.getEntry k h = List.getEntry k t.toListModel (contains_eq_containsKey hto ▸ h) := by
+  rw [getEntry_eq_getEntryₘ, getEntryₘ_eq_getEntry _ hto]
+  exact contains_eq_isSome_getEntry?ₘ hto ▸ h
+
+/-!
+### `getEntry!`
+-/
+
+theorem getEntry!ₘ_eq_getEntry! [Ord α] [TransOrd α] [BEq α] [LawfulBEqOrd α] [Inhabited ((a : α) × β a)] {k : α}
+    {t : Impl α β} (hto : t.Ordered) : t.getEntry!ₘ k = List.getEntry! k t.toListModel := by
+  simp [getEntry!ₘ, getEntry?ₘ_eq_getEntry? hto, getEntry!_eq_getEntry?]
+
+theorem getEntry!_eq_getEntry! [Ord α] [TransOrd α] [instBEq : BEq α] [LawfulBEqOrd α] [Inhabited ((a : α) × β a)] {k : α}
+    {t : Impl α β} (hto : t.Ordered) : t.getEntry! k = List.getEntry! k t.toListModel := by
+  rw [getEntry!_eq_getEntry!ₘ, getEntry!ₘ_eq_getEntry! hto]
+
+/-!
+### `getEntryD`
+-/
+
+theorem getEntryDₘ_eq_getEntryD [Ord α] [TransOrd α] [BEq α] [LawfulBEqOrd α] {k : α}
+    {t : Impl α β} {fallback : (a : α) × β a} (hto : t.Ordered) :
+    t.getEntryDₘ k fallback = List.getEntryD k fallback t.toListModel := by
+  simp [getEntryDₘ, getEntry?ₘ_eq_getEntry? hto, getEntryD_eq_getEntry?]
+
+theorem getEntryD_eq_getEntryD [Ord α] [TransOrd α] [instBEq : BEq α] [LawfulBEqOrd α] {k : α}
+    {t : Impl α β} {fallback : (a : α) × β a} (hto : t.Ordered) :
+    t.getEntryD k fallback = List.getEntryD k fallback t.toListModel := by
+  rw [getEntryD_eq_getEntryDₘ, getEntryDₘ_eq_getEntryD hto]
+
 namespace Const
 
 variable {β : Type v}

src/Std/Data/DTreeMap/Raw/Basic.lean

@@ -182,6 +182,22 @@ def get! [LawfulEqCmp cmp] (t : Raw α β cmp) (a : α) [Inhabited (β a)]  : β
 def getD [LawfulEqCmp cmp] (t : Raw α β cmp) (a : α) (fallback : β a) : β a :=
   letI : Ord α := ⟨cmp⟩; t.inner.getD a fallback
 
+@[inline, inherit_doc DTreeMap.getEntry?]
+def getEntry? (t : Raw α β cmp) (a : α) : Option ((a : α) × β a) :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry? a
+
+@[inline, inherit_doc DTreeMap.getEntry]
+def getEntry [LawfulEqCmp cmp] (t : Raw α β cmp) (a : α) (h : a ∈ t) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry a h
+
+@[inline, inherit_doc DTreeMap.getEntry!]
+def getEntry! [Inhabited ((a : α) × β a)](t : Raw α β cmp) (a : α) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntry! a
+
+@[inline, inherit_doc DTreeMap.getEntryD]
+def getEntryD (t : Raw α β cmp) (a : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  letI : Ord α := ⟨cmp⟩; t.inner.getEntryD a fallback
+
 @[inline, inherit_doc DTreeMap.getKey?]
 def getKey? (t : Raw α β cmp) (a : α) : Option α :=
   letI : Ord α := ⟨cmp⟩; t.inner.getKey? a