Modify value of HashMap

Stream: lean4

[Tomas Skrivan]

How can I easily modify a value of a HashMap and ensure linearity of the value?

Is this doing what I want?

unsafe def modify {α β} [BEq α] [Hashable α] (m : HashMap α β) (key : α) (f : β → β) : HashMap α β :=
  if let .some val := m.find? key then
    let m := m.insert key (unsafeCast ()) -- ensures linearity?
    m.insert key (f val)
  else
    m

However this is effectively doing three look-ups in the hash table which is a waste. Is there a better way?

[Eric Wieser]

(docs4#Std.HashMap)

[Tomas Skrivan]

I don't think any of those functions do what I'm asking for. Maybe you can put it together with functions from HashMap.Impl but I do not really understand the inner workings of the implementation.

[Mario Carneiro]

It's not currently supported. Here's a quick implementation, the real thing requires doing some proofs:

import Std.Data.HashMap

/--
`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 Std.AssocList.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)

namespace Std.HashMap

/-- Performs an in-place edit of the value, ensuring that the value is used linearly. -/
def Imp.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)⟩

variable [BEq α] [Hashable α]

/-- Performs an in-place edit of the value, ensuring that the value is used linearly. -/
unsafe def modifyImpl (self : HashMap α β) (a : α) (f : α → β → β) : HashMap α β :=
  ⟨self.1.modify a f, lcProof⟩

/-- Performs an in-place edit of the value, ensuring that the value is used linearly. -/
@[implemented_by modifyImpl]
opaque modify (self : HashMap α β) (a : α) (f : α → β → β) : HashMap α β

end Std.HashMap

[Eric Wieser]

Would α → Option β → Option β be a more sensible signature? That way you can handle removing or inserting during the update too (without computing the hash twice)

[Mario Carneiro]

that's alter, it is a lot more complicated to prove

[Tomas Skrivan]

Nice, thanks! I'm happy to use an unproven version now. Having alter would be also very useful.

[Mario Carneiro]

done in https://github.com/leanprover/std4/commit/ed1f8fe