From 36ee331ce2d0cb0cb2603c1dff54dea0fdd23d95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Wojciech=20R=C3=B3=C5=BCowski?= Date: Fri, 14 Nov 2025 09:09:53 +0000 Subject: [PATCH] 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 | 32 +++++++++ src/Std/Data/DTreeMap/Internal/Cell.lean | 11 ++++ src/Std/Data/DTreeMap/Internal/Lemmas.lean | 12 ++-- src/Std/Data/DTreeMap/Internal/Model.lean | 65 +++++++++++++++++++ src/Std/Data/DTreeMap/Internal/Queries.lean | 39 +++++++++++ src/Std/Data/DTreeMap/Internal/WF/Lemmas.lean | 65 +++++++++++++++++++ src/Std/Data/DTreeMap/Raw/Basic.lean | 16 +++++ 7 files changed, 236 insertions(+), 4 deletions(-) diff --git a/src/Std/Data/DTreeMap/Basic.lean b/src/Std/Data/DTreeMap/Basic.lean index 11691af0c5..1f07cc7386 100644 --- a/src/Std/Data/DTreeMap/Basic.lean +++ b/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. diff --git a/src/Std/Data/DTreeMap/Internal/Cell.lean b/src/Std/Data/DTreeMap/Internal/Cell.lean index 2be4d406ba..44f42f4530 100644 --- a/src/Std/Data/DTreeMap/Internal/Cell.lean +++ b/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 diff --git a/src/Std/Data/DTreeMap/Internal/Lemmas.lean b/src/Std/Data/DTreeMap/Internal/Lemmas.lean index 8dc647b02e..f54cdbbb01 100644 --- a/src/Std/Data/DTreeMap/Internal/Lemmas.lean +++ b/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 diff --git a/src/Std/Data/DTreeMap/Internal/Model.lean b/src/Std/Data/DTreeMap/Internal/Model.lean index 26605679dc..666d34a930 100644 --- a/src/Std/Data/DTreeMap/Internal/Model.lean +++ b/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 diff --git a/src/Std/Data/DTreeMap/Internal/Queries.lean b/src/Std/Data/DTreeMap/Internal/Queries.lean index 3bf013c47d..8936025576 100644 --- a/src/Std/Data/DTreeMap/Internal/Queries.lean +++ b/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 diff --git a/src/Std/Data/DTreeMap/Internal/WF/Lemmas.lean b/src/Std/Data/DTreeMap/Internal/WF/Lemmas.lean index c4bee9af7d..72e338ac4d 100644 --- a/src/Std/Data/DTreeMap/Internal/WF/Lemmas.lean +++ b/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} diff --git a/src/Std/Data/DTreeMap/Raw/Basic.lean b/src/Std/Data/DTreeMap/Raw/Basic.lean index 398261a199..072901cbeb 100644 --- a/src/Std/Data/DTreeMap/Raw/Basic.lean +++ b/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