HashMap extension

Stream: lean4

[Patrick Massot]

Playing with environment extensions, I realized I kept using extensions whose states were lists of pairs and searching in the list elements with a given first component. So I feel like I want extensions where the state is a hash map. But my understanding of environment extensions is very shallow. The following code seems to be doing what I want, but I'm not confident at all that I'm not doing something very stupid.

import Lean
import Std.Data.HashMap

open Lean Std

def HashMapExtension (α β : Type) [BEq α] [Hashable α] := SimplePersistentEnvExtension (α × β) (HashMap α β)

instance (α β : Type) [BEq α] [Hashable α] : Inhabited (HashMapExtension α β) :=
  inferInstanceAs (Inhabited (SimplePersistentEnvExtension (α × β) (HashMap α β)))

def mkHashMapExtension (name : Name) (α β : Type) [BEq α] [Hashable α]  : IO (HashMapExtension α β) :=
  registerSimplePersistentEnvExtension {
    name          := name,
    addImportedFn := fun as => HashMap.ofList (as.concatMap id).toList,
    addEntryFn    := fun s n => s.insert n.1 n.2 ,
    toArrayFn     := fun es => es.toArray
  }

-- Now I'll write some tests

initialize testExt : HashMapExtension Nat String ← mkHashMapExtension `my_test Nat String

elab "#add_test" name:num content:str : command =>
  modifyEnv (testExt.addEntry · (name.getNat, content.getString))

elab "#list_tests" : command => do
  let s := testExt.getState (← getEnv)
  dbg_trace s.toList

elab "#get_test" n:num : command => do
  let s := testExt.getState (← getEnv)
  dbg_trace s.find? n.getNat

[Patrick Massot]

@Sebastian Ullrich while you're here, do you have any comment about this thread?

[Sebastian Ullrich]

Looks good I think. You can use mkStateFromImportedEntries to simplify and optimize your addImportedFn.

[Patrick Massot]

I was especially worried about that addImportedFn indeed. Now I need to decipher what mkStateFromImportedEntries does...

[Patrick Massot]

Do you mean

import Lean
import Std.Data.HashMap

open Lean Std

def HashMapExtension (α β : Type) [BEq α] [Hashable α] := SimplePersistentEnvExtension (α × β) (HashMap α β)

instance (α β : Type) [BEq α] [Hashable α] : Inhabited (HashMapExtension α β) :=
  inferInstanceAs (Inhabited (SimplePersistentEnvExtension (α × β) (HashMap α β)))

def HashMapExtension.addEntryFn {α β : Type} [BEq α] [Hashable α] (s : HashMap α β) (n : α × β) : HashMap α β :=
s.insert n.1 n.2

def mkHashMapExtension (name : Name) (α β : Type) [BEq α] [Hashable α]  : IO (HashMapExtension α β) :=
  registerSimplePersistentEnvExtension {
    name          := name,
    addImportedFn := mkStateFromImportedEntries HashMapExtension.addEntryFn {},
    addEntryFn    := HashMapExtension.addEntryFn,
    toArrayFn     := fun es => es.toArray
  }

[Sebastian Ullrich]

Yes, though I would probably have inlined the new function as a lambda :)

[Patrick Massot]

Thanks!

[Patrick Massot]

I'm returning to this old thread. Now that I try to actually use this HashMap extension idea I see that it doesn't work. I have a file HashMapExt.lean with

import Lean

open Lean Std

def HashMapExtension (α β : Type) [BEq α] [Hashable α] := SimplePersistentEnvExtension (α × β) (HashMap α β)

instance (α β : Type) [BEq α] [Hashable α] : Inhabited (HashMapExtension α β) :=
  inferInstanceAs (Inhabited (SimplePersistentEnvExtension (α × β) (HashMap α β)))

def mkHashMapExtension (name : Name) (α β : Type) [BEq α] [Hashable α]  : IO (HashMapExtension α β) :=
  registerSimplePersistentEnvExtension {
    name          := name,
    addImportedFn := mkStateFromImportedEntries (λ s n => s.insert n.1 n.2) {},
    addEntryFn    := (λ s n => s.insert n.1 n.2),
    toArrayFn     := fun es => es.toArray
  }

namespace HashMapExtension

variable {α β : Type} [BEq α] [Hashable α] {m: Type → Type} [Monad m] [MonadEnv m]

def find? (ext : HashMapExtension α β) (a : α) : m $ Option β := do
  return (ext.getState (← getEnv)).find? a

def insert (ext : HashMapExtension α β) (a : α) (b : β) : m Unit :=
  modifyEnv (ext.modifyState · (·.insert a b))

def update (ext : HashMapExtension α β) (a : α) (b : β) : m Unit :=
  modifyEnv (ext.modifyState · (λ hm => (hm.erase a).insert a b))

end HashMapExtension

and then a second file test0.lean with

import HashMapExt

initialize testExt : HashMapExtension Nat String ← mkHashMapExtension `test Nat String

