feat: add minimal support for getEntry/getEntry?/getEntry!/getEntryD for DHashMap (#11076)

This PR adds `getEntry`/`getEntry?`/`getEntry!`/`getEntryD` operation on
DHashMap.

Author: wkrozowski

src/Std/Data/DHashMap/Basic.lean

@@ -190,6 +190,18 @@ end
 @[inline, inherit_doc Raw.getKeyD] def getKeyD (m : DHashMap α β) (a : α) (fallback : α) : α :=
   Raw₀.getKeyD ⟨m.1, m.2.size_buckets_pos⟩ a fallback
 
+@[inline, inherit_doc Raw.getEntry?] def getEntry? (m : DHashMap α β) (a : α) : Option ((a : α) × β a) :=
+  Raw₀.getEntry? ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[inline, inherit_doc Raw.getEntry] def getEntry (m : DHashMap α β) (a : α) (h : a ∈ m) : (a : α) × β a :=
+  Raw₀.getEntry ⟨m.1, m.2.size_buckets_pos⟩ a h
+
+@[inline, inherit_doc Raw.getEntry!] def getEntry! [Inhabited ((a : α) × β a)] (m : DHashMap α β) (a : α) : (a : α) × β a :=
+  Raw₀.getEntry! ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[inline, inherit_doc Raw.getEntryD] def getEntryD (m : DHashMap α β) (a : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  Raw₀.getEntryD ⟨m.1, m.2.size_buckets_pos⟩ a fallback
+
 @[inline, inherit_doc Raw.size] def size (m : DHashMap α β) : Nat :=
   m.1.size
 

src/Std/Data/DHashMap/Internal/AssocList/Basic.lean

@@ -107,6 +107,12 @@ def getCast? [BEq α] [LawfulBEq α] (a : α) : AssocList α β → Option (β a
   | cons k v es => if h : k == a then some (cast (congrArg β (eq_of_beq h)) v)
       else es.getCast? a
 
+/-- Internal implementation detail of the hash map -/
+def getEntry? [BEq α] (a : α) : (l : AssocList α β) → Option ((a : α) × β a)
+  | nil => none
+  | cons k v es => if k == a then some ⟨k, v⟩
+      else es.getEntry? a
+
 /-- Internal implementation detail of the hash map -/
 def contains [BEq α] (a : α) : AssocList α β → Bool
   | nil => false
@@ -122,6 +128,21 @@ def getCast [BEq α] [LawfulBEq α] (a : α) : (l : AssocList α β) → l.conta
   | cons k v es, h => if hka : k == a then cast (congrArg β (eq_of_beq hka)) v
       else es.getCast a (by rw [← h, contains, Bool.of_not_eq_true hka, Bool.false_or])
 
+/-- Internal implementation detail of the hash map -/
+def getEntry [BEq α] (a : α) : (l : AssocList α β) → l.contains a → (a : α) × β a
+  | cons k v es, h => if hka : k == a then ⟨k, v⟩
+      else es.getEntry a (by rw [← h, contains, Bool.of_not_eq_true hka, Bool.false_or])
+
+/-- Internal implementation detail of the hash map -/
+def getEntryD [BEq α] (a : α) (fallback : (a : α) × β a) : AssocList α β → (a : α) × β a
+  | nil => fallback
+  | cons k v es => if k == a then ⟨k, v⟩ else es.getEntryD a fallback
+
+/-- Internal implementation detail of the hash map -/
+def getEntry! [BEq α] (a : α) [Inhabited ((a : α) × β a)] : AssocList α β → (a : α) × β a
+  | nil => default
+  | cons k v es => if k == a then ⟨k, v⟩ else es.getEntry! a
+
 /-- Internal implementation detail of the hash map -/
 def getKey [BEq α] (a : α) : (l : AssocList α β) → l.contains a → α
   | cons k _ es, h => if hka : k == a then k

src/Std/Data/DHashMap/Internal/AssocList/Lemmas.lean

@@ -82,6 +82,38 @@ theorem get_eq {β : Type v} [BEq α] {l : AssocList α (fun _ => β)} {a : α}
   · simp [contains] at h
   next k v t ih => simp only [get, toList_cons, List.getValue_cons, ih]
 
+@[simp]
+theorem getEntry_eq [BEq α] {l : AssocList α β} {a : α} {h} :
+    l.getEntry a h = List.getEntry a l.toList (contains_eq.symm.trans h) := by
+  induction l
+  · simp [contains] at h
+  next k v t ih =>
+    simp only [getEntry, toList_cons, List.getEntry_cons, ih]
+
+@[simp]
+theorem getEntry?_eq [BEq α] {l : AssocList α β} {a : α} :
+    l.getEntry? a = List.getEntry? a l.toList  := by
+  induction l
+  · simp only [getEntry?, toList_nil, getEntry?_nil]
+  next k v t ih =>
+    simp only [getEntry?, ih, toList_cons, getEntry?_cons, Bool.ite_eq_cond_iff]
+
+@[simp]
+theorem getEntryD_eq [BEq α] {l : AssocList α β} {a : α} {fallback : (a : α) × β a} :
+    l.getEntryD a fallback = List.getEntryD a fallback l.toList := by
+  induction l
+  · simp only [getEntryD, toList_nil, getEntryD_nil]
+  next k v t ih =>
+    simp only [getEntryD, ih, toList_cons, getEntryD_cons, Bool.ite_eq_cond_iff]
+
+@[simp]
+theorem getEntry!_eq [BEq α] {l : AssocList α β} {a : α} [Inhabited ((a : α) × β a)] :
+    l.getEntry! a = List.getEntry! a l.toList := by
+  induction l
+  · simp only [getEntry!, toList_nil, getEntry!_nil]
+  next k v t ih =>
+    simp only [getEntry!, ih, toList_cons, List.getEntry!_cons, Bool.ite_eq_cond_iff]
+
 @[simp]
 theorem getCastD_eq [BEq α] [LawfulBEq α] {l : AssocList α β} {a : α} {fallback : β a} :
     l.getCastD a fallback = getValueCastD a l.toList fallback := by

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

@@ -368,6 +368,34 @@ def get [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (hma :
   let idx := mkIdx buckets.size h (hash a)
   buckets[idx.1].getCast a hma
 
+/-- Internal implementation detail of the hash map -/
+def getEntry [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (hma : m.contains a) :
+    (a : α) × β a :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getEntry a hma
+
+/-- Internal implementation detail of the hash map -/
+def getEntry? [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) :
+    Option ((a : α) × β a) :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getEntry? a
+
+/-- Internal implementation detail of the hash map -/
+def getEntryD [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : (a : α) × β a) :
+    (a : α) × β a :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getEntryD a fallback
+
+/-- Internal implementation detail of the hash map -/
+def getEntry! [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) [Inhabited ((a : α) × β a)] :
+    (a : α) × β a :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getEntry! a
+
 /-- Internal implementation detail of the hash map -/
 def getD [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : β a) :
     β a :=

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

@@ -294,6 +294,22 @@ def containsₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) : Bool :=
 def getₘ [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (h : m.containsₘ a) : β a :=
   (bucket m.1.buckets m.2 a).getCast a h
 
+/-- Internal implementation detail of the hash map -/
+def getEntryₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (h : m.containsₘ a) : (a : α) × β a :=
+  (bucket m.1.buckets m.2 a).getEntry a h
+
+/-- Internal implementation detail of the hash map -/
+def getEntry?ₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) : Option ((a : α) × β a) :=
+  (bucket m.1.buckets m.2 a).getEntry? a
+
+/-- Internal implementation detail of the hash map -/
+def getEntryDₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  (bucket m.1.buckets m.2 a).getEntryD a fallback
+
+/-- Internal implementation detail of the hash map -/
+def getEntry!ₘ [BEq α] [Hashable α] [Inhabited ((a : α) × β a)] (m : Raw₀ α β) (a : α) : (a : α) × β a :=
+  (bucket m.1.buckets m.2 a).getEntry! a
+
 /-- Internal implementation detail of the hash map -/
 def getDₘ [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : β a) : β a :=
   (m.get?ₘ a).getD fallback
@@ -452,6 +468,18 @@ theorem get?_eq_get?ₘ [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β)
 theorem get_eq_getₘ [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (h : m.contains a) :
     get m a h = getₘ m a (by exact h) := (rfl)
 
+theorem getEntry_eq_getEntryₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (h : m.contains a) :
+    getEntry m a h = getEntryₘ m a (by exact h) := (rfl)
+
+theorem getEntry?_eq_getEntry?ₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) :
+    getEntry? m a = getEntry?ₘ m a := (rfl)
+
+theorem getEntryD_eq_getEntryDₘ [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : (a : α) × β a) :
+    getEntryD m a fallback = getEntryDₘ m a fallback := (rfl)
+
+theorem getEntry!_eq_getEntry!ₘ [BEq α] [Hashable α] [Inhabited ((a : α) × β a)] (m : Raw₀ α β) (a : α) :
+    getEntry! m a = getEntry!ₘ m a := (rfl)
+
 theorem getD_eq_getDₘ [BEq α] [LawfulBEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : β a) :
     getD m a fallback = getDₘ m a fallback := by
   simp [getD, getDₘ, get?ₘ, List.getValueCastD_eq_getValueCast?, bucket]

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

@@ -140,6 +140,10 @@ private meta def queryMap : Std.DHashMap Name (fun _ => Name × Array (MacroM (T
      ⟨`getKey, (``getKey_eq_getKey, #[`(getKey_of_perm _)])⟩,
      ⟨`getKeyD, (``getKeyD_eq_getKeyD, #[`(getKeyD_of_perm _)])⟩,
      ⟨`getKey!, (``getKey!_eq_getKey!, #[`(getKey!_of_perm _)])⟩,
+     ⟨`getEntry, (``getEntry_eq_getEntry, #[`(getEntry_of_perm _)])⟩,
+     ⟨`getEntry?, (``getEntry?_eq_getEntry?, #[`(getEntry?_of_perm _)])⟩,
+     ⟨`getEntryD, (``getEntryD_eq_getEntryD, #[`(getEntryD_of_perm _)])⟩,
+     ⟨`getEntry!, (``getEntry!_eq_getEntry!, #[`(getEntry!_of_perm _)])⟩,
      ⟨`toList, (``Raw.toList_eq_toListModel, #[])⟩,
      ⟨`keys, (``Raw.keys_eq_keys_toListModel, #[`(perm_keys_congr_left)])⟩,
      ⟨`Const.toList, (``Raw.Const.toList_eq_toListModel_map, #[`(perm_map_congr_left)])⟩,
