feat: module header keyword for enabling module system

2025-04-14 18:08:19 +02:00 · 2025-04-14 18:08:19 +02:00 · da82cbd3d1
commit da82cbd3d1
parent 2386a3d7c7
13 changed files with 77 additions and 62 deletions
--- a/nix/buildLeanPackage.nix
+++ b/nix/buildLeanPackage.nix
@ -206,7 +206,7 @@ with builtins; let
  # subset of `modCandidates` that is transitively reachable from `roots`
  mods' = listToAttrs (map (e: { name = e.key; value = modCandidates.${e.key}; }) (genericClosure {
    startSet = map (m: { key = m; }) (concatMap expandGlob roots);
-    operator = e: if modDepsMap ? ${e.key} then map (m: { key = m.module; }) (filter (m: modCandidates ? ${m.module}) modDepsMap.${e.key}.imports) else [];
+    operator = e: if modDepsMap ? ${e.key} then map (m: { key = m.module; }) (filter (m: modCandidates ? ${m.module}) modDepsMap.${e.key}.result.imports) else [];
  }));
  allLinkFlags = lib.foldr (shared: acc: acc ++ [ "-L${shared}" "-l${shared.linkName or shared.name}" ]) linkFlags allNativeSharedLibs;

--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@ -132,6 +132,11 @@ def process (input : String) (env : Environment) (opts : Options) (fileName : Op
  let s ← IO.processCommands inputCtx { : Parser.ModuleParserState } (Command.mkState env {} opts)
  pure (s.commandState.env, s.commandState.messages)

+register_builtin_option experimental.module : Bool := {
+  defValue := false
+  descr := "Allow use of module system (experimental)"
+}
+
@[export lean_run_frontend]
 def runFrontend
    (input : String)
@ -151,7 +156,10 @@ def runFrontend
  let opts := Elab.async.setIfNotSet opts true
  let ctx := { inputCtx with }
  let processor := Language.Lean.process
-  let snap ← processor (fun _ => pure <| .ok { mainModuleName, opts, trustLevel, plugins }) none ctx
+  let snap ← processor (fun header => do
+    if !header.raw[0].isNone && !experimental.module.get opts then
+      throw <| IO.Error.userError "`module` keyword is experimental and not enabled here"
+    pure <| .ok { mainModuleName, opts, trustLevel, plugins }) none ctx
  let snaps := Language.toSnapshotTree snap
  let severityOverrides := errorOnKinds.foldl (·.insert · .error) {}
  let hasErrors ← snaps.runAndReport opts jsonOutput severityOverrides
--- a/src/Lean/Elab/Import.lean
+++ b/src/Lean/Elab/Import.lean
@ -10,19 +10,18 @@ import Lean.CoreM

 namespace Lean.Elab

-def headerToImports (header : Syntax) : Array Import :=
-  let imports := if header[0].isNone then #[{ module := `Init : Import }] else #[]
-  imports ++ header[1].getArgs.map fun stx =>
-    -- `stx` is of the form `(Module.import "import" "runtime"? id)
-    let runtime := !stx[1].isNone
-    let id      := stx[2].getId
-    { module := id, runtimeOnly := runtime }
+def headerToImports : TSyntax ``Parser.Module.header → Array Import
+  | `(Parser.Module.header| $[module%$moduleTk]? $[prelude%$preludeTk]? $[import $ns]*) =>
+    let imports := if preludeTk.isNone then #[{ module := `Init : Import }] else #[]
+    imports ++ ns.map ({ module := ·.getId })
+  | _ => unreachable!

-def processHeader (header : Syntax) (opts : Options) (messages : MessageLog)
+def processHeader (header : TSyntax ``Parser.Module.header) (opts : Options) (messages : MessageLog)
    (inputCtx : Parser.InputContext) (trustLevel : UInt32 := 0)
    (plugins : Array System.FilePath := #[]) (leakEnv := false)
    : IO (Environment × MessageLog) := do
-  let level := if experimental.module.get opts then
+  let isModule := !header.raw[0].isNone
+  let level := if isModule then
    if Elab.inServer.get opts then
      .server
    else
@ -35,7 +34,7 @@ def processHeader (header : Syntax) (opts : Options) (messages : MessageLog)
    pure (env, messages)
  catch e =>
    let env ← mkEmptyEnvironment
-    let spos := header.getPos?.getD 0
+    let spos := header.raw.getPos?.getD 0
    let pos  := inputCtx.fileMap.toPosition spos
    pure (env, messages.add { fileName := inputCtx.fileName, data := toString e, pos := pos })

--- a/src/Lean/Elab/ParseImportsFast.lean
+++ b/src/Lean/Elab/ParseImportsFast.lean
@ -10,9 +10,10 @@ namespace Lean
 namespace ParseImports

 structure State where
-  imports : Array Import := #[]
-  pos     : String.Pos := 0
-  error?  : Option String := none
+  imports  : Array Import := #[]
+  pos      : String.Pos := 0
+  error?   : Option String := none
+  isModule : Bool := false
  deriving Inhabited

 def Parser := String → State → State
@ -132,8 +133,8 @@ partial def whitespace : Parser := fun input s =>
  else
    false

-def State.pushModule (module : Name) (runtimeOnly : Bool) (s : State) : State :=
-  { s with imports := s.imports.push { module, runtimeOnly } }
+def State.pushModule (module : Name) (s : State) : State :=
+  { s with imports := s.imports.push { module } }

@[inline] def isIdRestCold (c : Char) : Bool :=
  c = '_' || c = '\'' || c == '!' || c == '?' || isLetterLike c || isSubScriptAlnum c
@ -141,7 +142,7 @@ def State.pushModule (module : Name) (runtimeOnly : Bool) (s : State) : State :=
@[inline] def isIdRestFast (c : Char) : Bool :=
  c.isAlphanum || (c != '.' && c != '\n' && c != ' ' && isIdRestCold c)

-partial def moduleIdent (runtimeOnly : Bool) : Parser := fun input s =>
+partial def moduleIdent : Parser := fun input s =>
  let rec parse (module : Name) (s : State) :=
    let i := s.pos
    if h : input.atEnd i then
@ -161,7 +162,7 @@ partial def moduleIdent (runtimeOnly : Bool) : Parser := fun input s =>
            let s := s.next input s.pos
            parse module s
          else
-            whitespace input (s.pushModule module runtimeOnly)
+            whitespace input (s.pushModule module)
      else if isIdFirst curr then
        let startPart := i
        let s         := takeWhile isIdRestFast input (s.next' input i h)
@ -171,7 +172,7 @@ partial def moduleIdent (runtimeOnly : Bool) : Parser := fun input s =>
          let s := s.next input s.pos
          parse module s
        else
-          whitespace input (s.pushModule module runtimeOnly)
+          whitespace input (s.pushModule module)
      else
        s.mkError "expected identifier"
  parse .anonymous s
@ -184,28 +185,31 @@ partial def moduleIdent (runtimeOnly : Bool) : Parser := fun input s =>
  | none => many p input s
  | some _ => { pos, error? := none, imports := s.imports.shrink size }

-@[inline] partial def preludeOpt (k : String) : Parser :=
-  keywordCore k (fun _ s => s.pushModule `Init false) (fun _ s => s)
-
 def main : Parser :=
-  preludeOpt "prelude" >>
-  many (keyword "import" >> keywordCore "runtime" (moduleIdent false) (moduleIdent true))
+  keywordCore "module" (fun _ s => { s with isModule := true }) (fun _ s => s) >>
+  keywordCore "prelude" (fun _ s => s.pushModule `Init) (fun _ s => s) >>
+  many (keyword "import" >> moduleIdent)

 end ParseImports

+deriving instance ToJson for Import
+
+structure ParseImportsResult where
+  imports  : Array Import
+  isModule : Bool
+  deriving ToJson
+
 /--
 Simpler and faster version of `parseImports`. We use it to implement Lake.
 -/
-def parseImports' (input : String) (fileName : String) : IO (Array Lean.Import) := do
+def parseImports' (input : String) (fileName : String) : IO ParseImportsResult := do
  let s := ParseImports.main input (ParseImports.whitespace input {})
  match s.error? with
-  | none => return s.imports
+  | none => return { s with }
  | some err => throw <| IO.userError s!"{fileName}: {err}"

-deriving instance ToJson for Import
-
 structure PrintImportResult where
-  imports? : Option (Array Import) := none
+  result?  : Option ParseImportsResult := none
  errors   : Array String := #[]
  deriving ToJson

@ -217,8 +221,8 @@ structure PrintImportsResult where
 def printImportsJson (fileNames : Array String) : IO Unit := do
  let rs ← fileNames.mapM fun fn => do
    try
-      let deps ← parseImports' (← IO.FS.readFile ⟨fn⟩) fn
-      return { imports? := some deps }
+      let res ← parseImports' (← IO.FS.readFile ⟨fn⟩) fn
+      return { result? := some res }
    catch e => return { errors := #[e.toString] }
  IO.println (toJson { imports := rs : PrintImportsResult } |>.compress)

--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@ -95,12 +95,11 @@ abbrev ConstMap := SMap Name ConstantInfo

 structure Import where
  module      : Name
-  runtimeOnly : Bool := false
  deriving Repr, Inhabited

 instance : Coe Name Import := ⟨({module := ·})⟩

-instance : ToString Import := ⟨fun imp => toString imp.module ++ if imp.runtimeOnly then " (runtime)" else ""⟩
+instance : ToString Import := ⟨fun imp => toString imp.module⟩

 /--
  A compacted region holds multiple Lean objects in a contiguous memory region, which can be read/written to/from disk.
@ -123,6 +122,8 @@ instance : Nonempty EnvExtensionEntry := EnvExtensionEntrySpec.property
 /-- Content of a .olean file.
   We use `compact.cpp` to generate the image of this object in disk. -/
 structure ModuleData where
+  /-- Participating in the module system? -/
+  isModule        : Bool
  imports         : Array Import
  /--
  `constNames` contains all constant names in `constants`.
@ -152,6 +153,8 @@ structure EnvironmentHeader where
  Name of the module being compiled.
  -/
  mainModule   : Name         := default
+  /-- Participating in the module system? -/
+  isModule     : Bool         := false
  /-- Direct imports -/
  imports      : Array Import := #[]
  /-- Compacted regions for all imported modules. Objects in compacted memory regions do no require any memory management. -/
@ -1581,17 +1584,11 @@ def mkModuleData (env : Environment) (level : OLeanLevel := .private) : IO Modul
  -- TODO: does not include cstage* constants from the old codegen
  --let constants := constNames.filterMap env.find?
  let constNames := constants.map (·.name)
-  return {
-    imports         := env.header.imports
+  return { env.header with
    extraConstNames := env.checked.get.extraConstNames.toArray
    constNames, constants, entries
  }

-register_builtin_option experimental.module : Bool := {
-  defValue := false
-  descr := "Enable module system (experimental)"
-}
-
@[export lean_write_module]
 def writeModule (env : Environment) (fname : System.FilePath) (split := false) : IO Unit := do
  if split then
@ -1699,7 +1696,7 @@ abbrev ImportStateM := StateRefT ImportState IO
 partial def importModulesCore (imports : Array Import) (level := OLeanLevel.private) :
    ImportStateM Unit := do
  for i in imports do
-    if i.runtimeOnly || (← get).moduleNameSet.contains i.module then
+    if (← get).moduleNameSet.contains i.module then
      continue
    modify fun s => { s with moduleNameSet := s.moduleNameSet.insert i.module }
    let mFile ← findOLean i.module
@ -1756,7 +1753,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) : IO Environment := do
+    (leakEnv loadExts : Bool) (isModule := false) : IO Environment := do
  let numConsts := s.moduleData.foldl (init := 0) fun numConsts mod =>
    numConsts + mod.constants.size + mod.extraConstNames.size
  let mut const2ModIdx : Std.HashMap Name ModuleIdx := Std.HashMap.emptyWithCapacity (capacity := numConsts)
@ -1783,7 +1780,7 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
      extraConstNames := {}
      extensions      := exts
      header     := {
-        trustLevel, imports
+        trustLevel, isModule, imports
        regions      := s.parts.flatMap (·.map (·.2))
        moduleNames  := s.moduleNames
        moduleData   := s.moduleData
@ -1847,7 +1844,8 @@ def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32
  withImporting do
    plugins.forM Lean.loadPlugin
    let (_, s) ← importModulesCore (level := level) imports |>.run
-    finalizeImport (leakEnv := leakEnv) (loadExts := loadExts) s imports opts trustLevel
+    finalizeImport (leakEnv := leakEnv) (loadExts := loadExts) (isModule := !level matches .private)
+      s imports opts trustLevel

 /--
 Creates environment object from imports and frees compacted regions after calling `act`. No live
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@ -361,7 +361,7 @@ General notes:
  the `sync` parameter on `parseCmd` and spawn an elaboration task when we leave it.
 -/
 partial def process
-    (setupImports : Syntax → ProcessingT IO (Except HeaderProcessedSnapshot SetupImportsResult))
+    (setupImports : TSyntax ``Parser.Module.header → ProcessingT IO (Except HeaderProcessedSnapshot SetupImportsResult))
    (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
  parseHeader old? |>.run (old?.map (·.ictx))
 where
@ -423,7 +423,7 @@ where
          result? := none
        }

-      let trimmedStx := stx.unsetTrailing
+      let trimmedStx := stx.raw.unsetTrailing
      -- semi-fast path: go to next snapshot if syntax tree is unchanged
      -- NOTE: We compare modulo `unsetTrailing` in order to ensure that changes in trailing
      -- whitespace do not invalidate the header. This is safe because we only pass the trimmed
@ -443,11 +443,11 @@ where
        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
        result? := some {
          parserState
-          processedSnap := (← processHeader trimmedStx parserState)
+          processedSnap := (← processHeader ⟨trimmedStx⟩ parserState)
        }
      }

-  processHeader (stx : Syntax) (parserState : Parser.ModuleParserState) :
+  processHeader (stx : TSyntax ``Parser.Module.header) (parserState : Parser.ModuleParserState) :
      LeanProcessingM (SnapshotTask HeaderProcessedSnapshot) := do
    let ctx ← read
    SnapshotTask.ofIO stx none (some ⟨0, ctx.input.endPos⟩) <|
@ -489,7 +489,7 @@ where
            ngen    := { namePrefix := `_import }
          }) (Elab.InfoTree.node
              (Elab.Info.ofCommandInfo { elaborator := `header, stx })
-              (stx[1].getArgs.toList.map (fun importStx =>
+              (stx.raw[2].getArgs.toList.map (fun importStx =>
                Elab.InfoTree.node (Elab.Info.ofCommandInfo {
                  elaborator := `import
                  stx := importStx
--- a/src/Lean/Parser/Module.lean
+++ b/src/Lean/Parser/Module.lean
@ -11,9 +11,14 @@ namespace Lean
 namespace Parser

 namespace Module
+def moduleTk   := leading_parser "module"
 def «prelude»  := leading_parser "prelude"
-def «import»   := leading_parser "import " >> optional "runtime" >> identWithPartialTrailingDot
-def header     := leading_parser optional («prelude» >> ppLine) >> many («import» >> ppLine) >> ppLine
+def «import»   := leading_parser "import " >> identWithPartialTrailingDot
+def header     := leading_parser optional (moduleTk >> ppLine >> ppLine) >>
+  optional («prelude» >> ppLine) >>
+  many («import» >> ppLine) >>
+  ppLine
+
 /--
  Parser for a Lean module. We never actually run this parser but instead use the imperative definitions below that
  return the same syntax tree structure, but add error recovery. Still, it is helpful to have a `Parser` definition
@ -64,7 +69,7 @@ where
      if let .original (trailing := trailing) .. := stx.getTailInfo then pure (some trailing)
        else none

-def parseHeader (inputCtx : InputContext) : IO (Syntax × ModuleParserState × MessageLog) := do
+def parseHeader (inputCtx : InputContext) : IO (TSyntax ``Module.header × ModuleParserState × MessageLog) := do
  let dummyEnv ← mkEmptyEnvironment
  let p   := andthenFn whitespace Module.header.fn
  let tokens := Module.updateTokens (getTokenTable dummyEnv)
@ -73,7 +78,7 @@ def parseHeader (inputCtx : InputContext) : IO (Syntax × ModuleParserState × M
  let mut messages : MessageLog := {}
  for (pos, stk, err) in s.allErrors do
    messages := messages.add <| mkErrorMessage inputCtx pos stk err
-  pure (stx, {pos := s.pos, recovering := s.hasError}, messages)
+  pure (⟨stx⟩, {pos := s.pos, recovering := s.hasError}, messages)

 private def mkEOI (pos : String.Pos) : Syntax :=
  let atom := mkAtom (SourceInfo.original "".toSubstring pos "".toSubstring pos) ""
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@ -368,7 +368,7 @@ def setupImports
    (meta        : DocumentMeta)
    (cmdlineOpts : Options)
    (chanOut     : Std.Channel JsonRpc.Message)
-    (stx         : Syntax)
+    (stx         : TSyntax ``Parser.Module.header)
    : Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot SetupImportsResult) := do
  let importsAlreadyLoaded ← importsLoadedRef.modifyGet ((·, true))
  if importsAlreadyLoaded then
--- a/src/lake/Lake/Build/Module.lean
+++ b/src/lake/Lake/Build/Module.lean
@ -35,8 +35,8 @@ building an `Array` product of its direct local imports.
 -/
 def Module.recParseImports (mod : Module) : FetchM (Job (Array Module)) := Job.async do
  let contents ← IO.FS.readFile mod.leanFile
-  let imports ← Lean.parseImports' contents mod.leanFile.toString
-  let mods ← imports.foldlM (init := OrdModuleSet.empty) fun set imp =>
+  let res ← Lean.parseImports' contents mod.leanFile.toString
+  let mods ← res.imports.foldlM (init := OrdModuleSet.empty) fun set imp =>
    findModule? imp.module <&> fun | some mod => set.insert mod | none => set
  return mods.toArray

--- a/src/lake/Lake/CLI/Main.lean
+++ b/src/lake/Lake/CLI/Main.lean
@ -542,8 +542,8 @@ private def evalLeanFile
      let deps ← ws.runBuild (withRegisterJob s!"setup ({mod.name})" do mod.deps.fetch) buildConfig
      return mkSpawnArgs ws path deps (some mod.rootDir) mod.leanArgs
    else
-      let imports ← Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
-      let imports := imports.filterMap (ws.findModule? ·.module)
+      let res ← Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
+      let imports := res.imports.filterMap (ws.findModule? ·.module)
      let deps ← ws.runBuild (buildImportsAndDeps leanFile imports) buildConfig
      return mkSpawnArgs ws path deps none ws.root.moreLeanArgs
  logVerbose (mkCmdLog spawnArgs)
--- a/src/lake/Lake/Load/Lean/Elab.lean
+++ b/src/lake/Lake/Load/Lean/Elab.lean
@ -35,13 +35,13 @@ def importModulesUsingCache (imports : Array Import) (opts : Options) (trustLeve
  return env

 /-- Like `Lean.Elab.processHeader`, but using `importEnvCache`. -/
-def processHeader (header : Syntax) (opts : Options)
+def processHeader (header : TSyntax ``Parser.Module.header) (opts : Options)
 (inputCtx : Parser.InputContext) : StateT MessageLog IO Environment := do
  try
    let imports := Elab.headerToImports header
    importModulesUsingCache imports opts 1024
  catch e =>
-    let pos := inputCtx.fileMap.toPosition <| header.getPos?.getD 0
+    let pos := inputCtx.fileMap.toPosition <| header.raw.getPos?.getD 0
    modify (·.add { fileName := inputCtx.fileName, data := toString e, pos })
    mkEmptyEnvironment

--- a/tests/lean/1240.lean
+++ b/tests/lean/1240.lean
@ -2,4 +2,4 @@ import Lean.Elab.Frontend

 open Lean Elab in
 #eval show IO _ from do
-  discard <| importModules #[Import.mk Name.anonymous false] {} 0
+  discard <| importModules #[Name.anonymous] {} 0
--- a/tests/lean/ctor_layout.lean.expected.out
+++ b/tests/lean/ctor_layout.lean.expected.out
@ -5,6 +5,7 @@ obj@2
 ---
 scalar#4@0:u32
 obj@0
+scalar#1@4:u8
 obj@1
 obj@2
 obj@3