feat: move non-essential metadata into .olean.server (#8068)

This PR ensures that for modules opted into the experimental module
system, we do not import module docstrings or declaration ranges.

Excluding declaration docstrings as well would require some more work to
make `[inherit_doc]` leave a mere reference to the other declaration
instead of copying its docstring eagerly.

script/benchReelabRss.lean

@@ -1,4 +1,5 @@
 import Lean.Data.Lsp
+import Lean.Elab.Import
 open Lean
 open Lean.Lsp
 open Lean.JsonRpc
@@ -7,9 +8,7 @@ open Lean.JsonRpc
 Tests language server memory use by repeatedly re-elaborate a given file.
 
 NOTE: only works on Linux for now.
-
-HACK: The line that is to be prepended with a space is hard-coded below to be sufficiently far down
-not to touch the imports for usual files.
+ot to touch the imports for usual files.
 -/
 
 def main (args : List String) : IO Unit := do
@@ -33,6 +32,8 @@ def main (args : List String) : IO Unit := do
     Ipc.writeRequest ⟨0, "initialize", { capabilities : InitializeParams }⟩
 
     let text ← IO.FS.readFile file
+    let (_, headerEndPos, _) ← Elab.parseImports text
+    let headerEndPos := FileMap.ofString text |>.leanPosToLspPos headerEndPos
     let mut requestNo : Nat := 1
     let mut versionNo : Nat := 1
     Ipc.writeNotification ⟨"textDocument/didOpen", {
@@ -40,15 +41,14 @@ def main (args : List String) : IO Unit := do
     for i in [0:iters.toNat!] do
       if i > 0 then
         versionNo := versionNo + 1
-        let pos := { line := 19, character := 0 }
         let params : DidChangeTextDocumentParams := {
           textDocument := {
             uri      := uri
             version? := versionNo
           }
           contentChanges := #[TextDocumentContentChangeEvent.rangeChange {
-            start := pos
-            «end» := pos
+            start := headerEndPos
+            «end» := headerEndPos
           } " "]
         }
         let params := toJson params

src/Lean/DeclarationRange.lean

@@ -14,7 +14,8 @@ import Lean.MonadEnv
 namespace Lean
 
 builtin_initialize builtinDeclRanges : IO.Ref (NameMap DeclarationRanges) ← IO.mkRef {}