@@ -2424,7 +2428,7 @@ theorem union_equiv_congr_left {m₃ : Raw₀ α β} [EquivBEq α] [LawfulHashab
     (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) (h₃ : m₃.val.WF) (equiv : m₁.1.Equiv m₂.1) :
     (m₁.union m₃).1.Equiv (m₂.union m₃).1 := by
   revert equiv
-  simp_to_model [union]
+  simp_to_model [Equiv, union]
   intro equiv
   apply List.insertList_perm_of_perm_first equiv
   wf_trivial
@@ -2433,18 +2437,18 @@ theorem union_equiv_congr_right {m₃ : Raw₀ α β} [EquivBEq α] [LawfulHasha
     (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) (h₃ : m₃.val.WF) (equiv : m₂.1.Equiv m₃.1) :
     (m₁.union m₂).1.Equiv (m₁.union m₃).1 := by
   revert equiv
-  simp_to_model [union]
+  simp_to_model [Equiv, union]
   intro equiv
   apply @List.insertList_perm_of_perm_second _ _ _ _ (toListModel m₂.val.buckets) (toListModel m₃.val.buckets) (toListModel m₁.val.buckets) equiv
   all_goals wf_trivial
 
 theorem union_insert_right_equiv_insert_union [EquivBEq α] [LawfulHashable α] {p : (a : α) × β a}
     (h₁ : m₁.val.WF) (h₂ : m₂.val.WF) :
     (m₁.union (m₂.insert p.fst p.snd)).1.Equiv ((m₁.union m₂).insert p.fst p.snd).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

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

@@ -405,11 +405,51 @@ theorem getₘ_eq_getValue [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α
   apply_bucket_with_proof hm a AssocList.getCast List.getValueCast AssocList.getCast_eq
     List.getValueCast_of_perm List.getValueCast_append_of_containsKey_eq_false
 
+theorem getEntryₘ_eq_getEntry [BEq α] [PartialEquivBEq α] [Hashable α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} {h : m.containsₘ a} :
+    m.getEntryₘ a h = List.getEntry a (toListModel m.1.buckets) (containsₘ_eq_containsKey hm ▸ h) :=
+  apply_bucket_with_proof hm a AssocList.getEntry List.getEntry AssocList.getEntry_eq getEntry_of_perm getEntry_append_of_containsKey_eq_false
+
+theorem getEntry?ₘ_eq_getEntry? [BEq α] [PartialEquivBEq α] [Hashable α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} :
+    m.getEntry?ₘ a = List.getEntry? a (toListModel m.1.buckets) :=
+  apply_bucket hm AssocList.getEntry?_eq getEntry?_of_perm getEntry?_append_of_containsKey_eq_false
+
 theorem get_eq_getValueCast [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
     {a : α} {h : m.contains a} :
     m.get a h = getValueCast a (toListModel m.1.buckets) (contains_eq_containsKey hm ▸ h) := by
   rw [get_eq_getₘ, getₘ_eq_getValue hm]
 
+theorem getEntry_eq_getEntry [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} {h : m.contains a} :
+    m.getEntry a h = List.getEntry a (toListModel m.1.buckets) (contains_eq_containsKey hm ▸ h) := by
+  rw [getEntry_eq_getEntryₘ, getEntryₘ_eq_getEntry hm]
+
+theorem getEntry?_eq_getEntry? [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} :
+    m.getEntry? a = List.getEntry? a (toListModel m.1.buckets) := by
+  rw [getEntry?_eq_getEntry?ₘ, getEntry?ₘ_eq_getEntry? hm]
+
+theorem getEntryDₘ_eq_getEntryD [BEq α] [PartialEquivBEq α] [Hashable α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} {fallback : (a : α) × β a} :
+    m.getEntryDₘ a fallback = List.getEntryD a fallback (toListModel m.1.buckets) :=
+  apply_bucket hm AssocList.getEntryD_eq getEntryD_of_perm getEntryD_append_of_containsKey_eq_false
+
+theorem getEntryD_eq_getEntryD [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} {fallback : (a : α) × β a} :
+    m.getEntryD a fallback = List.getEntryD a fallback (toListModel m.1.buckets) := by
+  rw [getEntryD_eq_getEntryDₘ, getEntryDₘ_eq_getEntryD hm]
+
+theorem getEntry!ₘ_eq_getEntry! [BEq α] [PartialEquivBEq α] [Hashable α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} [Inhabited ((a : α) × β a)] :
+    m.getEntry!ₘ a = List.getEntry! a (toListModel m.1.buckets) :=
+  apply_bucket hm AssocList.getEntry!_eq getEntry!_of_perm getEntry!_append_of_containsKey_eq_false
+
+theorem getEntry!_eq_getEntry! [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α] {m : Raw₀ α β} (hm : Raw.WFImp m.1)
+    {a : α} [Inhabited ((a : α) × β a)] :
+    m.getEntry! a = List.getEntry! a (toListModel m.1.buckets) := by
+  rw [getEntry!_eq_getEntry!ₘ, getEntry!ₘ_eq_getEntry! hm]
+
 theorem get!ₘ_eq_getValueCast! [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α β}
     (hm : Raw.WFImp m.1) {a : α} [Inhabited (β a)] :
     m.get!ₘ a = getValueCast! a (toListModel m.1.buckets) := by

src/Std/Data/DHashMap/Raw.lean

@@ -300,6 +300,41 @@ If no panic occurs the result is guaranteed to be pointer equal to the key in th
     Raw₀.getKey! ⟨m, h⟩ a
   else default -- will never happen for well-formed inputs
 
+/--
+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? [BEq α] [Hashable α] (m : Raw α β) (a : α) : Option ((a : α) × β a) :=
+  if h : 0 < m.buckets.size then
+    Raw₀.getEntry? ⟨m, h⟩ a
+  else none -- will never  happen for well-formed inputs
+
+/--
+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 [BEq α] [Hashable α] (m : Raw α β) (a : α) (h : a ∈ m) : (a : α) × β a :=
+  Raw₀.getEntry ⟨m, by change dite .. = true at h; split at h <;> simp_all⟩ a
+    (by change dite .. = true at h; split at h <;> simp_all)
+
+/--
+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 be `BEq` to the input `a`.
+-/
+@[inline] def getEntryD [BEq α] [Hashable α] (m : Raw α β) (a : α) (fallback : (a : α) × β a) : (a : α) × β a :=
+  if h : 0 < m.buckets.size then
+    Raw₀.getEntryD ⟨m, h⟩ a fallback
+  else fallback -- will never happen for well-formed inputs
+
+/--
+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! [BEq α] [Hashable α] [Inhabited ((a : α) × β a)] (m : Raw α β) (a : α) : (a : α) × β a :=
+  if h : 0 < m.buckets.size then
+    Raw₀.getEntry! ⟨m, h⟩ a
+  else default -- will never happen for well-formed inputs
+
 /--
 Returns `true` if the hash map contains no mappings.