feat: (DHashMap|HashMap|HashSet).(getKey?|getKey|getKey!|getKeyD) (leanprover#5244)

Author: hargoniX

src/Std/Data/DHashMap/Basic.lean

@@ -152,6 +152,18 @@ variable {β : Type v}
 
 end
 
+@[inline, inherit_doc Raw.getKey?] def getKey? (m : DHashMap α β) (a : α) : Option α :=
+  Raw₀.getKey? ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[inline, inherit_doc Raw.getKey] def getKey (m : DHashMap α β) (a : α) (h : a ∈ m) : α :=
+  Raw₀.getKey ⟨m.1, m.2.size_buckets_pos⟩ a h
+
+@[inline, inherit_doc Raw.getKey!] def getKey! [Inhabited α] (m : DHashMap α β) (a : α) : α :=
+  Raw₀.getKey! ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[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.size] def size (m : DHashMap α β) : Nat :=
   m.1.size
 

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

@@ -102,16 +102,31 @@ 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 getKey [BEq α] (a : α) : (l : AssocList α β) → l.contains a → α
+  | cons k _ es, h => if hka : k == a then k
+      else es.getKey a (by rw [← h, contains, Bool.of_not_eq_true hka, Bool.false_or])
+
 /-- Internal implementation detail of the hash map -/
 def getCast! [BEq α] [LawfulBEq α] (a : α) [Inhabited (β a)] : AssocList α β → β a
   | nil => panic! "key is not present in hash table"
   | cons k v es => if h : k == a then cast (congrArg β (eq_of_beq h)) v else es.getCast! a
 
+/-- Internal implementation detail of the hash map -/
+def getKey? [BEq α] (a : α) : AssocList α β → Option α
+  | nil => none
+  | cons k _ es => if k == a then some k else es.getKey? a
+
 /-- Internal implementation detail of the hash map -/
 def get! {β : Type v} [BEq α] [Inhabited β] (a : α) : AssocList α (fun _ => β) → β
   | nil => panic! "key is not present in hash table"
   | cons k v es => bif k == a then v else es.get! a
 
+/-- Internal implementation detail of the hash map -/
+def getKey! [BEq α] [Inhabited α] (a : α) : AssocList α β → α
+  | nil => panic! "key is not present in hash table"
+  | cons k _ es => if k == a then k else es.getKey! a
+
 /-- Internal implementation detail of the hash map -/
 def getCastD [BEq α] [LawfulBEq α] (a : α) (fallback : β a) : AssocList α β → β a
   | nil => fallback
@@ -123,6 +138,11 @@ def getD {β : Type v} [BEq α] (a : α) (fallback : β) : AssocList α (fun _ =
   | nil => fallback
   | cons k v es => bif k == a then v else es.getD a fallback
 
+/-- Internal implementation detail of the hash map -/
+def getKeyD [BEq α] (a : α) (fallback : α) : AssocList α β → α
+  | nil => fallback
+  | cons k _ es => if k == a then k else es.getKeyD a fallback
+
 /-- Internal implementation detail of the hash map -/
 def replace [BEq α] (a : α) (b : β a) : AssocList α β → AssocList α β
   | nil => nil

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

@@ -112,6 +112,32 @@ theorem get!_eq {β : Type v} [BEq α] [Inhabited β] {l : AssocList α (fun _ =
   · simp_all [get!, List.getValue!, List.getValue!, List.getValue?_cons,
       Bool.apply_cond Option.get!]
 
+@[simp]
+theorem getKey?_eq [BEq α] {l : AssocList α β} {a : α} :
+    l.getKey? a = List.getKey? a l.toList := by
+  induction l <;> simp_all [getKey?]
+
+@[simp]
+theorem getKey_eq [BEq α] {l : AssocList α β} {a : α} {h} :
+    l.getKey a h = List.getKey a l.toList (contains_eq.symm.trans h) := by
+  induction l
+  · simp [contains] at h
+  · next k v t ih => simp only [getKey, toList_cons, List.getKey_cons, ih]
+
+@[simp]
+theorem getKeyD_eq [BEq α] {l : AssocList α β} {a fallback : α} :
+    l.getKeyD a fallback = List.getKeyD a l.toList fallback := by
+  induction l
+  · simp [getKeyD, List.getKeyD]
+  · simp_all [getKeyD, List.getKeyD, Bool.apply_cond (fun x => Option.getD x fallback)]
+
+@[simp]
+theorem getKey!_eq [BEq α] [Inhabited α] {l : AssocList α β} {a : α} :
+    l.getKey! a = List.getKey! a l.toList := by
+  induction l
+  · simp [getKey!, List.getKey!]
+  · simp_all [getKey!, List.getKey!, Bool.apply_cond Option.get!]
+
 @[simp]
 theorem toList_replace [BEq α] {l : AssocList α β} {a : α} {b : β a} :
     (l.replace a b).toList = replaceEntry a b l.toList := by

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

@@ -100,9 +100,9 @@ Here is a summary of the steps required to add and verify a new operation:
   * Connect the implementation on lists and associative lists in `Internal.AssocList.Lemmas` via a
     lemma `AssocList.operation_eq`.
 3. Write the model implementation
-  * Write the model implementation `Raw₀.operationₘ` in `Internal.List.Model`
+  * Write the model implementation `Raw₀.operationₘ` in `Internal.Model`
   * Prove that the model implementation is equal to the actual implementation in
-    `Internal.List.Model` via a lemma `operation_eq_operationₘ`.
+    `Internal.Model` via a lemma `operation_eq_operationₘ`.
 4. Verify the model implementation
   * In `Internal.WF`, prove `operationₘ_eq_List.operation` (for access operations) or
     `wfImp_operationₘ` and `toListModel_operationₘ`
@@ -121,18 +121,18 @@ Here is a summary of the steps required to add and verify a new operation:
     might also have to prove that your list operation is invariant under permutation and add that to
     the tactic.
 7. State and prove the user-facing lemmas
-  * Restate all of your lemmas for `DHashMap.Raw` in `DHashMap.Lemmas` and prove them using the
+  * Restate all of your lemmas for `DHashMap.Raw` in `DHashMap.RawLemmas` and prove them using the
     provided tactic after hooking in your `operation_eq` and `operation_val` from step 5.
   * Restate all of your lemmas for `DHashMap` in `DHashMap.Lemmas` and prove them by reducing to
     `Raw₀`.
-  * Restate all of your lemmas for `HashMap.Raw` in `HashMap.Lemmas` and prove them by reducing to
+  * Restate all of your lemmas for `HashMap.Raw` in `HashMap.RawLemmas` and prove them by reducing to
     `DHashMap.Raw`.
   * Restate all of your lemmas for `HashMap` in `HashMap.Lemmas` and prove them by reducing to
     `DHashMap`.
-  * Restate all of your lemmas for `HashSet.Raw` in `HashSet.Lemmas` and prove them by reducing to
-    `DHashSet.Raw`.
+  * Restate all of your lemmas for `HashSet.Raw` in `HashSet.RawLemmas` and prove them by reducing to
+    `HashMap.Raw`.
   * Restate all of your lemmas for `HashSet` in `HashSet.Lemmas` and prove them by reducing to
-    `DHashSet`.
+    `HashMap`.
 
 This sounds like a lot of work (and it is if you have to add a lot of user-facing lemmas), but the
 framework is set up in such a way that each step is really easy and the proofs are all really short
@@ -420,6 +420,30 @@ variable {β : Type v}
 
 end
 
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey? [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) : Option α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let ⟨i, h⟩ := mkIdx buckets.size h (hash a)
+  buckets[i].getKey? a
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (hma : m.contains a) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKey a hma
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKeyD [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : α) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKeyD a fallback
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey! [BEq α] [Hashable α] [Inhabited α] (m : Raw₀ α β) (a : α) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKey! a
+
 end Raw₀
 
 /-- Internal implementation detail of the hash map -/