fix: lake: use server header for workspace modules (#9559)

This PR changes `lake setup-file` to use the server-provided header for
workspace modules.

This also reverts #9163 as the underlying issue is now fixed.

src/Lean/Server/FileWorker.lean

@@ -428,8 +428,8 @@ def setupImports
 
   return .ok {
     mainModuleName := doc.mod
-    isModule := header.isModule
-    imports := header.imports
+    isModule := strictOr setup.isModule header.isModule
+    imports := setup.imports?.getD header.imports
     opts
     importArts := setup.importArts
     plugins := setup.plugins

src/lake/Lake/Build/Module.lean

@@ -10,14 +10,16 @@ import Lean.Elab.ParseImportsFast
 import Lake.Build.Common
 import Lake.Build.Target
 
-/-! # Module Facet Builds
-Build function definitions for a module's builtin facets.
--/
+/-! # Module Build Definitions -/
 
 open System Lean
 
 namespace Lake
 
+/-! ## Facet Builds
+Build function definitions for a module's builtin facets.
+-/
+
 /-- Compute library directories and build external library Jobs of the given packages. -/
 @[deprecated "Deprecated without replacement" (since := "2025-03-28")]
 def recBuildExternDynlibs (pkgs : Array Package)
@@ -95,16 +97,21 @@ private structure ModuleImportData where
     | .ok impSet s => .ok impSet.toArray s
     | .error e s => .error e s
 
+private def computeTransImportsAux
+  (fileName : String) (imports : Array Module)
+: FetchM (Job (Array Module)) := do
+  collectImportsAux fileName imports fun imp =>
+    (true, ·) <$> imp.transImports.fetch
+
 /-- Recursively compute a module's transitive imports. -/
 def Module.recComputeTransImports (mod : Module) : FetchM (Job (Array Module)) := ensureJob do
-  collectImportsAux mod.leanFile.toString (← (← mod.imports.fetch).await) fun imp =>
-    (true, ·) <$> imp.transImports.fetch
+  inline <| computeTransImportsAux mod.relLeanFile.toString (← (← mod.imports.fetch).await)
 
 /-- The `ModuleFacetConfig` for the builtin `transImportsFacet`. -/
 def Module.transImportsFacetConfig : ModuleFacetConfig transImportsFacet :=
   mkFacetJobConfig recComputeTransImports (buildable := false)
 
-def computePrecompileImportsAux
+private def computePrecompileImportsAux
   (fileName : String) (imports : Array Module)
 : FetchM (Job (Array Module)) := do
   collectImportsAux fileName imports fun imp =>
@@ -115,7 +122,7 @@ def computePrecompileImportsAux
 
 /-- Recursively compute a module's precompiled imports. -/
 def Module.recComputePrecompileImports (mod : Module) : FetchM (Job (Array Module)) := ensureJob do
-  inline <| computePrecompileImportsAux mod.leanFile.toString (← (← mod.imports.fetch).await)
+  inline <| computePrecompileImportsAux mod.relLeanFile.toString (← (← mod.imports.fetch).await)
 
 /-- The `ModuleFacetConfig` for the builtin `precompileImportsFacet`. -/
 def Module.precompileImportsFacetConfig : ModuleFacetConfig precompileImportsFacet :=
@@ -881,10 +888,54 @@ def Module.initFacetConfigs : DNameMap ModuleFacetConfig :=
 @[inherit_doc Module.initFacetConfigs]
 abbrev initModuleFacetConfigs := Module.initFacetConfigs
 
-/-!
-Definitions to support `lake setup-file` builds.
+/-! ## Top-Level Builds
+Definitions to support `lake setup-file` and `lake lean` builds.
 -/
 
+/--
+Computes the module setup of a workspace module being edited by the Lean language server,
+building its imports and other dependencies. Used by `lake setup-file`.
+
+Due to its exclusive use as a top-level build, it does not construct a proper trace state.
+-/
+private def setupEditedModule
+  (mod : Module) (header : ModuleHeader)
+: FetchM (Job ModuleSetup) := do
+  let extraDepJob ← mod.lib.extraDep.fetch
+  let directImports ← header.imports.mapM fun imp => do
+    return ⟨imp, ← findModule? imp.module⟩
+  let fileName := mod.relLeanFile.toString
+  let localImports := directImports.filterMap (·.module?)
+  let impInfoJob ← fetchImportInfo fileName mod.name header
+  let precompileImports ←
+    if mod.shouldPrecompile then
+      (← computeTransImportsAux fileName localImports).await
+    else
+      (← computePrecompileImportsAux fileName localImports).await
+  let impLibsJob ← Job.collectArray (traceCaption := "import dynlibs") <$>
+    mod.fetchImportLibs precompileImports mod.shouldPrecompile
+  let externLibsJob ← Job.collectArray (traceCaption := "package external libraries") <$>
+    if mod.shouldPrecompile then mod.pkg.externLibs.mapM (·.dynlib.fetch) else pure #[]
+  let dynlibsJob ← mod.dynlibs.fetchIn mod.pkg "module dynlibs"
+  let pluginsJob ← mod.plugins.fetchIn mod.pkg "module plugins"
+  extraDepJob.bindM (sync := true) fun _ => do
+  impInfoJob.bindM (sync := true) fun info => do
+  impLibsJob.bindM (sync := true) fun impLibs => do
+  externLibsJob.bindM (sync := true) fun externLibs => do
+  dynlibsJob.bindM (sync := true) fun dynlibs => do
+  pluginsJob.mapM fun plugins => do
+    let {dynlibs, plugins} ← computeModuleDeps impLibs externLibs dynlibs plugins
+    let transImpArts ← fetchTransImportArts directImports info.directArts !header.isModule
+    return {
+      name := mod.name
+      isModule := header.isModule
+      imports? := none
+      importArts := transImpArts
+      dynlibs := dynlibs.map (·.path)
+      plugins := plugins.map (·.path)
+      options := mod.leanOptions
+    }
+
 /--
 Computes the module setup of Lean code external to the workspace,
 building its imports and other dependencies.
@@ -894,10 +945,9 @@ to build the dependencies of the file and generate the data for `lean --setup`.
 
 Due to its exclusive use as a top-level build, it does not construct a proper trace state.
 -/
-def setupExternalModule
+private def setupExternalModule
   (fileName : String) (header : ModuleHeader) (leanOpts : LeanOptions)
 : FetchM (Job ModuleSetup) := do
-  withRegisterJob s!"setup ({fileName})" do
   let root ← getRootPackage
   let extraDepJob ← root.extraDep.fetch
   let imports ← header.imports.mapM fun imp => do
@@ -927,3 +977,62 @@ def setupExternalModule
       plugins := plugins.map (·.path)
       options := leanOpts
     }
+
+/--
+Computes the module setup of edited Lean code for the Lean language server,
+building its imports and other dependencies. Used by `lake setup-file`.
+
+Due to its exclusive use as a top-level build, it does not construct a proper trace state.
+-/
+def setupServerModule
+  (fileName : String) (path : FilePath) (header? : Option ModuleHeader)
+: FetchM (Job ModuleSetup) :=
+  withRegisterJob s!"setup-file {fileName}" do
+  let header ← header?.getDM do
+    Lean.parseImports' (← IO.FS.readFile path) fileName
+  if let some mod ← findModuleBySrc? path then
+    logVerbose s!"file identified as module: {mod.name}"
+    setupEditedModule mod header
+  else
+    setupExternalModule fileName header (← getServerOptions)
+
+/--
+Computes the arguments required to evaluate the Lean file with `lean`,
+building its imports and other dependencies. Used by `lake lean`.
+
+Due to its exclusive use as a top-level build, it does not construct a proper trace state.
+-/
+def prepareLeanCommand
+  (leanFile : FilePath) (moreArgs : Array String := #[])
+: FetchM (Job IO.Process.SpawnArgs) :=
+  withRegisterJob s!"prepare lean {leanFile}" do
+  let some path ← resolvePath? leanFile
+    | error s!"file not found: {leanFile}"
+  if let some mod ← findModuleBySrc? path then
+    logVerbose s!"file identified as module: {mod.name}"
+    let setupJob ← mod.setup.fetch
+    setupJob.mapM (sync := true) fun setup => do
+      mkSpawnArgs path setup mod.leanArgs
+  else
+    let header ← Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
+    let setupJob ← setupExternalModule leanFile.toString header (← getLeanOptions)
+    setupJob.mapM (sync := true) fun setup => do
+      mkSpawnArgs path setup (← getLeanArgs)
+where
+  mkArgs leanFile setup cfgArgs := do
+    let args := cfgArgs.push leanFile.toString
+    let (h, setupFile) ← IO.FS.createTempFile
+    let contents := (toJson setup).compress
+    logVerbose s!"module setup: {contents}"
+    h.putStr contents
+    let args := args ++ #["--setup", setupFile.toString]
+    return args
+  mkSpawnArgs leanFile setup cfgArgs := do
+    let args ← mkArgs leanFile setup cfgArgs
+    let spawnArgs : IO.Process.SpawnArgs := {
+      args := args ++ moreArgs
+      cmd := (← getLean).toString
+      env := (← getAugmentedEnv)
+    }
+    logVerbose (mkCmdLog spawnArgs)
+    return spawnArgs

src/lake/Lake/CLI/Actions.lean

@@ -6,7 +6,7 @@ Authors: Mac Malone
 prelude
 import Lake.Build.Run
 import Lake.Build.Targets
-import Lake.Build.Common
+import Lake.Build.Module
 import Lake.CLI.Build
 
 namespace Lake
@@ -89,3 +89,18 @@ def Package.lint (pkg : Package) (args : List String := []) (buildConfig : Build
   else
     let pkgName := pkg.name.toString (escape := false)
     error s!"{pkgName}: invalid lint driver: unknown script or executable '{driver}'"
+
+/--
+Run `lean` on file using configuration from the workspace.
+
+Additional arguments can be passed to `lean` via `moreArgs` and the
+building of dependencies can be further configured via `buildConfig`.
+-/
+def Workspace.evalLeanFile
+  (ws : Workspace) (leanFile : FilePath)
+  (moreArgs : Array String := #[]) (buildConfig : BuildConfig := {})
+: IO UInt32 := do
+  let spawnArgs ← ws.runBuild (cfg := buildConfig) do
+    prepareLeanCommand leanFile moreArgs
+  let child ← IO.Process.spawn spawnArgs
+  child.wait

src/lake/Lake/CLI/Main.lean

@@ -537,47 +537,13 @@ protected def exe : CliM PUnit := do
   let exeFile ← ws.runBuild exe.fetch (mkBuildConfig opts)
   exit <| ← (env exeFile.toString args.toArray).run <| mkLakeContext ws
 
-private def evalLeanFile
-  (ws : Workspace) (leanFile : FilePath)
-  (moreArgs : Array String := #[]) (buildConfig : BuildConfig := {})
-: LoggerIO UInt32 := do
-  let some path ← resolvePath? leanFile
-    | error s!"file not found: {leanFile}"
-  let args ← do
-    if let some mod := ws.findModuleBySrc? path then
-      let setup ← ws.runBuild (cfg := buildConfig) do
-        withRegisterJob s!"{mod.name}:setup" do mod.setup.fetch
-      mkArgs path setup mod.leanArgs
-    else
-      let header ← Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
-      let setup ← ws.runBuild (cfg := buildConfig) do
-        setupExternalModule leanFile.toString header ws.leanOptions
-      mkArgs path setup ws.root.moreLeanArgs
-  let spawnArgs : IO.Process.SpawnArgs := {
-    args := args ++ moreArgs
-    cmd := ws.lakeEnv.lean.lean.toString
-    env := ws.augmentedEnvVars
-  }
-  logVerbose (mkCmdLog spawnArgs)
-  let child ← IO.Process.spawn spawnArgs
-  child.wait
-where
-  mkArgs leanFile setup cfgArgs := do
-    let args := cfgArgs.push leanFile.toString
-    let (h, setupFile) ← IO.FS.createTempFile
-    let contents := (toJson setup).compress
-    logVerbose s!"module setup: {contents}"
-    h.putStr contents
-    let args := args ++ #["--setup", setupFile.toString]
-    return args
-
 protected def lean : CliM PUnit := do
   processOptions lakeOption
   let leanFile ← takeArg "Lean file"
   let opts ← getThe LakeOptions
   noArgsRem do
   let ws ← loadWorkspace (← mkLoadConfig opts)
-  let rc ← evalLeanFile ws leanFile opts.subArgs.toArray (mkBuildConfig opts)
+  let rc ← ws.evalLeanFile leanFile opts.subArgs.toArray (mkBuildConfig opts)
   exit rc
 
 protected def translateConfig : CliM PUnit := do

src/lake/Lake/CLI/Serve.lean

@@ -25,14 +25,13 @@ def invalidConfigEnvVar := "LAKE_INVALID_CONFIG"
 /--
 Build the dependencies of a Lean file and print the computed module's setup as JSON.
 If `header?` is not not `none`, it will be used to determine imports instead of the
-file's own header (for modules external to the workspace).
+file's own header.
 
 Requires a configuration file to succeed. If no configuration file exists, it
 will exit silently with `noConfigFileCode` (i.e, 2).
 
-The `setup-file` command is used internally by the Lean 4 server.
+The `setup-file` command is used internally by the Lean server.
 -/
--- TODO: Use `header?` for modules within the workspace as well.
 def setupFile
   (loadConfig : LoadConfig) (leanFile : FilePath)
   (header? : Option ModuleHeader := none) (buildConfig : BuildConfig := {})
@@ -54,16 +53,9 @@ def setupFile
   else
     let some ws ← loadWorkspace loadConfig |>.toBaseIO buildConfig.outLv buildConfig.ansiMode
       | error "failed to load workspace"
-    if let some mod := ws.findModuleBySrc? path then
-      let setup ← ws.runBuild (cfg := buildConfig) do
-        withRegisterJob s!"{mod.name}:setup" do mod.setup.fetch
-      IO.println (toJson setup).compress
-    else
-      let header ← header?.getDM do
-        Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
-      let setup ← ws.runBuild (cfg := buildConfig) do
-        setupExternalModule leanFile.toString header ws.serverOptions
-      IO.println (toJson setup).compress
+    let setup ← ws.runBuild (cfg := buildConfig) do
+      setupServerModule leanFile.toString path header?
+    IO.println (toJson setup).compress
 
 /--
 Start the Lean LSP for the `Workspace` loaded from `config`

src/lake/Lake/Config/Monad.lean

@@ -8,7 +8,7 @@ import Lake.Config.Context
 import Lake.Config.Workspace
 
 open System
-open Lean (Name NameMap)
+open Lean (Name NameMap LeanOptions)
 
 /-! # Lake Configuration Monads
 Definitions and helpers for interacting with the Lake configuration monads.
@@ -82,6 +82,10 @@ def findPackage? (name : Name) : m (Option (NPackage name)) :=
 def findModule? (name : Name) : m (Option Module) :=
   (·.findModule? name) <$> getWorkspace
 
+@[inherit_doc Workspace.findModuleBySrc?, inline]
+def findModuleBySrc? (path : FilePath) : m (Option Module) :=
+  (·.findModuleBySrc? path) <$> getWorkspace
+
 @[inherit_doc Workspace.findLeanExe?, inline]
 def findLeanExe? (name : Name) : m (Option LeanExe) :=
   (·.findLeanExe? name) <$> getWorkspace
@@ -94,6 +98,18 @@ def findLeanLib? (name : Name) : m (Option LeanLib) :=
 def findExternLib? (name : Name) : m (Option ExternLib) :=
   (·.findExternLib? name) <$> getWorkspace
 
+@[inherit_doc Workspace.serverOptions, inline]
+def getServerOptions : m LeanOptions :=
+  (·.serverOptions) <$> getWorkspace
+
+@[inherit_doc Workspace.leanOptions, inline]
+def getLeanOptions : m LeanOptions :=
+  (·.leanOptions) <$> getWorkspace
+
+@[inherit_doc Workspace.leanArgs, inline]
+def getLeanArgs : m (Array String) :=
+  (·.leanArgs) <$> getWorkspace
+
 /-- Get the paths added to `LEAN_PATH` by the context's workspace. -/
 @[inline] def getLeanPath : m SearchPath :=
   (·.leanPath) <$> getWorkspace

src/lake/Lake/Config/Package.lean

@@ -566,7 +566,7 @@ namespace Package
 @[inline] def leanOptions (self : Package) : LeanOptions :=
   .ofArray self.config.leanOptions
 
-/-- The package's `moreLeanArgs` configuration appended to its `leanOptions` configuration. -/
+/-- The package's `moreLeanArgs` configuration. -/
 @[inline] def moreLeanArgs (self : Package) : Array String :=
   self.config.moreLeanArgs
 

src/lake/Lake/Config/Workspace.lean

@@ -69,6 +69,10 @@ namespace Workspace
 @[inline] def pkgsDir (self : Workspace) : FilePath :=
   self.root.pkgsDir
 
+/-- Arguments to pass to `lean` for files outside a library (e.g., via `lake lean`). -/
+@[inline] def leanArgs (self : Workspace) : Array String :=
+  self.root.moreLeanArgs
+
 /-- Options to pass to `lean` for files outside a library (e.g., via `lake lean`). -/
 @[inline] def leanOptions (self : Workspace) : LeanOptions :=
   self.root.leanOptions

src/lake/tests/setupFile/test.sh

@@ -46,13 +46,11 @@ BOGUS_JSON='{"isModule":false,"imports":[{"module":"Test.Bogus","isMeta":false,"
 test_err 'no such file or directory' setup-file ImportFoo.lean "$BOGUS_JSON"
 
 # Test that when a header is provided (via CLI or stdin),
-# the header is *NOT* used for an internal module and its imports are not built.
-# TODO: Use the provided header.
+# the header is used for an internal module and its imports are built.
 test_out '"isModule":false' setup-file Test.lean
-test_out '"isModule":false' setup-file Test.lean "$HEADER_JSON"
-echo "$HEADER_JSON" | test_out '"isModule":false' setup-file Test.lean -
-# If the provided import (Test.Bogus) was built, this would fail.
-test_run setup-file Test.lean "$BOGUS_JSON"
+test_out '"isModule":true' setup-file Test.lean "$HEADER_JSON"
+echo "$HEADER_JSON" | test_out '"isModule":true' setup-file Test.lean -
+test_err 'no such file or directory' setup-file Test.lean "$BOGUS_JSON"
 
 # Cleanup
 rm -f produced.out