feat: HashMap.modify

Author: digama0

Std/Data/Array/Lemmas.lean

@@ -107,6 +107,9 @@ theorem get_set (a : Array α) (i : Fin a.size) (j : Nat) (hj : j < a.size) (v :
     (a.set i v)[j]'(by simp [*]) = if i = j then v else a[j] := by
   if h : i.1 = j then subst j; simp [*] else simp [*]
 
+theorem set_set (a : Array α) (i : Fin a.size) (v v' : α) :
+    (a.set i v).set ⟨i, by simp [i.2]⟩ v' = a.set i v' := by simp [set, List.set_set]
+
 private theorem fin_cast_val (e : n = n') (i : Fin n) : e ▸ i = ⟨i.1, e ▸ i.2⟩ := by cases e; rfl
 
 theorem swap_def (a : Array α) (i j : Fin a.size) :

Std/Data/AssocList.lean

@@ -182,6 +182,23 @@ with key equal to `a` to have key `a` and value `b`.
 @[simp] theorem erase_toList [BEq α] (a : α) (l : AssocList α β) :
     (erase a l).toList = l.toList.eraseP (·.1 == a) := eraseP_toList ..
 
+/--
+`O(n)`. Replace the first entry `a', b` in the list
+with key equal to `a` to have key `a` and value `f a' b`.
+-/
+@[simp] def modify [BEq α] (a : α) (f : α → β → β) : AssocList α β → AssocList α β
+  | nil         => nil
+  | cons k v es => match k == a with
+    | true  => cons a (f k v) es
+    | false => cons k v (modify a f es)
+
+@[simp] theorem modify_toList [BEq α] (a : α) (l : AssocList α β) :
+    (modify a f l).toList =
+    l.toList.replaceF fun (k, v) => bif k == a then some (a, f k v) else none := by
+  simp [cond]
+  induction l with simp [List.replaceF]
+  | cons k v es ih => cases k == a <;> simp [ih]
+
 /-- The implementation of `ForIn`, which enables `for (k, v) in aList do ...` notation. -/
 @[specialize] protected def forIn [Monad m]
     (as : AssocList α β) (init : δ) (f : (α × β) → δ → m (ForInStep δ)) : m δ :=

Std/Data/HashMap/Basic.lean

@@ -39,6 +39,9 @@ Note: this is marked `noncomputable` because it is only intended for specificati
 -/
 noncomputable def size (data : Bucket α β) : Nat := .sum (data.1.data.map (·.toList.length))
 
+@[simp] theorem update_size (self : Bucket α β) (i d h) :
+    (self.update i d h).1.size = self.1.size := Array.size_uset ..
+
 /-- Map a function over the values in the map. -/
 @[specialize] def mapVal (f : α → β → γ) (self : Bucket α β) : Bucket α γ :=
   ⟨self.1.map (.mapVal f), by simp [self.2]⟩
@@ -181,6 +184,14 @@ def erase [BEq α] [Hashable α] (m : Imp α β) (a : α) : Imp α β :=
 @[inline] def mapVal (f : α → β → γ) (self : Imp α β) : Imp α γ :=
   { size := self.size, buckets := self.buckets.mapVal f }
 
+/-- Performs an in-place edit of the value, ensuring that the value is used linearly. -/
+def modify [BEq α] [Hashable α] (m : Imp α β) (a : α) (f : α → β → β) : Imp α β :=
+  let ⟨size, buckets⟩ := m
+  let ⟨i, h⟩ := mkIdx buckets.2 (hash a |>.toUSize)
+  let bkt := buckets.1[i]
+  let buckets := buckets.update i .nil h -- for linearity
+  ⟨size, buckets.update i (bkt.modify a f) ((Bucket.update_size ..).symm ▸ h)⟩
+
 /--
 Applies `f` to each key-value pair `a, b` in the map. If it returns `some c` then
 `a, c` is pushed into the new map; else the key is removed from the map.
@@ -223,6 +234,8 @@ inductive WF [BEq α] [Hashable α] : Imp α β → Prop where
   | insert : WF m → WF (insert m a b)
   /-- Removing an element from a well formed hash map yields a well formed hash map. -/
   | erase : WF m → WF (erase m a)
+  /-- Replacing an element in a well formed hash map yields a well formed hash map. -/
+  | modify : WF m → WF (modify m a f)
 
 theorem WF.empty [BEq α] [Hashable α] : WF (empty n : Imp α β) := by unfold empty; apply empty'
 
@@ -280,6 +293,13 @@ Removes key `a` from the map. If it does not exist in the map, the map is return
 -/
 @[inline] def erase (self : HashMap α β) (a : α) : HashMap α β := ⟨self.1.erase a, self.2.erase⟩
 
+/--
+Performs an in-place edit of the value, ensuring that the value is used linearly.
+The function `f` is passed the original key of the entry, along with the value in the map.
+-/
+def modify (self : HashMap α β) (a : α) (f : α → β → β) : HashMap α β :=
+  ⟨self.1.modify a f, self.2.modify⟩
+
 /-- Given a key `a`, returns a key-value pair in the map whose key compares equal to `a`. -/
 @[inline] def findEntry? (self : HashMap α β) (a : α) : Option (α × β) := self.1.findEntry? a
 

Std/Data/HashMap/WF.lean

@@ -21,14 +21,15 @@ namespace Bucket
 theorem update_data (self : Bucket α β) (i d h) :
     (self.update i d h).1.data = self.1.data.set i.toNat d := rfl
 
-@[simp] theorem update_size (self : Bucket α β) (i d h) :
-    (self.update i d h).1.size = self.1.size := Array.size_uset ..
-
 theorem exists_of_update (self : Bucket α β) (i d h) :
     ∃ l₁ l₂, self.1.data = l₁ ++ self.1[i] :: l₂ ∧ List.length l₁ = i.toNat ∧
       (self.update i d h).1.data = l₁ ++ d :: l₂ := by
   simp [Array.getElem_eq_data_get]; exact List.exists_of_set' h
 
+theorem update_update (self : Bucket α β) (i d d' h h') :
+    (self.update i d h).update i d' h' = self.update i d' h := by
+  simp [update]; congr 1; rw [Array.set_set]
+
 theorem size_eq (data : Bucket α β) :
   size data = .sum (data.1.data.map (·.toList.length)) := rfl
 
@@ -185,8 +186,8 @@ theorem insert_size [BEq α] [Hashable α] {m : Imp α β} {k v}
     refine have ⟨_, _, h₁, _, eq⟩ := Bucket.exists_of_update ..; eq ▸ ?_
     simp [h₁, Bucket.size_eq, Nat.succ_add]; rfl
 
-private theorem mem_replaceF {l : List (α × β)} {x : α × β} {p : α × β → Bool} :
-    x ∈ (l.replaceF fun a => bif p a then some (k, v) else none) → x.1 = k ∨ x ∈ l := by
+private theorem mem_replaceF {l : List (α × β)} {x : α × β} {p : α × β → Bool} {f : α × β → β} :
+    x ∈ (l.replaceF fun a => bif p a then some (k, f a) else none) → x.1 = k ∨ x ∈ l := by
   induction l with
   | nil => exact .inr
   | cons a l ih =>
@@ -200,16 +201,16 @@ private theorem mem_replaceF {l : List (α × β)} {x : α × β} {p : α × β
       | .inr h => exact (ih h).imp_right .inr
 
 private theorem pairwise_replaceF [BEq α] [PartialEquivBEq α]
-    {l : List (α × β)} {x : α × β} (hx₁ : x ∈ l) (hx₂ : x.fst == k)
+    {l : List (α × β)} {f : α × β → β}
     (H : l.Pairwise fun a b => ¬(a.fst == b.fst)) :
-    (l.replaceF fun a => bif a.fst == k then some (k, v) else none)
+    (l.replaceF fun a => bif a.fst == k then some (k, f a) else none)
       |>.Pairwise fun a b => ¬(a.fst == b.fst) := by
-  induction hx₁ with
-  | head => simp_all; exact (H.1 · · ∘ PartialEquivBEq.trans hx₂)
-  | tail _ _ ih =>
+  induction l with
+  | nil => simp [H]
+  | cons a l ih =>
     simp at H ⊢
-    generalize e : cond .. = z; revert e
-    unfold cond; split <;> (intro h; subst h; simp)
+    generalize e : cond .. = z; unfold cond at e; revert e
+    split <;> (intro h; subst h; simp)
     · next e => exact ⟨(H.1 · · ∘ PartialEquivBEq.trans e), H.2⟩
     · next e =>
       refine ⟨fun a h => ?_, ih H.2⟩
@@ -223,7 +224,7 @@ theorem insert_WF [BEq α] [Hashable α] {m : Imp α β} {k v}
   · next h₁ =>
     simp at h₁; have ⟨x, hx₁, hx₂⟩ := h₁
     refine h.update (fun H => ?_) (fun H a h => ?_)
-    · simp; exact pairwise_replaceF hx₁ hx₂ H
+    · simp; exact pairwise_replaceF H
     · simp [AssocList.All] at H h ⊢
       match mem_replaceF h with
       | .inl rfl => rfl
@@ -261,13 +262,32 @@ theorem erase_WF [BEq α] [Hashable α] {m : Imp α β} {k}
     · exact H _ (List.mem_of_mem_eraseP h)
   · exact h
 
+theorem modify_size [BEq α] [Hashable α] {m : Imp α β} {k}
+    (h : m.size = m.buckets.size) :
+    (modify m k f).size = (modify m k f).buckets.size := by
+  dsimp [modify, cond]; rw [Bucket.update_update]
+  simp [h, Bucket.size]
+  refine have ⟨_, _, h₁, _, eq⟩ := Bucket.exists_of_update ..; eq ▸ ?_
+  simp [h, h₁, Bucket.size_eq]
+
+theorem modify_WF [BEq α] [Hashable α] {m : Imp α β} {k}
+    (h : m.buckets.WF) : (modify m k f).buckets.WF := by
+  dsimp [modify, cond]; rw [Bucket.update_update]
+  refine h.update (fun H => ?_) (fun H a h => ?_) <;> simp at h ⊢
+  · exact pairwise_replaceF H
+  · simp [AssocList.All] at H h ⊢
+    match mem_replaceF h with
+    | .inl rfl => rfl
+    | .inr h => exact H _ h
+
 theorem WF.out [BEq α] [Hashable α] {m : Imp α β} (h : m.WF) :
     m.size = m.buckets.size ∧ m.buckets.WF := by
   induction h with
   | mk h₁ h₂ => exact ⟨h₁, h₂⟩
   | @empty' _ h => exact ⟨(Bucket.mk_size h).symm, .mk' h⟩
   | insert _ ih => exact ⟨insert_size ih.1, insert_WF ih.2⟩
   | erase _ ih => exact ⟨erase_size ih.1, erase_WF ih.2⟩
+  | modify _ ih => exact ⟨modify_size ih.1, modify_WF ih.2⟩
 
 theorem WF_iff [BEq α] [Hashable α] {m : Imp α β} :
     m.WF ↔ m.size = m.buckets.size ∧ m.buckets.WF :=

Std/Data/List/Basic.lean

@@ -390,26 +390,23 @@ def indexOf [BEq α] (a : α) : List α → Nat := findIdx (a == ·)
     | none => x :: replaceF f xs
     | some a => a :: xs
 
-/-- Tail recursive version of `replaceF`. -/
+/-- Tail-recursive version of `replaceF`. -/
 @[inline] def replaceFTR (f : α → Option α) (l : List α) : List α := go l #[] where
-  /-- Auxiliary for `replaceFTR`:
-  `replaceFTR.go f l xs acc = acc.toList ++ replaceF f xs` if `f` returns `some`, else `l`. -/
-  go : List α → Array α → List α
-  | [], _ => l
+  /-- Auxiliary for `replaceFTR`: `replaceFTR.go f xs acc = acc.toList ++ replaceF f xs`. -/
+  @[specialize] go : List α → Array α → List α
+  | [], acc => acc.toList
   | x :: xs, acc => match f x with
     | none => go xs (acc.push x)
-    | some a => acc.toListAppend (a :: xs)
+    | some a' => acc.toListAppend (a' :: xs)
 
 @[csimp] theorem replaceF_eq_replaceFTR : @replaceF = @replaceFTR := by
-  funext α f l; simp [replaceFTR]
-  suffices ∀ xs acc, l = acc.data ++ xs →
-      replaceFTR.go f l xs acc = acc.data ++ xs.replaceF f from
-    (this l #[] (by simp)).symm
-  intro xs; induction xs with intro acc
-  | nil => simp [replaceF, replaceFTR.go]
-  | cons x xs IH =>
-    simp [replaceF, replaceFTR.go]; split <;> simp [*]
-    · intro h; rw [IH]; simp; simp; exact h
+  funext α p l; simp [replaceFTR]
+  let rec go (acc) : ∀ xs, replaceFTR.go p xs acc = acc.data ++ xs.replaceF p
+  | [] => by simp [replaceFTR.go, replaceF]
+  | x::xs => by
+    simp [replaceFTR.go, replaceF]; cases p x <;> simp
+    · rw [go _ xs]; simp
+  exact (go #[] _).symm
 
 /-- Inserts an element into a list without duplication. -/
 @[inline] protected def insert [DecidableEq α] (a : α) (l : List α) : List α :=

Std/Data/List/Lemmas.lean

@@ -819,6 +819,11 @@ theorem set_comm (a b : α) : ∀ {n m : Nat} (l : List α), n ≠ m →
   | n+1, m+1, x :: t, h =>
     congrArg _ <| set_comm a b t fun h' => h <| Nat.succ_inj'.mpr h'
 
+theorem set_set (a b : α) : ∀ (l : List α) (n : Nat), (l.set n a).set n b = l.set n b
+  | [], _ => by simp
+  | _ :: _, 0 => by simp [set]
+  | _ :: _, _+1 => by simp [set, set_set]
+
 @[simp] theorem get_set_eq (l : List α) (i : Nat) (a : α) (h : i < (l.set i a).length) :
     (l.set i a).get ⟨i, h⟩ = a := by
   rw [← Option.some_inj, ← get?_eq_get, get?_set_eq, get?_eq_get] <;> simp_all
@@ -1280,8 +1285,51 @@ theorem Pairwise.imp {α R S} (H : ∀ {a b}, R a b → S a b) :
 
 /-! ### replaceF -/
 
+theorem replaceF_nil : [].replaceF p = [] := rfl
+
+theorem replaceF_cons (a : α) (l : List α) :
+    (a :: l).replaceF p = match p a with
+      | none => a :: replaceF p l
+      | some a' => a' :: l := rfl
+
+theorem replaceF_cons_of_some {l : List α} (p) (h : p a = some a') :
+    (a :: l).replaceF p = a' :: l := by
+  simp [replaceF_cons, h]
+
+theorem replaceF_cons_of_none {l : List α} (p) (h : p a = none) :
+    (a :: l).replaceF p = a :: l.replaceF p := by simp [replaceF_cons, h]
+
+theorem replaceF_of_forall_none {l : List α} (h : ∀ a, a ∈ l → p a = none) : l.replaceF p = l := by
+  induction l with
+  | nil => rfl
+  | cons _ _ ih => simp [h _ (.head ..), ih (forall_mem_cons.1 h).2]
+
+theorem exists_of_replaceF : ∀ {l : List α} {a a'} (al : a ∈ l) (pa : p a = some a'),
+    ∃ a a' l₁ l₂,
+      (∀ b ∈ l₁, p b = none) ∧ p a = some a' ∧ l = l₁ ++ a :: l₂ ∧ l.replaceF p = l₁ ++ a' :: l₂
+  | b :: l, a, a', al, pa =>
+    match pb : p b with
+    | some b' => ⟨b, b', [], l, forall_mem_nil _, pb, by simp [pb]⟩
+    | none =>
+      match al with
+      | .head .. => nomatch pb.symm.trans pa
+      | .tail _ al =>
+        let ⟨c, c', l₁, l₂, h₁, h₂, h₃, h₄⟩ := exists_of_replaceF al pa
+        ⟨c, c', b::l₁, l₂, (forall_mem_cons ..).2 ⟨pb, h₁⟩,
+          h₂, by rw [h₃, cons_append], by simp [pb, h₄]⟩
+
+theorem exists_or_eq_self_of_replaceF (p) (l : List α) :
+    l.replaceF p = l ∨ ∃ a a' l₁ l₂,
+      (∀ b ∈ l₁, p b = none) ∧ p a = some a' ∧ l = l₁ ++ a :: l₂ ∧ l.replaceF p = l₁ ++ a' :: l₂ :=
+  if h : ∃ a ∈ l, (p a).isSome then
+    let ⟨_, ha, pa⟩ := h
+    .inr (exists_of_replaceF ha (Option.get_mem pa))
+  else
+    .inl <| replaceF_of_forall_none fun a ha =>
+      Option.not_isSome_iff_eq_none.1 fun h' => h ⟨a, ha, h'⟩
+
 @[simp] theorem length_replaceF : length (replaceF f l) = length l := by
-  induction l <;> simp; split <;> simp [*]
+  induction l <;> simp [replaceF]; split <;> simp [*]
 
 /-! ### disjoint -/