Missing `HashMap.getThenInsert`

Stream: lean4

[Jovan Gerbscheid]

I want to get a value of a HashMap, and set it to a new value at the same time. There is no standard function for doing both at the same time (there is Std.HashMap.getThenInsertIfNew? which is not what I want), so instead I manually insert and get? separately. But the generated code does the insert before the get?, which means that the insert cannot be done in-place. What is the recommended solution? mwe:

import Std

set_option trace.Compiler.result true
def foo (ref : IO.Ref (Std.HashMap Nat Bool)) : IO (Option Bool) := do
  ref.modifyGet fun map => (map[0]?, map.insert 0 true)

Which compiles to

    def foo ref a.1 : EStateM.Result IO.Error lcRealWorld (Option Bool) :=
      let _x.2 := @ST.Prim.Ref.take ◾ ◾ ref a.1;
      cases _x.2 : EStateM.Result IO.Error lcRealWorld (Option Bool)
      | EStateM.Result.ok a.3 a.4 =>
        let _x.5 := 0;
        let _x.6 := true;
        let _x.7 := Std.DHashMap.Internal.Raw₀.insert._at_.foo.spec_0._redArg a.3 _x.5 _x.6;
        let _x.8 := @ST.Prim.Ref.set ◾ ◾ ref _x.7 a.4;
        cases _x.8 : EStateM.Result IO.Error lcRealWorld (Option Bool)
        | EStateM.Result.ok a.9 a.10 =>
          let _x.11 := Std.DHashMap.Internal.Raw₀.Const.get?._at_.foo.spec_6._redArg a.3 _x.5;
          let _x.12 := EStateM.Result.ok ◾ ◾ ◾ _x.11 a.10;
          return _x.12
        | EStateM.Result.error a.13 a.14 =>
          ⊥
      | EStateM.Result.error a.15 a.16 =>
        ⊥

[Markus Himmel]

Oh no, looks like lean4#6058 has returned. This is very unfortunate. I guess you can "launder" the map like so:

import Std

set_option trace.Compiler.result true
def foo (ref : IO.Ref (Std.HashMap Nat Bool)) : IO (Option Bool) := do
  ref.modifyGet fun map => let (v, m) := inner map; (v, m.insert 0 true)
where
  inner (m : Std.HashMap Nat Bool) : Option Bool × Std.HashMap Nat Bool :=
    (m[0]?, m)

This has a small overhead, but probably much less than a failure of linearity. Like in lean4#6058, I still think that this should be considered a bug in the compiler, so I think you can create a new lean4 issue referencing the old one.

Either way, adding an explicit getThenInsert is on the (long) containers TODO list. We might get to it this year, or we might not.

[Jovan Gerbscheid]

Thanks!

In my case the hashmap is actually very small, so it doesn't really matter. Do you know why there is the getThenInsertIfNew? function? It seems a bit arbitrary to have that one and not the other.

[Markus Himmel]

When we made the new Std hash map in early 2024 we made a long list of all of the functions that we want to have eventually and picked some subset of them to do immediately, leaving the others for the future (because there just wasn't enough time to do all of them). One of the criteria for immediate inclusion was that all functions that were present in the Lean and Batteries variants of the hash map needed to have an equivalent in the new hash map. The old Lean.HashMap had a function Lean.HashMap.insertIfNew which behaved like the new getThenInsertIfNew? , so that is why we added it immediately even though otherwise we probably would have postponed both getThenInsertIfNew? and getThenInsert?.

[Jovan Gerbscheid]

I suppose that the example that I posted can similarly apply to something like modifyGet fun arr => (arr.back!, arr.pop), so the issue is not fully resolved by having a getThenInsert?. Looking more closely at the compiled code (on nightly Lean, because we have zero cost BaseIO since very recently), I see the following:

inc x_3;
let x_7 : obj := Std.DHashMap.Internal.Raw₀.insert._at_.foo.spec_0._redArg x_3 x_4 x_6;
let x_8 : tagged := ST.Prim.Ref.set ◾ ◾ x_1 x_7 ◾;
let x_9 : tobj := Std.DHashMap.Internal.Raw₀.Const.get?._at_.foo.spec_6._redArg x_3 x_4;
dec x_3;

Would it be reasonable if the compiler could be able to figure out that it can save on inc and dec by swapping the order of operations? Because it could instead produce this:

let x_9 : tobj := Std.DHashMap.Internal.Raw₀.Const.get?._at_.foo.spec_6._redArg x_3 x_4;
let x_7 : obj := Std.DHashMap.Internal.Raw₀.insert._at_.foo.spec_0._redArg x_3 x_4 x_6;
let x_8 : tagged := ST.Prim.Ref.set ◾ ◾ x_1 x_7 ◾;

[Henrik Böving]

Reordering let operations at this level of the IR is generally unsound so no that is not reasonable.

[Jovan Gerbscheid]

Oh right, because then it could accidentally swap IO operations. Could this instead be done at an earlier level of the IR?