elab "Test" n:num s:str : command => do
  testExt.insert n.getNat s.getString

open Lean

elab "PrintTests" : command => do
  logInfo $ repr $ (testExt.getState (← getEnv)).toList

and test1.lean

import test0

Test 3 "Foo"

PrintTests

so far everything seems to work. But when I create test2.lean with

import test1

Test 4 "Bar"

PrintTests

I see only the entry created in test2, not two entries as I expected.

[Patrick Massot]

Sebastian, do you see what I'm doing wrong?

[Sebastian Ullrich]

Not directly, I'd have to take a closer look

[Patrick Massot]

Should I create an issue then?

[Patrick Massot]

For people who want to investigate this, issue.zip contains a folder issue with a Lean project showing my problem in the Issue lib.

[Patrick Massot]

(deleted)

[Gabriel Ebner]

Your issue is with the HashMapExtension.insert and update functions.

[Gabriel Ebner]

The SimplePersistentEnvExtension.modifyState should never be used because the state is not persisted.

[Gabriel Ebner]

(It's obvious that it can't be persisted: what is being written to the olean file is the entries, while modifyState applies some opaque function to the state. It doesn't know which entries to store to recreate the state modification.)

[Gabriel Ebner]

This should fix your issues:

def insert (ext : HashMapExtension α β) (a : α) (b : β) : m Unit :=
  modifyEnv (ext.addEntry · (a, b))

def update (ext : HashMapExtension α β) (a : α) (b : β) : m Unit :=
  modifyEnv (ext.addEntry · (a, b))

[Patrick Massot]

Indeed it works much better!

[Patrick Massot]

Thanks a lot Gabriel!

[Patrick Massot]

SimplePersistentEnvExtension.modifyState is indeed rather misleading. Should I PR a docstring with a warning or do you think it should be completely removed?

[Gabriel Ebner]

Yes. At least a docstring would be good.

[Patrick Massot]

It has nothing to do with "Simple". Already PersistentExtension.modifyState is implemented as

def modifyState {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : Environment) (f : σ → σ) : Environment :=
  ext.toEnvExtension.modifyState env fun ps => { ps with state := f (ps.state) }

where ext.toEnvExtension very much sounds like to first step is to forget about persistance...

[Gabriel Ebner]

Yeah, that function should get a docstring as well.

[Gabriel Ebner]

I believe the only use of these dangerous functions is to implement ScopedEnvExtension .

[Patrick Massot]

Actually the same projection occurs in addEntry so it doesn't explain why your version works

[Gabriel Ebner]

Maybe it's best to just remove both sets of them, and then have ScopedEnvExtension break the abstraction barrier explicitly.

[Gabriel Ebner]

Patrick Massot said:

Actually the same projection occurs in addEntry so it doesn't explain why your version works

If you look at the addEntry function, it calls the addEntryFn field. And the addEntryFn is set by registerSimplePersistentEnvExtension to a function that updates both the state as well as the entries.

[Patrick Massot]

I opened https://github.com/leanprover/lean4/pull/1649