feat: verify keys method on HashMaps

src/Init/Data/List/Lemmas.lean

@@ -2843,6 +2843,10 @@ theorem contains_iff_exists_mem_beq [BEq α] {l : List α} {a : α} :
     l.contains a ↔ ∃ a' ∈ l, a == a' := by
   induction l <;> simp_all
 
+theorem List.contains_iff_mem [BEq α] [LawfulBEq α] {l : List α} {a : α} :
+    l.contains a ↔ a ∈ l := by
+  simp
+
 /-! ## Sublists -/
 
 /-! ### partition

src/Init/Data/List/Perm.lean

@@ -114,6 +114,14 @@ theorem Perm.length_eq {l₁ l₂ : List α} (p : l₁ ~ l₂) : length l₁ = l
   | swap => rfl
   | trans _ _ ih₁ ih₂ => simp only [ih₁, ih₂]
 
+theorem Perm.contains_eq [BEq α] {l₁ l₂ : List α} (h : l₁ ~ l₂) {a : α} :
+    l₁.contains a = l₂.contains a := by
+  induction h with
+  | nil => rfl
+  | cons => simp_all
+  | swap => simp only [contains_cons, ← Bool.or_assoc, Bool.or_comm]
+  | trans => simp_all
+
 theorem Perm.eq_nil {l : List α} (p : l ~ []) : l = [] := eq_nil_of_length_eq_zero p.length_eq
 
 theorem Perm.nil_eq {l : List α} (p : [] ~ l) : [] = l := p.symm.eq_nil.symm

src/Std/Data/DHashMap/Internal/List/Associative.lean

@@ -849,6 +849,12 @@ theorem isEmpty_eraseKey [BEq α] {l : List ((a : α) × β a)} {k : α} :
 theorem keys_eq_map (l : List ((a : α) × β a)) : keys l = l.map (·.1) := by
   induction l using assoc_induction <;> simp_all
 
+theorem length_keys_eq_length (l : List ((a : α) × β a)) : (keys l).length = l.length := by 
+  induction l using assoc_induction <;> simp_all
+
+theorem isEmpty_keys_eq_isEmpty (l : List ((a : α) × β a)) : (keys l).isEmpty = l.isEmpty := by 
+  induction l using assoc_induction <;> simp_all
+
 theorem containsKey_eq_keys_contains [BEq α] [PartialEquivBEq α] {l : List ((a : α) × β a)}
     {a : α} : containsKey a l = (keys l).contains a := by
   induction l using assoc_induction

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

@@ -84,7 +84,10 @@ private def queryNames : Array Name :=
     ``get?_eq_getValueCast?, ``Const.get?_eq_getValue?, ``get_eq_getValueCast,
     ``Const.get_eq_getValue, ``get!_eq_getValueCast!, ``getD_eq_getValueCastD,
     ``Const.get!_eq_getValue!, ``Const.getD_eq_getValueD, ``getKey?_eq_getKey?,
-    ``getKey_eq_getKey, ``getKeyD_eq_getKeyD, ``getKey!_eq_getKey!]
+    ``getKey_eq_getKey, ``getKeyD_eq_getKeyD, ``getKey!_eq_getKey!,
+    ``Raw.length_keys_eq_length_keys, ``Raw.isEmpty_keys_eq_isEmpty_keys,
+    ``Raw.contains_keys_eq_contains_keys, ``Raw.mem_keys_iff_contains_keys,
+    ``Raw.pairwise_keys_iff_pairwise_keys]
 
 private def modifyNames : Array Name :=
   #[``toListModel_insert, ``toListModel_erase, ``toListModel_insertIfNew]
@@ -811,6 +814,31 @@ theorem getThenInsertIfNew?_snd {k : α} {v : β} :
 
 end Const
 
