refactor: environment extension state splitting (#8653)

Replaces the previous `export/saveEntriesFn` split with a stricly more
general function such that `exportEntriesFn` could be deprecated at a
later point. Also gives the new function access to the `Environment`
while we're at it. Also gives `getModuleEntries` access to more olean
levels in preparation for `meta import`.

Author: Kha

src/Lean/DeclarationRange.lean

@@ -15,7 +15,10 @@ namespace Lean
 
 builtin_initialize builtinDeclRanges : IO.Ref (NameMap DeclarationRanges) ← IO.mkRef {}
 builtin_initialize declRangeExt : MapDeclarationExtension DeclarationRanges ←
-  mkMapDeclarationExtension (exportEntriesFn := fun _ => #[])
+  mkMapDeclarationExtension (exportEntriesFn := fun _ s level =>
+    if level < .server then
+      #[]
+    else s.toArray)
 
 def addBuiltinDeclarationRanges (declName : Name) (declRanges : DeclarationRanges) : IO Unit :=
   builtinDeclRanges.modify (·.insert declName declRanges)
@@ -24,7 +27,7 @@ def addDeclarationRanges [Monad m] [MonadEnv m] (declName : Name) (declRanges :
   modifyEnv fun env => declRangeExt.insert env declName declRanges
 
 def findDeclarationRangesCore? [Monad m] [MonadEnv m] (declName : Name) : m (Option DeclarationRanges) :=
-  return declRangeExt.find? (includeServer := true) (← getEnv) declName
+  return declRangeExt.find? (level := .server) (← getEnv) declName
 
 def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT BaseIO m] (declName : Name) : m (Option DeclarationRanges) := do
   let env ← getEnv

src/Lean/DocString/Extension.lean

@@ -58,8 +58,9 @@ structure ModuleDoc where
 private builtin_initialize moduleDocExt : SimplePersistentEnvExtension ModuleDoc (PersistentArray ModuleDoc) ← registerSimplePersistentEnvExtension {
   addImportedFn := fun _ => {}
   addEntryFn    := fun s e => s.push e
-  toArrayFn     := fun es => es.toArray
-  exported      := false
+  exportEntriesFnEx? := some fun _ _ es => fun
+    | .exported => #[]
+    | _         => es.toArray
 }
 
 def addMainModuleDoc (env : Environment) (doc : ModuleDoc) : Environment :=
@@ -70,7 +71,7 @@ def getMainModuleDoc (env : Environment) : PersistentArray ModuleDoc :=
 
 def getModuleDoc? (env : Environment) (moduleName : Name) : Option (Array ModuleDoc) :=
   env.getModuleIdx? moduleName |>.map fun modIdx =>
-    moduleDocExt.getModuleEntries (includeServer := true) env modIdx
+    moduleDocExt.getModuleEntries (level := .server) env modIdx
 
 def getDocStringText [Monad m] [MonadError m] (stx : TSyntax `Lean.Parser.Command.docComment) : m String :=
   match stx.raw[1] with

src/Lean/Elab/RecommendedSpelling.lean

@@ -28,7 +28,7 @@ open Lean.Parser.Command
 def allRecommendedSpellings : MetaM (Array RecommendedSpelling) := do
   let all := recommendedSpellingExt.toEnvExtension.getState (← getEnv)
       |>.importedEntries
-      |>.push (recommendedSpellingExt.exportEntriesFn (recommendedSpellingExt.getState (← getEnv)))
+      |>.push (recommendedSpellingExt.exportEntriesFn (← getEnv) (recommendedSpellingExt.getState (← getEnv)) .exported)
   return all.flatMap id
 
 end Lean.Elab.Term.Doc

src/Lean/Elab/Tactic/Doc.lean

@@ -98,7 +98,7 @@ Displays all available tactic tags, with documentation.
 @[builtin_command_elab printTacTags] def elabPrintTacTags : CommandElab := fun _stx => do
   let all :=
     tacticTagExt.toEnvExtension.getState (← getEnv)
-      |>.importedEntries |>.push (tacticTagExt.exportEntriesFn (tacticTagExt.getState (← getEnv)))
+      |>.importedEntries |>.push (tacticTagExt.exportEntriesFn (← getEnv) (tacticTagExt.getState (← getEnv)) .exported)
   let mut mapping : NameMap NameSet := {}
   for arr in all do
     for (tac, tag) in arr do
@@ -148,7 +148,7 @@ def allTacticDocs : MetaM (Array TacticDoc) := do
   let env ← getEnv
   let all :=
     tacticTagExt.toEnvExtension.getState (← getEnv)
-      |>.importedEntries |>.push (tacticTagExt.exportEntriesFn (tacticTagExt.getState (← getEnv)))
+      |>.importedEntries |>.push (tacticTagExt.exportEntriesFn (← getEnv) (tacticTagExt.getState (← getEnv)) .exported)
   let mut tacTags : NameMap NameSet := {}
   for arr in all do
     for (tac, tag) in arr do

src/Lean/EnvExtension.lean

@@ -21,13 +21,10 @@ structure SimplePersistentEnvExtensionDescr (α σ : Type) where
   addEntryFn    : σ → α → σ
   addImportedFn : Array (Array α) → σ
   toArrayFn     : List α → Array α := fun es => es.toArray
+  exportEntriesFnEx? :
+    Option (Environment → σ → List α → OLeanLevel → Array α) := none
   asyncMode     : EnvExtension.AsyncMode := .mainOnly
   replay?       : Option ((newEntries : List α) → (newState : σ) → σ → List α × σ) := none
-  /--
-  Whether entries should be imported into other modules. Entries are always accessible in the
-  language server via `getModuleEntries (includeServer := true)`.
-  -/
-  exported      : Bool := true
 
 /--
 Returns a function suitable for `SimplePersistentEnvExtensionDescr.replay?` that replays all new
@@ -47,8 +44,9 @@ def registerSimplePersistentEnvExtension {α σ : Type} [Inhabited σ] (descr :
     addImportedFn   := fun as => pure ([], descr.addImportedFn as),
     addEntryFn      := fun s e => match s with
       | (entries, s) => (e::entries, descr.addEntryFn s e),
-    exportEntriesFn := fun s => if descr.exported then descr.toArrayFn s.1.reverse else #[],
-    saveEntriesFn := fun s => descr.toArrayFn s.1.reverse,
+    exportEntriesFnEx env s level := match descr.exportEntriesFnEx? with
+      | some fn => fn env s.2 s.1.reverse level
+      | none    => descr.toArrayFn s.1.reverse
     statsFn := fun s => format "number of local entries: " ++ format s.1.length
     asyncMode := descr.asyncMode
     replay? := descr.replay?.map fun replay oldState newState _ (entries, s) =>
@@ -129,14 +127,15 @@ structure MapDeclarationExtension (α : Type) extends PersistentEnvExtension (Na
 deriving Inhabited
 
 def mkMapDeclarationExtension (name : Name := by exact decl_name%)
-    (exportEntriesFn : NameMap α → Array (Name × α) := (·.toArray)) : IO (MapDeclarationExtension α) :=
+    (exportEntriesFn : Environment → NameMap α → OLeanLevel → Array (Name × α) :=
+      fun _ s _ => s.toArray) :
+    IO (MapDeclarationExtension α) :=
   .mk <$> registerPersistentEnvExtension {
     name            := name,
     mkInitial       := pure {}
     addImportedFn   := fun _ => pure {}
     addEntryFn      := fun s (n, v) => s.insert n v
-    saveEntriesFn   := fun s => s.toArray
-    exportEntriesFn
+    exportEntriesFnEx env s level := exportEntriesFn env s level
     asyncMode       := .async
     replay?         := some fun _ newState newConsts s =>
       newConsts.foldl (init := s) fun s c =>
@@ -156,10 +155,10 @@ def insert (ext : MapDeclarationExtension α) (env : Environment) (declName : Na
     ext.addEntry env (declName, val)
 
 def find? [Inhabited α] (ext : MapDeclarationExtension α) (env : Environment) (declName : Name)
-    (includeServer := false) : Option α :=
+    (level := OLeanLevel.exported) : Option α :=
   match env.getModuleIdxFor? declName with
   | some modIdx =>
-    match (ext.getModuleEntries (includeServer := includeServer) env modIdx).binSearch (declName, default) (fun a b => Name.quickLt a.1 b.1) with
+    match (ext.getModuleEntries (level := level) env modIdx).binSearch (declName, default) (fun a b => Name.quickLt a.1 b.1) with
     | some e => some e.2
     | none   => none
   | none => (ext.findStateAsync env declName).find? declName

src/Lean/Environment.lean

@@ -526,11 +526,10 @@ structure Environment where
   -/
   private base : VisibilityMap Kernel.Environment
   /--
-  Additional imported environment extension state for use in the language server. This field is
-  identical to `base.extensions` in other contexts. Access via
-  `getModuleEntries (includeServer := true)`.
+  Additional imported environment extension state for use in the language server, `none` elsewhere.
+  Access via `getModuleEntries (level := .server)`.
   -/
-  private serverBaseExts : Array EnvExtensionState := base.private.extensions
+  private serverBaseExts? : Option (Array EnvExtensionState) := none
   /--
   Kernel environment task that is fulfilled when all asynchronously elaborated declarations are
   finished, containing the resulting environment. Also collects the environment extension state of
@@ -1457,6 +1456,19 @@ structure ImportM.Context where
 
 abbrev ImportM := ReaderT Lean.ImportM.Context IO
 
+/-- The level of information to save/load. Each level includes all previous ones. -/
+inductive OLeanLevel where
+  /-- Information from exported contexts. -/
+  | exported
+  /-- Environment extension state for the language server. -/
+  | server
+  /-- Private module data. -/
+  | «private»
+deriving DecidableEq, Ord, Repr
+
+instance : LE OLeanLevel := leOfOrd
+instance : LT OLeanLevel := ltOfOrd
+
 /--
 An environment extension with support for storing/retrieving entries from a .olean file.
  - α is the type of the entries that are stored in .olean files.
@@ -1507,15 +1519,17 @@ structure PersistentEnvExtension (α : Type) (β : Type) (σ : Type) where
   name            : Name
   addImportedFn   : Array (Array α) → ImportM σ
   addEntryFn      : σ → β → σ
-  /-- Function to transform state into data that should always be imported into other modules. -/
-  exportEntriesFn : σ → Array α
   /--
-  Function to transform state into data that should be imported into other modules when the module
-  system is disabled. When it is enabled, the data is loaded only in the language server and
-  accessible via `getModuleEntries (includeServer := true)`. Conventionally, this is a superset of
-  the data returned by `exportEntriesFn`.
+  Function to transform state into data that should be imported into other modules. When using the
+  module system without `import all`, `OLeanLevel.exported` is imported, else `OLeanLevel.private`.
+  Additionally, when using the module system in the language server, the `OLeanLevel.server` data is
+  accessible via `getModuleEntries (level := .server)`. By convention, each level should include all
+  data of previous levels.
+
+  This function is run after elaborating the file and joining all asynchronous threads. It is run
+  once for each level when the module system is enabled, otherwise once for `private`.
   -/
-  saveEntriesFn   : σ → Array α
+  exportEntriesFn : Environment → σ → OLeanLevel → Array α
   statsFn         : σ → Format
 
 instance {α σ} [Inhabited σ] : Inhabited (PersistentEnvExtensionState α σ) :=
@@ -1527,20 +1541,26 @@ instance {α β σ} [Inhabited σ] : Inhabited (PersistentEnvExtension α β σ)
      name := default,
      addImportedFn := fun _ => default,
      addEntryFn := fun s _ => s,
-     exportEntriesFn := fun _ => #[],
-     saveEntriesFn := fun _ => #[],
+     exportEntriesFn := fun _ _ _ => #[],
      statsFn := fun _ => Format.nil
   }
 
 namespace PersistentEnvExtension
 
 /--
-Returns the data saved by `ext.exportEntriesFn/saveEntriesFn` when `m` was elaborated. See docs on
-the functions for details.
+Returns the data saved by `ext.exportEntriesFn` when `m` was elaborated. See docs on the function for
+details. When using the module system, `level` can be used to select the level of data to retrieve,
+but is limited to the maximum level actually imported: `exported` on the cmdline and `server` in the
+language server.
 -/
 def getModuleEntries {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ)
-    (env : Environment) (m : ModuleIdx) (includeServer := false) : Array α :=
-  let exts := if includeServer then env.serverBaseExts else env.base.private.extensions
+    (env : Environment) (m : ModuleIdx) (level := OLeanLevel.exported) : Array α :=
+  let exts := match level with
+    | .exported => env.base.private.extensions
+    | .server   => match env.serverBaseExts? with
+      | some exts => exts
+      | none      => panic! "cannot use `server` level outside the language server"
+    | level     => panic! s!"unsupported level {repr level}, expected `exported` or `server`"
   -- safety: as in `getStateUnsafe`
   unsafe (ext.toEnvExtension.getStateImpl exts).importedEntries[m]!
 
@@ -1573,19 +1593,36 @@ end PersistentEnvExtension
 
 builtin_initialize persistentEnvExtensionsRef : IO.Ref (Array (PersistentEnvExtension EnvExtensionEntry EnvExtensionEntry EnvExtensionState)) ← IO.mkRef #[]
 
-structure PersistentEnvExtensionDescr (α β σ : Type) where
-  name            : Name := by exact decl_name%
-  mkInitial       : IO σ
-  addImportedFn   : Array (Array α) → ImportM σ
-  addEntryFn      : σ → β → σ
-  exportEntriesFn : σ → Array α
-  saveEntriesFn   : σ → Array α := exportEntriesFn
-  statsFn         : σ → Format := fun _ => Format.nil
-  asyncMode       : EnvExtension.AsyncMode := .mainOnly
-  replay?         : Option (ReplayFn σ) := none
+-- Helper structure to enable cyclic default values of `exportEntriesFn` and `exportEntriesFnEx`.
+structure PersistentEnvExtensionDescrCore (α β σ : Type) where
+  name              : Name := by exact decl_name%
+  mkInitial         : IO σ
+  addImportedFn     : Array (Array α) → ImportM σ
+  addEntryFn        : σ → β → σ
+  exportEntriesFnEx : Environment → σ → OLeanLevel → Array α
+  statsFn           : σ → Format := fun _ => Format.nil
+  asyncMode         : EnvExtension.AsyncMode := .mainOnly
+  replay?           : Option (ReplayFn σ) := none
 
-attribute [inherit_doc PersistentEnvExtension.exportEntriesFn] PersistentEnvExtensionDescr.exportEntriesFn
-attribute [inherit_doc PersistentEnvExtension.saveEntriesFn] PersistentEnvExtensionDescr.saveEntriesFn
+attribute [inherit_doc PersistentEnvExtension.exportEntriesFn]
+  PersistentEnvExtensionDescrCore.exportEntriesFnEx
+
+/--
+Auxiliary function to signal to the structure instance elaborator that `default` should be used as
+the default value for a field but only if `_otherField` has been given, which is added as an
+artifical dependency.
+-/
+def useDefaultIfOtherFieldGiven (default : α) (_otherField : β) : α :=
+  default
+
+structure PersistentEnvExtensionDescr (α β σ : Type) extends PersistentEnvExtensionDescrCore α β σ where
+  -- The cyclic default values force the user to specify at least one of the two following fields.
+  /--
+  Obsolete simpler version of `exportEntriesFnEx`. Its value is ignored if the latter is also
+  specified.
+  -/
+  exportEntriesFn : σ → Array α := useDefaultIfOtherFieldGiven (fun _ => #[]) exportEntriesFnEx
+  exportEntriesFnEx := fun _ s _ => exportEntriesFn s
 
 unsafe def registerPersistentEnvExtensionUnsafe {α β σ : Type} [Inhabited σ] (descr : PersistentEnvExtensionDescr α β σ) : IO (PersistentEnvExtension α β σ) := do
   let pExts ← persistentEnvExtensionsRef.get
@@ -1604,8 +1641,7 @@ unsafe def registerPersistentEnvExtensionUnsafe {α β σ : Type} [Inhabited σ]
     name            := descr.name,
     addImportedFn   := descr.addImportedFn,
     addEntryFn      := descr.addEntryFn,
-    exportEntriesFn := descr.exportEntriesFn,
-    saveEntriesFn := descr.saveEntriesFn,
+    exportEntriesFn := descr.exportEntriesFnEx,
     statsFn         := descr.statsFn
   }
   persistentEnvExtensionsRef.modify fun pExts => pExts.push (unsafeCast pExt)
@@ -1661,16 +1697,6 @@ unsafe def Environment.freeRegions (env : Environment) : IO Unit :=
     TODO: statically check for this. -/
   env.header.regions.forM CompactedRegion.free
 
-/-- The level of information to save/load. Each level includes all previous ones. -/
-inductive OLeanLevel where
-  /-- Information from exported contexts. -/
-  | exported
-  /-- Environment extension state for the language server. -/
-  | server
-  /-- Private module data. -/
-  | «private»
-deriving DecidableEq
-
 def OLeanLevel.adjustFileName (base : System.FilePath) : OLeanLevel → System.FilePath
   | .exported => base
   | .server   => base.addExtension "server"
@@ -1688,7 +1714,7 @@ def mkModuleData (env : Environment) (level : OLeanLevel := .private) : IO Modul
     if asyncMode matches .async then
       asyncMode := .sync
     let state := pExt.getState (asyncMode := asyncMode) env
-    (pExt.name, if level = .exported then pExt.exportEntriesFn state else pExt.saveEntriesFn state)
+    (pExt.name, pExt.exportEntriesFn env state level)
   let kenv := env.toKernelEnv
   let env := env.setExporting (level != .private)
   let constNames := kenv.constants.foldStage2 (fun names name _ => names.push name) #[]
@@ -2040,8 +2066,9 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
     realizedImportedConsts? := none
   }
   env := env.setCheckedSync { env.base.private with extensions := (← setImportedEntries env.base.private.extensions moduleData) }
-  let serverData := modules.filterMap (·.serverData? level)
-  env := { env with serverBaseExts := (← setImportedEntries env.base.private.extensions serverData) }
+  if level ≥ .server then
+    let serverData := modules.filterMap (·.serverData? level)
+    env := { env with serverBaseExts? := some (← setImportedEntries env.base.private.extensions serverData) }
   if leakEnv then
     /- Mark persistent a first time before `finalizePersistenExtensions`, which
        avoids costly MT markings when e.g. an interpreter closure (which

tests/pkg/module/Module.lean

@@ -1,4 +1,3 @@
-import Lean
 import Module.Basic
 import Module.Imported
 import Module.ImportedAll
@@ -11,14 +10,7 @@ import Module.NonModule
 
 open Lean
 
-/-! Non-essential metadata should only be accessible at level >= .server -/
-
-#eval show IO Unit from do
-  let env ← importModules (level := .exported) #[`Module.Basic] {}
-  assert! env.header.isModule
-  let _ ← Core.CoreM.toIO (ctx := { fileName := "module.lean", fileMap := default }) (s := { env }) do
-    assert! (← findDeclarationRanges? ``f).isNone
-    assert! (getModuleDoc? (← getEnv) `Module.Basic).any (·.size == 0)
+/-! Non-essential metadata should be accessible at level >= .server -/
 
 #eval show IO Unit from do
   let env ← importModules (level := .server) #[`Module.Basic] {}