feat: faster, linear HashMap.alter and modify

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

@@ -200,26 +200,29 @@ theorem toList_filter {f : (a : α) → β a → Bool} {l : AssocList α β} :
     · exact ih _
 
 theorem toList_alter [BEq α] [LawfulBEq α] {a : α} {f : Option (β a) → Option (β a)} {l : AssocList α β } :
-    (l.alter a f).toList = alterKey a f l.toList := by
+    Perm (l.alter a f).toList (alterKey a f l.toList) := by
   rw [alterKey]
   split
   · next heq =>
     induction l
     · simp only [toList_nil, getValueCast?_nil] at heq
       simp only [alter, heq, toList_nil, eraseKey_nil]
+      exact .nil
     · next ih =>
       simp only [toList_cons, getValueCast?_cons, beq_iff_eq] at heq
       split at heq
       · next heq₂ =>
         simp only [alter, dif_pos heq₂, heq, toList_cons, eraseKey_cons_of_beq heq₂]
+        exact Perm.rfl
       · next heq₂ =>
         simp only [alter, beq_iff_eq, dif_neg heq₂, toList_cons]
-        simp only [eraseKey_cons_of_false (Bool.not_eq_true _ ▸ heq₂), List.cons.injEq, true_and]
-        exact ih heq
+        simp only [eraseKey_cons_of_false (Bool.not_eq_true _ ▸ heq₂), true_and]
+        exact Perm.cons _ <| ih heq
   · next heq =>
     induction l
     · rw [toList_nil, getValueCast?_nil] at heq
       simp only [alter, heq, toList_cons, toList_nil, insertEntry_nil]
+      exact Perm.rfl
     · next ih =>
       simp only [toList_cons, getValueCast?_cons, beq_iff_eq] at heq
       split at heq