+@[simp]
+theorem length_keys_eq_size [EquivBEq α] [LawfulHashable α] (h : m.1.WF) :
+    m.1.keys.length = m.1.size := by
+  simp_to_model using List.length_keys_eq_length
+
+@[simp]
+theorem contains_keys [EquivBEq α] [LawfulHashable α] (h : m.1.WF) {k : α} :
+    m.1.keys.contains k = m.contains k := by
+  simp_to_model using List.containsKey_eq_keys_contains.symm
+
+@[simp]
+theorem isEmpty_keys_eq_isEmpty [EquivBEq α] [LawfulHashable α] (h: m.1.WF):
+    m.1.keys.isEmpty = m.1.isEmpty:= by
+  simp_to_model using List.isEmpty_keys_eq_isEmpty
+
+@[simp]
+theorem mem_keys_iff_contains [LawfulBEq α] [LawfulHashable α] (h : m.1.WF) {k : α} :
+    k ∈ m.1.keys ↔ m.contains k := by
+  simp_to_model 
+  rw [List.containsKey_eq_keys_contains]
+
+theorem distinct_keys [EquivBEq α] [LawfulHashable α] (h : m.1.WF) :
+    m.1.keys.Pairwise (fun a b => (a == b) = false) := by
+  simp_to_model using (Raw.WF.out h).distinct.distinct
+
 end Raw₀
 
 end Std.DHashMap.Internal

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

