withImportModules and SimplePersistentEnvExtension

Stream: lean4

[Jon Eugster]

I can't get persistent env extensions working. The following procedure used to work in maybe Lean v4.7.0 or earlier, but I only get an empty output now.

In the first file, I define the env extension (not sure if .sync is correct):

File1.lean: define the env extension

import Lean
open Lean

initialize myExt : SimplePersistentEnvExtension String (Array String) ←
  registerSimplePersistentEnvExtension {
    name := `my_keys
    asyncMode := .sync
    addEntryFn := Array.push
    addImportedFn := Array.flatMap id }

Then in a second file I want to add something to the extension:

File2.lean: add something to the env extension

import MyModule.File1
open Lean

def test1  [Monad m] [MonadEnv m] : m Unit := do
  modifyEnv (myExt.addEntry · "test1")

#eval test1

Finally I call lake build MyModule.File2 to ensure the .oleans have been and try to use withImportModules in a third file to access the state at file 2:

File3.lean: use withImportModules to access the env at the version of File2

import MyModule.File1
open Lean

unsafe def test2 : IO Unit := do
  initSearchPath (← findSysroot)
  -- unsafe Lean.enableInitializersExecution        -- saw this somewhere?
  -- searchPathRef.set compile_time_search_path%    -- old code?
  withImportModules #[`MyModule.File2] {} (trustLevel := 1024) fun env => do
    let val := myExt.getState env
    IO.println val.toList

#eval test2 -- got [], expected ["test1"]

But I get [] instead of ["test1"]. What am I missing here?

(For context, the command lake exe i18n -t at lean-i18n is completely broken. It used to export all translation keys which are stored in this env extension, so file 2 is intended to be the main module, like Mathlib.lean)

[Yves Jäckle]

I had a similar problem, where using a version of withImportModules that calls importModules with loadExts set to true solved it. Didn't test this on your MWE though :sweat_smile:

[Robin Arnez]

From the doc-string of docs#Lean.withImportModules:

Creates environment object from imports and frees compacted regions after calling act. No live references to the environment object or imported objects may exist after act finishes. As this cannot be ruled out after loading environment extensions, docs#Lean.importModules​'s loadExts is always unset using this function.

and docs#Lean.importModules:

[...] If loadExts is true, we initialize the environment extensions using the imported data [...]

In other words, withImportModules doesn't load extension data because docs#Lean.Environment.freeRegions isn't safe when using extension data

[Robin Arnez]

You'll have to use docs#Lean.importModules directly

[Thomas Murrills]

(For what it’s worth, the PR which changed this behavior was lean4#6325; see also the notes on breaking changes in that PR’s diff.)

[Jon Eugster]

Ah I see, thanks for the pointers and the link to the PR!

[Jon Eugster]

So just to be explicit, I just change the mentioned flag

/-- same as `Lean.withImportModules` but with `(loadExts := true)`. -/
unsafe def MyModule.withImportModules {α : Type} (imports : Array Import) (opts : Options)
    (act : Environment → IO α) (trustLevel : UInt32 := 0) : IO α := do
  let env ← importModules (loadExts := true) imports opts trustLevel
  try act env finally env.freeRegions

add a random unsafe Lean.enableInitializersExecution before its usage and happily accept it's working, without thinking much about the details or consequences...? Seems fine by me :D

[Michael Rothgang]

CC @Anne Baanen as you encountered this question for the lint-style options

[Michael Rothgang]

TL;DR There may be a more robust way

[Anne Baanen]

The importModules (loadExts := true) approach seems to work fine but I did not trust that I understood the implications for garbage collection. So I ended up writing an elaborator that inserts the current state of the env extension as an expression.

[Anne Baanen]

But now I can't find that code actually, it has not been merged and it is not in one of my open PRs either. Where did I put that? :sweat_smile:

[Anne Baanen]

Ah, I totally missed that the last commit in #29513 got eaten. Fixed now! (Thanks to git reflog for finding it.)