@@ -229,18 +232,14 @@ theorem toList_alter [BEq α] [LawfulBEq α] {a : α} {f : Option (β a) → Opt
         cases eq_of_beq heq₂; rfl
       · next heq₂ =>
         simp only [alter, beq_iff_eq, dif_neg heq₂, toList_cons]
-        -- oh no, this does not work: alter inserts at the end, while alterKey inserts at the front
-        done
-
-
--- theorem containsKey_alter [BEq α] [LawfulBEq α] {a : α} {f : Option (β a) → Option (β a)}
---     (l : AssocList α β ) : containsKey a (alter a f l)
-
-theorem toList_modify [BEq α] [LawfulBEq α] {a : α} {f : β a → β a} {l : AssocList α β } :
-    (l.modify a f).toList = modifyKey a f l.toList := sorry
+        refine insertEntry_cons_of_false (Bool.not_eq_true _ ▸ heq₂) |>.symm |> Perm.trans ?_
+        exact Perm.cons _ <| ih heq
 
 theorem modify_eq_alter [BEq α] [LawfulBEq α] {a : α} {f : β a → β a} {l : AssocList α β } :
-    modify a f l = alter a (·.map f) l := sorry:
+    modify a f l = alter a (·.map f) l := by
+  induction l
+  · rfl
+  · next ih => simp only [modify, beq_iff_eq, alter, Option.map_some', ih]
 
 theorem foldl_apply {l : AssocList α β} {acc : List δ} (f : (a : α) → β a → δ) :
     l.foldl (fun acc k v => f k v :: acc) acc =

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

@@ -972,6 +972,14 @@ theorem insertEntry_nil [BEq α] {k : α} {v : β k} :
     insertEntry k v ([] : List ((a : α) × β a)) = [⟨k, v⟩] := by
   simp [insertEntry]
 
+theorem insertEntry_cons_of_false [BEq α] {l : List ((a : α) × β a)} {k k' : α} {v : β k} {v' : β k'}
+    (h : (k' == k) = false) :
+    Perm (insertEntry k v (⟨k', v'⟩ :: l)) (⟨k', v'⟩ :: insertEntry k v l) := by
+  simp only [insertEntry, containsKey_cons, h, Bool.false_or, cond_eq_if]
+  split
+  · rw [replaceEntry_cons_of_false h]
+  · apply Perm.swap
+
 theorem insertEntry_cons_of_beq [BEq α] {l : List ((a : α) × β a)} {k k' : α} {v : β k} {v' : β k'}
     (h : k' == k) : insertEntry k v (⟨k', v'⟩ :: l) = ⟨k, v⟩ :: l := by
   simp_all only [insertEntry, containsKey_cons, Bool.true_or, cond_true, replaceEntry_cons_of_true]
@@ -1838,11 +1846,23 @@ theorem eraseKey_append_of_containsKey_right_eq_false [BEq α] {l l' : List ((a
     · rw [cond_true, cond_true]
 
 /-- Internal implementation detail of the hash map -/
-def alterKey [BEq α] [LawfulBEq α] (k : α) (f : Option (β k) → Option (β k))
-    (l : List ((a : α) × β a)) : List ((a : α) × β a) :=
-  match f (getValueCast? k l) with
+def alterKey [BEq α] [LawfulBEq α] (k : α) (f : Option (β k) → Option (β k)) :
+    List ((a : α) × β a) → List ((a : α) × β a) :=
+-- | [] => match f none with
+--   | none => []
+--   | some v => ⟨k, v⟩ :: []
+-- | ⟨k', v'⟩ :: tail =>
+--   if h : k' == k then
+--     match f (some (cast (congrArg β <| eq_of_beq h) v')) with
+--     | none => tail
+--     | some v => ⟨k, v⟩ :: tail
+--   else
+--     ⟨k', v'⟩ :: alterKey k f tail
+
+  fun l => match f (getValueCast? k l) with
   | none => eraseKey k l
   | some v => insertEntry k v l
+
   -- if h : containsKey k l then
   --   match f (some (getValueCast k l h)) with
   --   | none => eraseKey k l

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

@@ -165,13 +165,14 @@ theorem apply_bucket_with_proof {γ : α → Type w} [BEq α] [Hashable α] [Par
 /-- This is the general theorem to show that modification operations are correct. -/
 theorem toListModel_updateBucket [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α]
     {m : Raw₀ α β} (hm : Raw.WFImp m.1) {a : α} {f : AssocList α β → AssocList α β}
-    {g : List ((a : α) × β a) → List ((a : α) × β a)} (hfg : ∀ {l}, (f l).toList = g l.toList)
+    {g : List ((a : α) × β a) → List ((a : α) × β a)} (hfg : ∀ {l}, Perm (f l).toList (g l.toList))
     (hg₁ : ∀ {l l'}, DistinctKeys l → Perm l l' → Perm (g l) (g l'))
     (hg₂ : ∀ {l l'}, containsKey a l' = false → g (l ++ l') = g l ++ l') :
     Perm (toListModel (updateBucket m.1.buckets m.2 a f)) (g (toListModel m.1.2)) := by
   obtain ⟨l, h₁, h₂, h₃⟩ := exists_bucket_of_update m.1.buckets m.2 a f
   refine h₂.trans (Perm.trans ?_ (hg₁ hm.distinct h₁).symm)
-  rw [hfg, hg₂]
+  refine Perm.append_right l hfg |>.trans ?_
+  rw [hg₂]
   exact h₃ hm.buckets_hash_self _ rfl
 
 /-- This is the general theorem to show that mapping operations (like `map` and `filter`) are
@@ -445,7 +446,8 @@ theorem modify_eq_alter [BEq α] [Hashable α] [LawfulBEq α] (m : Raw₀ α β)
     · next h =>
       simp only [AssocList.contains_eq] at h
       simp only [AssocList.modify_eq_alter, Array.set_set, AssocList.contains_eq,
-        AssocList.toList_alter, ← modifyKey_eq_alterKey, containsKey_modifyKey_iff, h, ↓reduceIte]
+        containsKey_of_perm AssocList.toList_alter, ← modifyKey_eq_alterKey,
+        containsKey_modifyKey_iff, h, ↓reduceIte]
     · rfl
 
 theorem modify_eq_modifyₘ [BEq α] [Hashable α] [LawfulBEq α] (m : Raw₀ α β) (a : α)

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

@@ -407,7 +407,7 @@ end
 theorem toListModel_replaceₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (m : Raw₀ α β)
    (h : Raw.WFImp m.1) (a : α) (b : β a) :
   Perm (toListModel (m.replaceₘ a b).1.buckets) (replaceEntry a b (toListModel m.1.2)) :=
-  toListModel_updateBucket h AssocList.toList_replace List.replaceEntry_of_perm
+  toListModel_updateBucket h (.of_eq AssocList.toList_replace) List.replaceEntry_of_perm
     List.replaceEntry_append_of_containsKey_right_eq_false
 
 theorem isHashSelf_replaceₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (m : Raw₀ α β)
@@ -427,7 +427,7 @@ theorem wfImp_replaceₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α
 theorem toListModel_consₘ [BEq α] [Hashable α] [PartialEquivBEq α] [LawfulHashable α]
     (m : Raw₀ α β) (h : Raw.WFImp m.1) (a : α) (b : β a) :
     Perm (toListModel (m.consₘ a b).1.buckets) (⟨a, b⟩ :: (toListModel m.1.2)) :=
-  toListModel_updateBucket h rfl (fun _ => Perm.cons _) (fun _ => cons_append _ _ _)
+  toListModel_updateBucket h Perm.rfl (fun _ => Perm.cons _) (fun _ => cons_append _ _ _)
 
 theorem isHashSelf_consₘ [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (m : Raw₀ α β)
     (h : Raw.WFImp m.1) (a : α) (b : β a) : IsHashSelf (m.consₘ a b).1.buckets := by
@@ -486,12 +486,26 @@ theorem wfImp_insert [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] {m
 
 /-! # `alter` -/
 
+theorem mem_of_perm {l l' : List α} {x : α} (h : l ~ l') : x ∈ l ↔ x ∈ l' := by
+  induction h
+  case nil =>
+    simp only [not_mem_nil]
+  case cons y t' t a ih =>
+    simp only [mem_cons, ih]
+  case swap =>
+    simp only [mem_cons]
+    apply Iff.intro <;> intro x <;> cases x with
+    | inl y => simp only [y, true_or, or_true]
+    | inr y => cases y <;> simp only [*, true_or, or_true]
+  case trans ih₁ ih₂ =>
+    exact Iff.trans ih₁ ih₂
+
 theorem isHashSelf_alterₘ [BEq α] [Hashable α] [LawfulBEq α] (m : Raw₀ α β)
     (h : Raw.WFImp m.1) (a : α) (f : Option (β a) → Option (β a)) : IsHashSelf (m.alterₘ a f).1.buckets := by
   dsimp only [alterₘ, withComputedSize]
   split
   · apply h.buckets_hash_self.updateBucket (fun l p hp => ?_)
-    rw [AssocList.toList_alter] at hp
+    rw [mem_of_perm AssocList.toList_alter] at hp
     by_cases h : p.fst = a
     · exact .inr <| congrArg hash h
     · rw [mem_alterKey_of_key_ne _ h] at hp
@@ -528,7 +542,7 @@ theorem toListModel_alterₘ [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ 
       rw [insertEntry_of_containsKey_eq_false]
       simpa only [← containsₘ_eq_containsKey h] using hc
 
-theorem toListModel_updateBucket_alter [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α β} (h : Raw.WFImp m.1) {a : α}
+theorem toListModel_updateBucket'_alter [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α β} (h : Raw.WFImp m.1) {a : α}
     {f : Option (β a) → Option (β a)} :
     Perm (toListModel (updateBucket m.1.buckets m.2 a (AssocList.alter a f))) (alterKey a f (toListModel m.1.buckets)) := by
   apply toListModel_updateBucket h
@@ -542,7 +556,7 @@ theorem wfImp_updateBucket_alter [BEq α] [Hashable α] [LawfulBEq α] {m : Raw
   buckets_hash_self := by
     -- TODO: deduplicate with isHashSelf_alterₘ
     apply h.buckets_hash_self.updateBucket (fun l p hp => ?_)
-    rw [AssocList.toList_alter] at hp
+    rw [mem_of_perm AssocList.toList_alter] at hp
     by_cases h : p.fst = a
     · exact .inr <| congrArg hash h
     · rw [mem_alterKey_of_key_ne _ h] at hp
@@ -551,7 +565,7 @@ theorem wfImp_updateBucket_alter [BEq α] [Hashable α] [LawfulBEq α] {m : Raw
     rw [size_withComputedSize, computeSize_eq]
     rfl
   distinct := by
-    apply DistinctKeys.perm <| toListModel_updateBucket_alter h
+    apply DistinctKeys.perm <| toListModel_updateBucket'_alter h
     apply distinctKeys_alterKey
     exact h.distinct
 
@@ -571,7 +585,7 @@ theorem wfImp_alterₘ [BEq α] [Hashable α] [LawfulBEq α] {m : Raw₀ α β}
       simp only [length_alterKey, h.size_eq, dif_pos h₁]
       rw [containsₘ_eq_containsKey (by apply wfImp_updateBucket_alter h)]
       simp only [buckets_withComputedSize]
-      simp only [containsKey_of_perm <| toListModel_updateBucket_alter h]
+      simp only [containsKey_of_perm <| toListModel_updateBucket'_alter h]
       rw [← getValueCast?_eq_some_getValueCast h₁]
       conv => lhs; congr; rw [containsKey_alterKey_iff (a := a) (f := f) h.distinct];
     · next h₁ =>
@@ -699,7 +713,7 @@ theorem Const.wfImp_getThenInsertIfNew? {β : Type v} [BEq α] [Hashable α] [Eq
 theorem toListModel_eraseₘaux [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (m : Raw₀ α β)
     (a : α) (h : Raw.WFImp m.1) :
     Perm (toListModel (m.eraseₘaux a).1.buckets) (eraseKey a (toListModel m.1.buckets)) :=
-  toListModel_updateBucket h AssocList.toList_erase List.eraseKey_of_perm
+  toListModel_updateBucket h (Perm.of_eq AssocList.toList_erase) List.eraseKey_of_perm
     List.eraseKey_append_of_containsKey_right_eq_false
 
 theorem isHashSelf_eraseₘaux [BEq α] [Hashable α] [EquivBEq α] [LawfulHashable α] (m : Raw₀ α β)