chore: basic functionality tests for modify/alter

src/Std/Data/DHashMap/Basic.lean

@@ -247,6 +247,39 @@ instance [BEq α] [Hashable α] : ForIn m (DHashMap α β) ((a : α) × β a) wh
     (m : DHashMap α (fun _ => β)) : Array β :=
   m.1.valuesArray
 
+/--
+Modifies in place the value associated with a given key.
+
+This function ensures that the value is used linearly.
+It is currently implemented in terms of `get?`, `erase`, and `insert`,
+but will later become a primitive operation.
+(It is provided already to help avoid non-linear code.)
+-/
+@[inline] def modify [LawfulBEq α] (m : DHashMap α β) (a : α) (f : β a → β a) : DHashMap α β :=
+  match m.get? a with
+  | none => m
+  | some b => m.erase a |>.insert a (f b)
+
+/--
+Modifies in place the value associated with a given key,
+allowing creating new values and deleting values via an `Option` valued replacement function.
+
+This function ensures that the value is used linearly.
+It is currently implemented in terms of `get?`, `erase`, and `insert`,
+but will later become a primitive operation.
+(It is provided already to help avoid non-linear code.)
+-/
+@[inline] def alter [LawfulBEq α] (m : DHashMap α β) (a : α) (f : Option (β a) → Option (β a)) : DHashMap α β :=
+  match m.get? a with
+  | none =>
+    match f none with
+    | none => m
+    | some b => m.insert a b
+  | some b =>
+    match f (some b) with
+    | none => m.erase a
+    | some b => m.erase a |>.insert a b
+
 @[inline, inherit_doc Raw.insertMany] def insertMany {ρ : Type w}
     [ForIn Id ρ ((a : α) × β a)] (m : DHashMap α β) (l : ρ) : DHashMap α β :=
   ⟨(Raw₀.insertMany ⟨m.1, m.2.size_buckets_pos⟩ l).1,

src/Std/Data/HashMap/Basic.lean

@@ -245,6 +245,22 @@ instance [BEq α] [Hashable α] {m : Type w → Type w} : ForIn m (HashMap α β
     Array β :=
   m.inner.valuesArray
 
+@[inline, inherit_doc DHashMap.modify] def modify (m : HashMap α β) (a : α) (f : β → β) : HashMap α β :=
+  match m.get? a with
+  | none => m
+  | some b => m.erase a |>.insert a (f b)
+
+@[inline, inherit_doc DHashMap.alter] def alter (m : HashMap α β) (a : α) (f : Option β → Option β) : HashMap α β :=
+  match m.get? a with
+  | none =>
+    match f none with
+    | none => m
+    | some b => m.insert a b
+  | some b =>
+    match f (some b) with
+    | none => m.erase a
+    | some b => m.erase a |>.insert a b
+
 @[inline, inherit_doc DHashMap.Const.insertMany] def insertMany {ρ : Type w}
     [ForIn Id ρ (α × β)] (m : HashMap α β) (l : ρ) : HashMap α β :=
   ⟨DHashMap.Const.insertMany m.inner l⟩
@@ -290,7 +306,13 @@ def Array.groupByKey [BEq α] [Hashable α] (key : β → α) (xs : Array β)
     : Std.HashMap α (Array β) := Id.run do
   let mut groups := ∅
   for x in xs do
-    let group := groups.getD (key x) #[]
-    groups := groups.erase (key x) -- make `group` referentially unique
-    groups := groups.insert (key x) (group.push x)
+    groups := groups.alter (key x) (·.getD #[] |>.push x)
   return groups
+
+/--
+Groups all elements `x`, `y` in `xs` with `key x == key y` into the same list
+`(xs.groupByKey key).find! (key x)`. Groups preserve the relative order of elements in `xs`.
+-/
+def List.groupByKey [BEq α] [Hashable α] (key : β → α) (xs : List β) :
+    Std.HashMap α (List β) :=
+  xs.foldr (init := ∅) fun x acc => acc.alter (key x) (fun v => x :: v.getD [])

tests/lean/run/hashmap.lean

@@ -336,6 +336,34 @@ def addValueToState (_ : Nat) (v : Nat) : StateM Nat PUnit := do
 #guard_msgs in
 #eval m ∪ {(4, 5), (37, 37)}
 
+/-- info: Std.HashMap.ofList [(3, 6), (2, 4), (1, 3)] -/
+#guard_msgs in
+#eval m.modify 1 (fun v => v + 1)
+
+/-- info: Std.HashMap.ofList [(3, 6), (2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.modify 4 (fun v => v + 1)
+
+/-- info: Std.HashMap.ofList [(3, 6), (2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.alter 4 (fun v? => v?.map (· + 2))
+
+/-- info: Std.HashMap.ofList [(4, 7), (3, 6), (2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.alter 4 (fun _ => some 7)
+
+/-- info: Std.HashMap.ofList [(3, 6), (2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.alter 4 (fun _ => none)
+
+/-- info: Std.HashMap.ofList [(2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.alter 3 (fun _ => none)
+
+/-- info: Std.HashMap.ofList [(3, 37), (2, 4), (1, 2)] -/
+#guard_msgs in
+#eval m.alter 3 (fun _ => some 37)
+
 end HashMap
 
 namespace HashSet.Raw