@@ -59,6 +59,77 @@ theorem isEmpty_eq_isEmpty [BEq α] [Hashable α] {m : Raw α β} (h : Raw.WFImp
   rw [Raw.isEmpty, Bool.eq_iff_iff, List.isEmpty_iff_length_eq_zero, size_eq_length h,
     Nat.beq_eq_true_eq]
 
+theorem AssocList.foldlM_apply {l : AssocList α β} {acc : List γ} (f : (a : α) → β a → γ) :
+    l.foldlM (m := Id) (fun acc k v => f k v :: acc) acc =
+      (l.toList.map (fun p => f p.1 p.2)).reverse ++ acc := by
+  induction l generalizing acc <;> simp_all [AssocList.foldlM]
+
+theorem AssocList.foldlM_cons {l : AssocList α β} {acc : List ((a : α) × β a)} :
+    l.foldlM (m := Id) (fun acc k v => ⟨k, v⟩ :: acc) acc = l.toList.reverse ++ acc := by
+  simp [AssocList.foldlM_apply]
+
+theorem AssocList.foldlM_cons_key {l : AssocList α β} {acc : List α} :
+    l.foldlM (m := Id) (fun acc k _ => k :: acc) acc = (List.keys l.toList).reverse ++ acc := by
+  rw [AssocList.foldlM_apply, keys_eq_map]
+
+-- TODO (Markus): clean up
+theorem toList_perm_toListModel {m : Raw α β} : Perm m.toList (toListModel m.buckets) := by
+  rw [Raw.toList, toListModel, List.flatMap_eq_foldl, Raw.fold, Raw.foldM,
+    Array.foldlM_eq_foldlM_toList, ← List.foldl_eq_foldlM, Id.run]
+  simp only [AssocList.foldlM_cons]
+  suffices ∀ (l : List (AssocList α β)) (l₁ l₂), Perm l₁ l₂ →
+      Perm (l.foldl (fun acc m => m.toList.reverse ++ acc) l₁)
+        (l.foldl (fun acc m => acc ++ m.toList) l₂) by
+    simpa using this m.buckets.toList [] [] (Perm.refl _)
+  intros l l₁ l₂ h
+  induction l generalizing l₁ l₂
+  · simpa
+  · next l t ih =>
+    simp only [foldl_cons]
+    apply ih
+    refine (reverse_perm _).append_right _ |>.trans perm_append_comm |>.trans ?_
+    exact h.append_right l.toList
+
+-- TODO (Markus): clean up
+theorem keys_perm_keys_toListModel {m : Raw α β} :
+    Perm m.keys (List.keys (toListModel m.buckets)) := by
+  rw [Raw.keys, List.keys_eq_map]
+  refine Perm.trans ?_ (toList_perm_toListModel.map _)
+  simp only [Raw.toList, Raw.fold, Raw.foldM, Array.foldlM_eq_foldlM_toList, ← List.foldl_eq_foldlM,
+    Id.run]
+  generalize m.buckets.toList = l
+  suffices ∀ l', foldl
+      (fun acc l => AssocList.foldlM (m := Id) (fun acc k x => k :: acc) acc l) (l'.map (·.1)) l =
+      map Sigma.fst (foldl
+        (fun acc l => AssocList.foldlM (m := Id) (fun acc k v => ⟨k, v⟩ :: acc) acc l) l' l) from
+    this [] ▸ Perm.refl _
+  intro l'
+  simp [AssocList.foldlM_cons, AssocList.foldlM_cons_key]
+  induction l generalizing l' with
+  | nil => simp
+  | cons h t ih => simp [List.keys_eq_map, ← ih]
+
+theorem length_keys_eq_length_keys {m : Raw α β} :
+    m.keys.length = (List.keys (toListModel m.buckets)).length :=
+  keys_perm_keys_toListModel.length_eq
+
+theorem isEmpty_keys_eq_isEmpty_keys {m : Raw α β} :
+    m.keys.isEmpty = (List.keys (toListModel m.buckets)).isEmpty :=
+  keys_perm_keys_toListModel.isEmpty_eq
+
+theorem contains_keys_eq_contains_keys [BEq α] {m : Raw α β} {k : α} :
+    m.keys.contains k = (List.keys (toListModel m.buckets)).contains k :=
+  keys_perm_keys_toListModel.contains_eq
+
+theorem mem_keys_iff_contains_keys [BEq α] [LawfulBEq α] {m : Raw α β} {k : α} :
+    k ∈ m.keys ↔ (List.keys (toListModel m.buckets)).contains k := by
+  rw [← List.contains_iff_mem, contains_keys_eq_contains_keys]
+
+theorem pairwise_keys_iff_pairwise_keys [BEq α] [PartialEquivBEq α] {m : Raw α β} :
+    m.keys.Pairwise (fun a b => (a == b) = false) ↔
+      (List.keys (toListModel m.buckets)).Pairwise (fun a b => (a == b) = false) :=
+  keys_perm_keys_toListModel.pairwise_iff BEq.symm_false
+
 end Raw
 
 namespace Raw₀

src/Std/Data/DHashMap/Lemmas.lean

@@ -943,4 +943,29 @@ theorem getThenInsertIfNew?_snd {k : α} {v : β} :
 
 end Const
 
+@[simp]
+theorem length_keys_eq_size [EquivBEq α] [LawfulHashable α] :
+    m.keys.length = m.size :=
+  Raw₀.length_keys_eq_size ⟨m.1, m.2.size_buckets_pos⟩ m.2
+
+@[simp]
+theorem contains_keys [EquivBEq α] [LawfulHashable α] {k : α} :
+    m.keys.contains k = m.contains k :=
+  Raw₀.contains_keys ⟨m.1, _⟩ m.2
+
+@[simp]
+theorem isEmpty_keys_eq_isEmpty [EquivBEq α] [LawfulHashable α]:
+    m.keys.isEmpty = m.isEmpty  :=
+  Raw₀.isEmpty_keys_eq_isEmpty ⟨m.1, m.2.size_buckets_pos⟩ m.2
+
+@[simp]
+theorem mem_keys_iff_mem [LawfulBEq α] [LawfulHashable α] {k : α} :
+    k ∈ m.keys ↔ k ∈ m := by 
+  rw [mem_iff_contains]
+  exact Raw₀.mem_keys_iff_contains ⟨m.1, _⟩ m.2
+
+theorem distinct_keys [EquivBEq α] [LawfulHashable α] :
+    m.keys.Pairwise (fun a b => (a == b) = false) := 
+  Raw₀.distinct_keys ⟨m.1, m.2.size_buckets_pos⟩ m.2
+
 end Std.DHashMap

src/Std/Data/DHashMap/RawLemmas.lean

@@ -1022,6 +1022,31 @@ theorem getThenInsertIfNew?_snd (h : m.WF) {k : α} {v : β} :
 
 end Const
 
+@[simp]
+theorem length_keys_eq_size [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keys.length = m.size := by
+  simp_to_raw using Raw₀.length_keys_eq_size ⟨m, h.size_buckets_pos⟩ h
+
+@[simp]
+theorem contains_keys [EquivBEq α] [LawfulHashable α] (h : m.WF) {k : α} :
+    m.keys.contains k = m.contains k := by
+  simp_to_raw using Raw₀.contains_keys ⟨m, _⟩ h
+
+@[simp]
+theorem isEmpty_keys_eq_isEmpty [EquivBEq α] [LawfulHashable α] (h : m.WF):
+    m.keys.isEmpty = m.isEmpty := by
+  simp_to_raw using Raw₀.isEmpty_keys_eq_isEmpty ⟨m, h.size_buckets_pos⟩
+
+@[simp]
+theorem mem_keys_iff_mem [LawfulBEq α] [LawfulHashable α] (h : m.WF) {k : α} :
+    k ∈ m.keys ↔ k ∈ m := by 
+  rw [mem_iff_contains]
+  simp_to_raw using Raw₀.mem_keys_iff_contains ⟨m, _⟩ h
+
+theorem distinct_keys [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keys.Pairwise (fun a b => (a == b) = false) := by
+  simp_to_raw using Raw₀.distinct_keys ⟨m, h.size_buckets_pos⟩ h
+
 end Raw
 
 end Std.DHashMap

src/Std/Data/HashMap/Lemmas.lean

@@ -677,6 +677,30 @@ instance [EquivBEq α] [LawfulHashable α] : LawfulGetElem (HashMap α β) α β
     rw [getElem!_eq_get!_getElem?]
     split <;> simp_all
 
+@[simp]
+theorem length_keys_eq_size [EquivBEq α] [LawfulHashable α] :
+    m.keys.length = m.size :=
+  DHashMap.length_keys_eq_size
+
+@[simp]
+theorem contains_keys [EquivBEq α] [LawfulHashable α] {k : α} :
+    m.keys.contains k = m.contains k :=
+  DHashMap.contains_keys
+
+@[simp]
+theorem isEmpty_keys_eq_isEmpty [EquivBEq α] [LawfulHashable α]:
+    m.keys.isEmpty = m.isEmpty :=
+  DHashMap.isEmpty_keys_eq_isEmpty
+
+@[simp]
+theorem mem_keys_iff_mem [LawfulBEq α] [LawfulHashable α] {k : α} :
+    k ∈ m.keys ↔ k ∈ m := 
+  DHashMap.mem_keys_iff_mem
+
+theorem distinct_keys [EquivBEq α] [LawfulHashable α] :
+    m.keys.Pairwise (fun a b => (a == b) = false) := 
+  DHashMap.distinct_keys
+
 end
 
 end Std.HashMap

src/Std/Data/HashMap/RawLemmas.lean

@@ -682,6 +682,30 @@ theorem getThenInsertIfNew?_snd (h : m.WF) {k : α} {v : β} :
     (getThenInsertIfNew? m k v).2 = m.insertIfNew k v :=
   ext (DHashMap.Raw.Const.getThenInsertIfNew?_snd h.out)
 
+@[simp]
+theorem length_keys_eq_size [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keys.length = m.size :=
+  DHashMap.Raw.length_keys_eq_size h.out
+
+@[simp]
+theorem contains_keys [EquivBEq α] [LawfulHashable α] (h : m.WF) {k : α} :
+    m.keys.contains k = m.contains k :=
+  DHashMap.Raw.contains_keys h.out
+
+@[simp]
+theorem isEmpty_keys_eq_isEmpty [EquivBEq α] [LawfulHashable α] (h : m.WF):
+    m.keys.isEmpty = m.isEmpty :=
+  DHashMap.Raw.isEmpty_keys_eq_isEmpty h.out
+
+@[simp]
+theorem mem_keys_iff_mem [LawfulBEq α] [LawfulHashable α] (h : m.WF) {k : α} :
+    k ∈ m.keys ↔ k ∈ m := 
+  DHashMap.Raw.mem_keys_iff_mem h.out
+
+theorem distinct_keys [EquivBEq α] [LawfulHashable α] (h : m.WF) :
+    m.keys.Pairwise (fun a b => (a == b) = false) := 
+  DHashMap.Raw.distinct_keys h.out
+
 end Raw
 
 end Std.HashMap