-builtin_initialize declRangeExt : MapDeclarationExtension DeclarationRanges ← mkMapDeclarationExtension
+builtin_initialize declRangeExt : MapDeclarationExtension DeclarationRanges ←
+  mkMapDeclarationExtension (exportEntriesFn := fun _ => #[])
 
 def addBuiltinDeclarationRanges (declName : Name) (declRanges : DeclarationRanges) : IO Unit :=
   builtinDeclRanges.modify (·.insert declName declRanges)
@@ -23,7 +24,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? (← getEnv) declName
+  return declRangeExt.find? (includeServer := true) (← 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

@@ -59,6 +59,7 @@ private builtin_initialize moduleDocExt : SimplePersistentEnvExtension ModuleDoc
   addImportedFn := fun _ => {}
   addEntryFn    := fun s e => s.push e
   toArrayFn     := fun es => es.toArray
+  exported      := false
 }
 
 def addMainModuleDoc (env : Environment) (doc : ModuleDoc) : Environment :=
@@ -68,7 +69,8 @@ def getMainModuleDoc (env : Environment) : PersistentArray ModuleDoc :=
   moduleDocExt.getState env
 
 def getModuleDoc? (env : Environment) (moduleName : Name) : Option (Array ModuleDoc) :=
-  env.getModuleIdx? moduleName |>.map fun modIdx => moduleDocExt.getModuleEntries env modIdx
+  env.getModuleIdx? moduleName |>.map fun modIdx =>
+    moduleDocExt.getModuleEntries (includeServer := true) env modIdx
 
 def getDocStringText [Monad m] [MonadError m] (stx : TSyntax `Lean.Parser.Command.docComment) : m String :=
   match stx.raw[1] with

src/Lean/EnvExtension.lean

@@ -23,6 +23,11 @@ structure SimplePersistentEnvExtensionDescr (α σ : Type) where
   toArrayFn     : List α → Array α := fun es => es.toArray
   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
@@ -42,7 +47,8 @@ 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 => descr.toArrayFn s.1.reverse,
+    exportEntriesFn := fun s => if descr.exported then descr.toArrayFn s.1.reverse else #[],
+    saveEntriesFn := fun s => 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) =>

src/Lean/Environment.lean

@@ -1689,9 +1689,10 @@ private def ImportedModule.mainModule? (self : ImportedModule) : Option ModuleDa
   self.parts[if baseMod.isModule && self.importAll then 2 else 0]?.map (·.1)
 
 /-- The module data that should be used for server purposes. -/
-private def ImportedModule.serverData? (self : ImportedModule) : Option ModuleData := do
+private def ImportedModule.serverData? (self : ImportedModule) (level : OLeanLevel) :
+    Option ModuleData := do
   let (baseMod, _) ← self.parts[0]?
-  self.parts[if baseMod.isModule then 1 else 0]?.map (·.1)
+  self.parts[if baseMod.isModule && level != .exported then 1 else 0]?.map (·.1)
 
 structure ImportState where
   private moduleNameMap : Std.HashMap Name ImportedModule := {}
@@ -1771,7 +1772,7 @@ Constructs environment from `importModulesCore` results.
 See also `importModules` for parameter documentation.
 -/
 def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0)
-    (leakEnv loadExts : Bool) (isModule := false) : IO Environment := do
+    (leakEnv loadExts : Bool) (level := OLeanLevel.private) : IO Environment := do
   let modules := s.moduleNames.filterMap (s.moduleNameMap[·]?)
   let moduleData ← modules.mapM fun mod => do
     let some data := mod.mainModule? |
@@ -1803,7 +1804,8 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
       extraConstNames := {}
       extensions      := exts
       header     := {
-        trustLevel, isModule, imports, moduleData
+        trustLevel, imports, moduleData
+        isModule := level == .exported
         regions      := modules.flatMap (·.parts.map (·.2))
         moduleNames  := s.moduleNames
       }
@@ -1811,7 +1813,7 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
     realizedImportedConsts? := none
   }
   env := env.setCheckedSync { env.base with extensions := (← setImportedEntries env.base.extensions moduleData) }
-  let serverData := modules.filterMap (·.serverData?)
+  let serverData := modules.filterMap (·.serverData? level)
   env := { env with serverBaseExts := (← setImportedEntries env.base.extensions serverData) }
   if leakEnv then
     /- Mark persistent a first time before `finalizePersistenExtensions`, which
@@ -1871,7 +1873,7 @@ def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32
   withImporting do
     plugins.forM Lean.loadPlugin
     let (_, s) ← importModulesCore (forceImportAll := level == .private) imports |>.run
-    finalizeImport (leakEnv := leakEnv) (loadExts := loadExts) (isModule := level != .private)
+    finalizeImport (leakEnv := leakEnv) (loadExts := loadExts) (level := level)
       s imports opts trustLevel
 
 /--

tests/lean/run/module.lean

@@ -1,3 +0,0 @@
-module
-
-/-! # Module system basic tests -/

tests/pkg/module/Module.lean

@@ -0,0 +1,23 @@
+module
+
+import Lean
+import Module.Basic
+
+/-! # Module system basic tests -/
+
+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)
+
+#eval show IO Unit from do
+  let env ← importModules (level := .server) #[`Module.Basic] {}
+  let _ ← Core.CoreM.toIO (ctx := { fileName := "module.lean", fileMap := default }) (s := { env }) do
+    assert! (← findDeclarationRanges? ``f).isSome
+    assert! (getModuleDoc? (← getEnv) `Module.Basic).any (·.size == 1)

tests/pkg/module/Module/Basic.lean

@@ -0,0 +1,9 @@
+module
+
+/-! Module docstring -/
+
+/-- A definition. -/
+def f := 1
+
+/-- A theorem. -/
+theorem t : f = f := rfl

tests/pkg/module/lakefile.toml

@@ -0,0 +1,6 @@
+name = "module"
+defaultTargets = ["Module"]
+
+[[lean_lib]]
+name = "Module"
+leanOptions = { experimental.module = true }

tests/pkg/module/lean-toolchain

@@ -0,0 +1 @@
+lean4

tests/pkg/module/test.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+
+rm -rf .lake/build
+LEAN_ABORT_ON_PANIC=1 lake build