From dd64678f0799c09e336a4e55a7f3310ce276174b Mon Sep 17 00:00:00 2001
From: Mac Malone <tydeu@hatpress.net>
Date: Mon, 23 Jun 2025 14:00:14 -0400
Subject: [PATCH] feat: server support for new module setup (#8699)

This PR adds support to the server for the new module setup process by
changing how `lake setup-file` is used.

In the new server setup, `lake setup-file` is invoked with the file name
of the edited module passed as a CLI argument and with the parsed header
passed to standard input in JSON form. Standard input is used to avoid
potentially exceeding the CLI length limits on Windows. Lake will build
the module's imports along with any other dependencies and then return
the module's workspace configuration via JSON (now in the form of
`ModuleSetup`). The server then post-processes this configuration a bit
and returns it back to the Lean language processor.

The server's header is currently only fully respected by Lake for
external modules (files that are not part of any workspace library). For
workspace modules, the saved module header is currently used to build
imports (as has been done since #7909). A follow-up Lake PR will align
both cases to follow the server's header.

Lean search paths (e.g., `LEAN_PATH`, `LEAN_SRC_PATH`) are no longer
negotiated between the server and Lake. These environment variables are
already configured during sever setup by `lake serve` and do not change
on a per-file basis. Lake can also pre-resolve the `.olean` files of
imports via the `importArts` field of `ModuleSetup`, limiting the
potential utility of communicating `LEAN_PATH`.
---
 src/Lean/Elab/Frontend.lean                   |  6 +-
 src/Lean/Elab/Import.lean                     |  7 ++-
 src/Lean/Environment.lean                     | 23 ++-----
 src/Lean/Language/Lean.lean                   |  6 +-
 .../Server/Completion/ImportCompletion.lean   |  1 -
 src/Lean/Server/FileWorker.lean               | 18 +++---
 src/Lean/Server/FileWorker/SetupFile.lean     | 59 ++++++++----------
 src/Lean/Server/Watchdog.lean                 |  2 -
 src/Lean/Setup.lean                           | 51 +++++++++++++---
 src/Lean/Util.lean                            |  1 -
 src/Lean/Util/FileSetupInfo.lean              | 15 -----
 src/Lean/Util/LeanOptions.lean                |  2 +-
 src/Lean/Util/Paths.lean                      | 24 --------
 src/lake/Lake/Build/Module.lean               |  4 +-
 src/lake/Lake/CLI/Main.lean                   | 18 +++---
 src/lake/Lake/CLI/Serve.lean                  | 60 +++++++------------
 src/lake/Lake/Config/Workspace.lean           |  1 -
 src/lake/examples/deps/test.sh                |  2 +-
 src/lake/tests/setupFile/test.sh              | 28 +++++++--
 tests/pkg/setup/setup.json                    |  6 +-
 20 files changed, 157 insertions(+), 177 deletions(-)
 delete mode 100644 src/Lean/Util/FileSetupInfo.lean
 delete mode 100644 src/Lean/Util/Paths.lean

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 1f5342ca6f..f7a54c767c 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -158,10 +158,10 @@ def runFrontend
       return .ok {
         trustLevel
         mainModuleName := setup.name
-        isModule := setup.isModule
-        imports := setup.imports
+        isModule := strictOr setup.isModule stx.isModule
+        imports := setup.imports?.getD stx.imports
         plugins := plugins ++ setup.plugins
-        modules := setup.modules
+        importArts := setup.importArts
         -- override cmdline options with setup options
         opts := opts.mergeBy (fun _ _ hOpt => hOpt) setup.options.toOptions
       }
diff --git a/src/Lean/Elab/Import.lean b/src/Lean/Elab/Import.lean
index 570b3d2e76..756b1f1c3a 100644
--- a/src/Lean/Elab/Import.lean
+++ b/src/Lean/Elab/Import.lean
@@ -5,7 +5,6 @@ Authors: Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
 import Lean.Parser.Module
-import Lean.Util.Paths
 import Lean.CoreM
 
 namespace Lean.Elab
@@ -29,13 +28,17 @@ def HeaderSyntax.imports (stx : HeaderSyntax) (includeInit : Bool := true) : Arr
       | _ => unreachable!
   | _ => unreachable!
 
+def HeaderSyntax.toModuleHeader (stx : HeaderSyntax) : ModuleHeader where
+  isModule := stx.isModule
+  imports := stx.imports
+
 abbrev headerToImports := @HeaderSyntax.imports
 
 def processHeaderCore
     (startPos : String.Pos) (imports : Array Import) (isModule : Bool)
     (opts : Options) (messages : MessageLog) (inputCtx : Parser.InputContext)
     (trustLevel : UInt32 := 0) (plugins : Array System.FilePath := #[]) (leakEnv := false)
-    (mainModule := Name.anonymous) (arts : NameMap ModuleArtifacts := {})
+    (mainModule := Name.anonymous) (arts : NameMap ImportArtifacts := {})
     : IO (Environment × MessageLog) := do
   let level := if isModule then
     if Elab.inServer.get opts then
diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean
index 21ae429546..5c456744b4 100644
--- a/src/Lean/Environment.lean
+++ b/src/Lean/Environment.lean
@@ -1876,18 +1876,6 @@ abbrev ImportStateM := StateRefT ImportState IO
 @[inline] nonrec def ImportStateM.run (x : ImportStateM α) (s : ImportState := {}) : IO (α × ImportState) :=
   x.run s
 
-def ModuleArtifacts.oleanParts (arts : ModuleArtifacts) : Array System.FilePath := Id.run do
-  let mut fnames := #[]
-  -- Opportunistically load all available parts.
-  -- Producer (e.g., Lake) should limit parts to the proper import level.
-  if let some mFile := arts.olean? then
-    fnames := fnames.push mFile
-    if let some sFile := arts.oleanServer? then
-      fnames := fnames.push sFile
-      if let some pFile := arts.oleanPrivate? then
-        fnames := fnames.push pFile
-  return fnames
-
 private def findOLeanParts (mod : Name) : IO (Array System.FilePath) := do
   let mFile ← findOLean mod
   unless (← mFile.pathExists) do
@@ -1904,7 +1892,7 @@ private def findOLeanParts (mod : Name) : IO (Array System.FilePath) := do
   return fnames
 
 partial def importModulesCore
-    (imports : Array Import) (isModule := false) (arts : NameMap ModuleArtifacts := {}) :
+    (imports : Array Import) (isModule := false) (arts : NameMap ImportArtifacts := {}) :
     ImportStateM Unit := do
   go imports (importAll := true) (isExported := isModule) (isMeta := false)
   if isModule then
@@ -1977,10 +1965,9 @@ where go (imports : Array Import) (importAll isExported isMeta : Bool) := do
       continue
     let fnames ←
       if let some arts := arts.find? i.module then
-        let fnames := arts.oleanParts
-        if fnames.isEmpty then
-          findOLeanParts i.module
-        else pure fnames
+        -- Opportunistically load all available parts.
+        -- Producer (e.g., Lake) should limit parts to the proper import level.
+        pure arts.oleanParts
       else
         findOLeanParts i.module
     let parts ← readModuleDataParts fnames
@@ -2146,7 +2133,7 @@ as if no `module` annotations were present in the imports.
 -/
 def importModules (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0)
     (plugins : Array System.FilePath := #[]) (leakEnv := false) (loadExts := false)
-    (level := OLeanLevel.private) (arts : NameMap ModuleArtifacts := {})
+    (level := OLeanLevel.private) (arts : NameMap ImportArtifacts := {})
     : IO Environment := profileitIO "import" opts do
   for imp in imports do
     if imp.module matches .anonymous then
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 88d6b8ebff..6b039e4861 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -291,8 +291,8 @@ structure SetupImportsResult where
   opts : Options
   /-- Kernel trust level. -/
   trustLevel : UInt32 := 0
-  /-- Pre-resolved artifacts of related modules (e.g., this module's transitive imports). -/
-  modules : NameMap ModuleArtifacts := {}
+  /-- Pre-resolved artifacts of transitively imported modules. -/
+  importArts : NameMap ImportArtifacts := {}
   /-- Lean plugins to load as part of the environment setup. -/
   plugins : Array System.FilePath := #[]
 
@@ -491,7 +491,7 @@ where
       -- allows `headerEnv` to be leaked, which would live until the end of the process anyway
       let (headerEnv, msgLog) ← Elab.processHeaderCore (leakEnv := true)
         stx.startPos setup.imports setup.isModule setup.opts .empty ctx.toInputContext
-        setup.trustLevel setup.plugins setup.mainModuleName setup.modules
+        setup.trustLevel setup.plugins setup.mainModuleName setup.importArts
       let stopTime := (← IO.monoNanosNow).toFloat / 1000000000
       let diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
       if msgLog.hasErrors then
diff --git a/src/Lean/Server/Completion/ImportCompletion.lean b/src/Lean/Server/Completion/ImportCompletion.lean
index d34af24c73..3984e33cce 100644
--- a/src/Lean/Server/Completion/ImportCompletion.lean
+++ b/src/Lean/Server/Completion/ImportCompletion.lean
@@ -5,7 +5,6 @@ Authors: Marc Huisinga
 -/
 prelude
 import Lean.Data.NameTrie
-import Lean.Util.Paths
 import Lean.Util.LakePath
 import Lean.Server.Completion.CompletionItemData
 import Lean.Parser.Module
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 9261be3aa2..31f4c92b13 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -14,7 +14,6 @@ import Lean.Environment
 import Lean.Data.Lsp
 import Lean.Data.Json.FromToJson
 
-import Lean.Util.FileSetupInfo
 import Lean.LoadDynlib
 import Lean.Language.Lean
 
@@ -372,7 +371,7 @@ def setupImports
     (doc         : DocumentMeta)
     (cmdlineOpts : Options)
     (chanOut     : Std.Channel JsonRpc.Message)
-    (stx         : TSyntax ``Parser.Module.header)
+    (stx         : Elab.HeaderSyntax)
     : Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot SetupImportsResult) := do
   let importsAlreadyLoaded ← importsLoadedRef.modifyGet ((·, true))
   if importsAlreadyLoaded then
@@ -384,9 +383,9 @@ def setupImports
     -- should not be visible to user as task is already canceled
     return .error { diagnostics := .empty, result? := none, metaSnap := default }
 
+  let header := stx.toModuleHeader
   chanOut.sync.send <| mkInitialIleanInfoUpdateNotification doc <| collectImports stx
-  let imports := Elab.headerToImports stx
-  let fileSetupResult ← setupFile doc imports fun stderrLine => do
+  let fileSetupResult ← setupFile doc header fun stderrLine => do
     let progressDiagnostic := {
       range      := ⟨⟨0, 0⟩, ⟨1, 0⟩⟩
       -- make progress visible anywhere in the file
@@ -414,8 +413,10 @@ def setupImports
     }
   | _ => pure ()
 
+  let setup := fileSetupResult.setup
+
   -- override cmdline options with file options
-  let opts := cmdlineOpts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
+  let opts := cmdlineOpts.mergeBy (fun _ _ fileOpt => fileOpt) setup.options.toOptions
 
   -- default to async elaboration; see also `Elab.async` docs
   let opts := Elab.async.setIfNotSet opts true
@@ -424,10 +425,11 @@ def setupImports
 
   return .ok {
     mainModuleName := doc.mod
-    isModule := Elab.HeaderSyntax.isModule stx
-    imports
+    isModule := strictOr setup.isModule header.isModule
+    imports := setup.imports?.getD header.imports
     opts
-    plugins := fileSetupResult.plugins
+    importArts := setup.importArts
+    plugins := setup.plugins
   }
 
 /- Worker initialization sequence. -/
diff --git a/src/Lean/Server/FileWorker/SetupFile.lean b/src/Lean/Server/FileWorker/SetupFile.lean
index a00e00885c..606c02d5fd 100644
--- a/src/Lean/Server/FileWorker/SetupFile.lean
+++ b/src/Lean/Server/FileWorker/SetupFile.lean
@@ -6,7 +6,6 @@ Authors: Sebastian Ullrich, Marc Huisinga
 prelude
 import Init.System.IO
 import Lean.Server.Utils
-import Lean.Util.FileSetupInfo
 import Lean.Util.LakePath
 import Lean.LoadDynlib
 import Lean.Server.ServerTask
@@ -24,20 +23,22 @@ structure LakeSetupFileOutput where
 partial def runLakeSetupFile
     (m                 : DocumentMeta)
     (lakePath filePath : System.FilePath)
-    (imports           : Array Import)
+    (header            : ModuleHeader)
     (handleStderr      : String → IO Unit)
     : IO LakeSetupFileOutput := do
-  let mut args := #["setup-file", filePath.toString] ++ imports.map (toString ·.module)
+  let mut args := #["setup-file", filePath.toString, "-"]
   if m.dependencyBuildMode matches .never then
     args := args.push "--no-build" |>.push "--no-cache"
   let spawnArgs : Process.SpawnArgs := {
-    stdin  := Process.Stdio.null
+    stdin  := Process.Stdio.piped
     stdout := Process.Stdio.piped
     stderr := Process.Stdio.piped
     cmd    := lakePath.toString
     args
   }
   let lakeProc ← Process.spawn spawnArgs
+  let (stdin, lakeProc) ← lakeProc.takeStdin
+  stdin.putStrLn (toJson header).compress
 
   let rec processStderr (acc : String) : IO String := do
     let line ← lakeProc.stderr.getLine
@@ -68,61 +69,53 @@ inductive FileSetupResultKind where
 structure FileSetupResult where
   /-- Kind of outcome. -/
   kind        : FileSetupResultKind
-  /-- Additional options from successful setup, or else empty. -/
-  fileOptions : Options
-  /-- Lean plugins from successful setup, or else empty. -/
-  plugins     : Array System.FilePath
+  /-- Configuration from a successful setup, or else the default. -/
+  setup       : ModuleSetup
 
-def FileSetupResult.ofSuccess (fileOptions : Options)
-    (plugins : Array System.FilePath) : IO FileSetupResult := do return {
+def FileSetupResult.ofSuccess (setup : ModuleSetup) : IO FileSetupResult := do return {
   kind          := FileSetupResultKind.success
-  fileOptions, plugins
+  setup
 }
 
-def FileSetupResult.ofNoLakefile : IO FileSetupResult := do return {
+def FileSetupResult.ofNoLakefile (m : DocumentMeta) (header : ModuleHeader) : IO FileSetupResult := do return {
   kind          := FileSetupResultKind.noLakefile
-  fileOptions   := Options.empty
-  plugins       := #[]
+  setup         := {name := m.mod, isModule := header.isModule}
 }
 
-def FileSetupResult.ofImportsOutOfDate : IO FileSetupResult := do return {
+def FileSetupResult.ofImportsOutOfDate (m : DocumentMeta) (header : ModuleHeader) : IO FileSetupResult := do return {
   kind          := FileSetupResultKind.importsOutOfDate
-  fileOptions   := Options.empty
-  plugins       := #[]
+  setup         := {name := m.mod, isModule := header.isModule}
 }
 
-def FileSetupResult.ofError (msg : String) : IO FileSetupResult := do return {
+def FileSetupResult.ofError (m : DocumentMeta) (header : ModuleHeader) (msg : String) : IO FileSetupResult := do return {
   kind          := FileSetupResultKind.error msg
-  fileOptions   := Options.empty
-  plugins       := #[]
+  setup         := {name := m.mod, isModule := header.isModule}
 }
 
 /-- Uses `lake setup-file` to compile dependencies on the fly and add them to `LEAN_PATH`.
 Compilation progress is reported to `handleStderr`. Returns the search path for
 source files and the options for the file. -/
-partial def setupFile (m : DocumentMeta) (imports : Array Import) (handleStderr : String → IO Unit) : IO FileSetupResult := do
+partial def setupFile (m : DocumentMeta) (header : ModuleHeader) (handleStderr : String → IO Unit) : IO FileSetupResult := do
   let some filePath := System.Uri.fileUriToPath? m.uri
-    | return ← FileSetupResult.ofNoLakefile -- untitled files have no lakefile
+    | return ← FileSetupResult.ofNoLakefile m header -- untitled files have no lakefile
 
   let lakePath ← determineLakePath
   if !(← System.FilePath.pathExists lakePath) then
-    return ← FileSetupResult.ofNoLakefile
+    return ← FileSetupResult.ofNoLakefile m header
 
-  let result ← runLakeSetupFile m lakePath filePath imports handleStderr
+  let result ← runLakeSetupFile m lakePath filePath header handleStderr
 
   let cmdStr := " ".intercalate (toString result.spawnArgs.cmd :: result.spawnArgs.args.toList)
 
   match result.exitCode with
   | 0 =>
-    let Except.ok (info : FileSetupInfo) := Json.parse result.stdout >>= fromJson?
-      | return ← FileSetupResult.ofError s!"Invalid output from `{cmdStr}`:\n{result.stdout}\nstderr:\n{result.stderr}"
-    initSearchPath (← getBuildDir) info.paths.oleanPath
-    info.paths.loadDynlibPaths.forM loadDynlib
-    let pluginPaths ← info.paths.pluginPaths.mapM realPathNormalized
-    FileSetupResult.ofSuccess info.setupOptions.toOptions pluginPaths
+    let Except.ok (setup : ModuleSetup) := Json.parse result.stdout >>= fromJson?
+      | return ← FileSetupResult.ofError m header s!"Invalid output from `{cmdStr}`:\n{result.stdout}\nstderr:\n{result.stderr}"
+    setup.dynlibs.forM loadDynlib
+    FileSetupResult.ofSuccess setup
   | 2 => -- exit code for lake reporting that there is no lakefile
-    FileSetupResult.ofNoLakefile
+    FileSetupResult.ofNoLakefile m header
   | 3 => -- exit code for `--no-build`
-    FileSetupResult.ofImportsOutOfDate
+    FileSetupResult.ofImportsOutOfDate m header
   | _ =>
-    FileSetupResult.ofError s!"`{cmdStr}` failed:\n{result.stdout}\nstderr:\n{result.stderr}"
+    FileSetupResult.ofError m header s!"`{cmdStr}` failed:\n{result.stdout}\nstderr:\n{result.stderr}"
diff --git a/src/Lean/Server/Watchdog.lean b/src/Lean/Server/Watchdog.lean
index 9ef954f877..f089c12822 100644
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@@ -11,8 +11,6 @@ import Std.Data.TreeMap
 import Init.Data.ByteArray
 import Lean.Data.RBMap
 
-import Lean.Util.Paths
-
 import Lean.Data.FuzzyMatching
 import Lean.Data.Json
 import Lean.Data.Lsp
diff --git a/src/Lean/Setup.lean b/src/Lean/Setup.lean
index 51962088e4..c84684a13e 100644
--- a/src/Lean/Setup.lean
+++ b/src/Lean/Setup.lean
@@ -38,6 +38,26 @@ structure ModuleHeader where
   isModule : Bool
   deriving Repr, Inhabited, ToJson, FromJson
 
+/--
+Module data files used for an `import` statement.
+This structure is designed for efficient JSON serialization.
+-/
+structure ImportArtifacts where
+  oleanParts : Array System.FilePath
+  deriving Repr, Inhabited
+
+instance : ToJson ImportArtifacts := ⟨(toJson ·.oleanParts)⟩
+instance : FromJson ImportArtifacts := ⟨(.mk <$> fromJson? ·)⟩
+
+def ImportArtifacts.olean? (arts : ImportArtifacts) :=
+  arts.oleanParts[0]?
+
+def ImportArtifacts.oleanServer? (arts : ImportArtifacts) :=
+  arts.oleanParts[1]?
+
+def ImportArtifacts.oleanPrivate? (arts : ImportArtifacts) :=
+  arts.oleanParts[2]?
+
 /-- Files containing data for a single module. -/
 structure ModuleArtifacts where
   lean? : Option System.FilePath := none
@@ -49,19 +69,32 @@ structure ModuleArtifacts where
   bc? : Option System.FilePath := none
   deriving Repr, Inhabited, ToJson, FromJson
 
-/--
-A module's setup information as described by a JSON file.
-Supersedes the module's header when the `--setup` CLI option is used.
--/
+def ModuleArtifacts.oleanParts (arts : ModuleArtifacts) : Array System.FilePath := Id.run do
+  let mut fnames := #[]
+  if let some mFile := arts.olean? then
+    fnames := fnames.push mFile
+    if let some sFile := arts.oleanServer? then
+      fnames := fnames.push sFile
+      if let some pFile := arts.oleanPrivate? then
+        fnames := fnames.push pFile
+  return fnames
+
+/-- A module's setup information as described by a JSON file. -/
 structure ModuleSetup where
   /-- Name of the module. -/
   name : Name
-  /-- Whether the module is participating in the module system. -/
+  /--
+  Whether the module, by default, participates in the module system.
+  Even if `false`, a module can still choose to participate by using `module` in its header.
+  -/
   isModule : Bool := false
-  /-- The module's direct imports. -/
-  imports : Array Import := #[]
-  /-- Pre-resolved artifacts of related modules (e.g., this module's transitive imports). -/
-  modules : NameMap ModuleArtifacts := {}
+  /--
+  The module's direct imports.
+  If `none`, uses the imports from the module header.
+  -/
+  imports? : Option (Array Import) := none
+  /-- Pre-resolved artifacts of transitively imported modules. -/
+  importArts : NameMap ImportArtifacts := {}
   /-- Dynamic libraries to load with the module. -/
   dynlibs : Array System.FilePath := #[]
   /-- Plugins to initialize with the module. -/
diff --git a/src/Lean/Util.lean b/src/Lean/Util.lean
index 85eae33284..81e0d41292 100644
--- a/src/Lean/Util.lean
+++ b/src/Lean/Util.lean
@@ -29,7 +29,6 @@ import Lean.Util.SCC
 import Lean.Util.TestExtern
 import Lean.Util.OccursCheck
 import Lean.Util.HasConstCache
-import Lean.Util.FileSetupInfo
 import Lean.Util.Heartbeats
 import Lean.Util.SearchPath
 import Lean.Util.SafeExponentiation
diff --git a/src/Lean/Util/FileSetupInfo.lean b/src/Lean/Util/FileSetupInfo.lean
deleted file mode 100644
index a5701c6f30..0000000000
--- a/src/Lean/Util/FileSetupInfo.lean
+++ /dev/null
@@ -1,15 +0,0 @@
-/-
-Copyright (c) 2023 Lean FRO, LLC. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Marc Huisinga
--/
-prelude
-import Lean.Util.LeanOptions
-
-/--
-Information shared between Lake and Lean when calling `lake setup-file`.
--/
-structure Lean.FileSetupInfo where
-  paths        : LeanPaths
-  setupOptions : LeanOptions
-  deriving FromJson, ToJson
diff --git a/src/Lean/Util/LeanOptions.lean b/src/Lean/Util/LeanOptions.lean
index d67cf15683..fec30882ba 100644
--- a/src/Lean/Util/LeanOptions.lean
+++ b/src/Lean/Util/LeanOptions.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Marc Huisinga
 -/
 prelude
-import Lean.Util.Paths
+import Lean.Data.Json
 
 namespace Lean
 
diff --git a/src/Lean/Util/Paths.lean b/src/Lean/Util/Paths.lean
deleted file mode 100644
index 50ae3f5cb4..0000000000
--- a/src/Lean/Util/Paths.lean
+++ /dev/null
@@ -1,24 +0,0 @@
-/-
-Copyright (c) 2021 Sebastian Ullrich. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-
-Authors: Sebastian Ullrich
--/
-prelude
-import Lean.Data.Json
-import Lean.Util.Path
-
-/-! Communicating Lean search paths between processes -/
-
-namespace Lean
-
-open System
-
-structure LeanPaths where
-  oleanPath       : SearchPath
-  srcPath         : SearchPath
-  loadDynlibPaths : Array FilePath := #[]
-  pluginPaths     : Array FilePath := #[]
-  deriving ToJson, FromJson
-
-end Lean
diff --git a/src/lake/Lake/Build/Module.lean b/src/lake/Lake/Build/Module.lean
index 3808772fd7..d97864a6b6 100644
--- a/src/lake/Lake/Build/Module.lean
+++ b/src/lake/Lake/Build/Module.lean
@@ -266,8 +266,8 @@ def Module.recFetchSetup (mod : Module) : FetchM (Job ModuleSetup) := ensureJob
     return {
       name := mod.name
       isModule := header.isModule
-      imports := header.imports
-      modules := {} -- TODO: modules
+      imports? := none
+      importArts := {} -- TODO
       dynlibs := dynlibs.map (·.path.toString)
       plugins := plugins.map (·.path.toString)
       options := mod.leanOptions
diff --git a/src/lake/Lake/CLI/Main.lean b/src/lake/Lake/CLI/Main.lean
index 94f371883c..431711bf2a 100644
--- a/src/lake/Lake/CLI/Main.lean
+++ b/src/lake/Lake/CLI/Main.lean
@@ -19,7 +19,7 @@ import Lake.CLI.Serve
 -- # CLI
 
 open System
-open Lean (Json toJson fromJson? LeanPaths NameMap)
+open Lean (Json toJson fromJson? NameMap)
 
 namespace Lake
 
@@ -445,10 +445,14 @@ protected def setupFile : CliM PUnit := do
   let loadConfig ← mkLoadConfig opts
   let buildConfig := mkBuildConfig opts
   let filePath ← takeArg "file path"
-  -- Additional arguments (imports) are ignored
-  -- TODO: Forbid them once the language server is updated
-  -- noArgsRem do
-  setupFile loadConfig filePath buildConfig
+  let header? ← takeArg?
+  noArgsRem do
+  let header ← header?.mapM  fun header => do
+    let header ← if header == "-" then IO.getStdin >>= (·.getLine) else pure header
+    match Json.parse header >>= fromJson? with
+    | .ok header => pure header
+    | .error e => error s!"failed to parse header JSON: {e}"
+  setupFile loadConfig filePath header buildConfig
 
 protected def test : CliM PUnit := do
   processOptions lakeOption
@@ -544,8 +548,8 @@ private def evalLeanFile
         withRegisterJob s!"{mod.name}:setup" do mod.setup.fetch
       mkArgs path setup (some mod.rootDir) mod.leanArgs
     else
-      let setup ← mkModuleSetup
-        ws leanFile.toString (← IO.FS.readFile path) ws.leanOptions buildConfig
+      let header ← Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
+      let setup ← mkModuleSetup ws leanFile.toString header ws.leanOptions buildConfig
       mkArgs path setup none ws.root.moreLeanArgs
   let spawnArgs : IO.Process.SpawnArgs := {
     args := args ++ moreArgs
diff --git a/src/lake/Lake/CLI/Serve.lean b/src/lake/Lake/CLI/Serve.lean
index 68a74962f6..10af416fb4 100644
--- a/src/lake/Lake/CLI/Serve.lean
+++ b/src/lake/Lake/CLI/Serve.lean
@@ -7,7 +7,6 @@ prelude
 import Lake.Load
 import Lake.Build
 import Lake.Util.MainM
-import Lean.Util.FileSetupInfo
 
 open Lean
 open System (FilePath)
@@ -23,52 +22,46 @@ and falls back to plain `lean --server`.
 -/
 def invalidConfigEnvVar := "LAKE_INVALID_CONFIG"
 
-private def mkLeanPaths (ws : Workspace) (setup : ModuleSetup) : LeanPaths where
-  oleanPath := ws.leanPath
-  srcPath := ws.leanSrcPath
-  loadDynlibPaths := setup.dynlibs
-  pluginPaths := setup.plugins
-
 def mkModuleSetup
-  (ws : Workspace) (fileName : String) (input : String) (opts : LeanOptions)
+  (ws : Workspace) (fileName : String) (header : ModuleHeader) (opts : LeanOptions)
   (buildConfig : BuildConfig)
 : IO ModuleSetup := do
-  let header ← Lean.parseImports' input fileName
   let {dynlibs, plugins} ← ws.runBuild (buildImportsAndDeps fileName header.imports) buildConfig
   return {
     name := ← Lean.moduleNameOfFileName fileName none
     isModule := header.isModule
-    imports := header.imports
-    modules := {} -- TODO
+    imports? := none
+    importArts := {} -- TODO
     dynlibs := dynlibs.map (·.path.toString)
     plugins := plugins.map (·.path.toString)
     options := opts
   }
 
 /--
-Build a list of imports of a file and print the `.olean` and source directories
-of every used package, as well as the server options for the file.
-If no configuration file exists, exit silently with `noConfigFileCode` (i.e, 2).
+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).
+
+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.
 -/
+-- TODO: Use `header?` for modules within the workspace as well.
 def setupFile
-  (loadConfig : LoadConfig) (path : FilePath) (buildConfig : BuildConfig := {})
+  (loadConfig : LoadConfig) (leanFile : FilePath)
+  (header? : Option ModuleHeader := none) (buildConfig : BuildConfig := {})
 : MainM PUnit := do
-  let path ← resolvePath path
+  let path ← resolvePath leanFile
   let configFile ← realConfigFile loadConfig.configFile
   if configFile.toString.isEmpty then
     exit noConfigFileCode
   else if configFile == path then do
-    let info : FileSetupInfo := {
-      paths := {
-        oleanPath := loadConfig.lakeEnv.leanPath
-        srcPath := loadConfig.lakeEnv.leanSrcPath
-        pluginPaths := #[loadConfig.lakeEnv.lake.sharedLib]
-      }
-      setupOptions := ⟨∅⟩
+    let setup : ModuleSetup := {
+      name := configModuleName
+      plugins :=  #[loadConfig.lakeEnv.lake.sharedLib]
     }
-    IO.println (toJson info).compress
+    IO.println (toJson setup).compress
   else if let some errLog := (← IO.getEnv invalidConfigEnvVar) then
     IO.eprint errLog
     IO.eprintln s!"Failed to configure the Lake workspace. Please restart the server after fixing the error above."
@@ -79,19 +72,13 @@ def setupFile
     if let some mod := ws.findModuleBySrc? path then
       let setup ← ws.runBuild (cfg := buildConfig) do
         withRegisterJob s!"{mod.name}:setup" do mod.setup.fetch
-      let info : FileSetupInfo := {
-        paths := mkLeanPaths ws setup
-        setupOptions := mod.serverOptions
-      }
-      IO.println (toJson info).compress
+      IO.println (toJson setup).compress
     else
+      let header ← header?.getDM do
+        Lean.parseImports' (← IO.FS.readFile path) leanFile.toString
       let setup ← mkModuleSetup
-        ws path.toString (← IO.FS.readFile path) ws.serverOptions buildConfig
-      let info : FileSetupInfo := {
-        paths := mkLeanPaths ws setup
-        setupOptions := setup.options
-      }
-      IO.println (toJson info).compress
+        ws leanFile.toString header ws.serverOptions buildConfig
+      IO.println (toJson setup).compress
 
 /--
 Start the Lean LSP for the `Workspace` loaded from `config`
@@ -102,8 +89,7 @@ def serve (config : LoadConfig) (args : Array String) : IO UInt32 := do
     let (ws?, log) ← (loadWorkspace config).captureLog
     log.replay (logger := MonadLog.stderr)
     if let some ws := ws? then
-      let ctx := mkLakeContext ws
-      pure (← LakeT.run ctx getAugmentedEnv, ws.root.moreGlobalServerArgs)
+      pure (ws.augmentedEnvVars, ws.root.moreGlobalServerArgs)
     else
       IO.eprintln "warning: package configuration has errors, falling back to plain `lean --server`"
       pure (config.lakeEnv.baseVars.push (invalidConfigEnvVar, log.toString), #[])
diff --git a/src/lake/Lake/Config/Workspace.lean b/src/lake/Lake/Config/Workspace.lean
index 12da46249f..54dec97cf8 100644
--- a/src/lake/Lake/Config/Workspace.lean
+++ b/src/lake/Lake/Config/Workspace.lean
@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Mac Malone
 -/
 prelude
-import Lean.Util.Paths
 import Lake.Config.FacetConfig
 import Lake.Config.TargetConfig
 import Lake.Config.Env
diff --git a/src/lake/examples/deps/test.sh b/src/lake/examples/deps/test.sh
index 5265ef4901..d597b3a197 100755
--- a/src/lake/examples/deps/test.sh
+++ b/src/lake/examples/deps/test.sh
@@ -17,7 +17,7 @@ $LAKE -d foo build --update
 ./bar/.lake/build/bin/bar
 
 # Test setup-file works (i.e., does not error)
-$LAKE -d foo setup-file ./foo/Foo.lean A B
+$LAKE -d foo setup-file ./foo/Foo.lean
 
 # Test `lake clean`
 test -d a/.lake/build
diff --git a/src/lake/tests/setupFile/test.sh b/src/lake/tests/setupFile/test.sh
index bec9d5a70b..963ff3cc96 100755
--- a/src/lake/tests/setupFile/test.sh
+++ b/src/lake/tests/setupFile/test.sh
@@ -8,24 +8,42 @@ source ../common.sh
 #---
 
 # Test that, by default, no plugins are used.
-test_out '"pluginPaths":[]' setup-file ImportFoo.lean
+test_out '"plugins":[]' setup-file ImportFoo.lean
 
 # Test that, by default, no dynlibs are used.
-test_out '"loadDynlibPaths":[]' setup-file ImportFoo.lean
+test_out '"dynlibs":[]' setup-file ImportFoo.lean
 
 # Test that, generally, no options are set.
-test_out '"setupOptions":{}' setup-file ImportFoo.lean
+test_out '"options":{}' setup-file ImportFoo.lean
 
 # Test that a more specific configuration will be used if
 # Lake can identify the module corresponding to the path.
-test_out '"setupOptions":{"weak.foo":"bar"}' setup-file Test.lean
+test_out '"options":{"weak.foo":"bar"}' setup-file Test.lean
 
 # Test that `setup-file` on an invalid Lean configuration file succeeds.
 test_run -f invalid.lean setup-file invalid.lean
 
 # Test that `setup-file` on a configuration file uses the Lake plugin,
 # even if the file is invalid and/or is not using a `Lake` import.
-test_not_out '"pluginPaths":[]' -f invalid.lean setup-file invalid.lean
+test_not_out '"plugins":[]' -f invalid.lean setup-file invalid.lean
+
+# Test that when a header is provided (via CLI or stdin),
+# it will be used instead of the file's header for an external module.
+test_out '"isModule":false' setup-file ImportFoo.lean
+HEADER_JSON='{"isModule":true,"imports":[]}'
+test_out '"isModule":true' setup-file ImportFoo.lean "$HEADER_JSON"
+echo "$HEADER_JSON" | test_out '"isModule":true' setup-file ImportFoo.lean -
+BOGUS_JSON='{"isModule":false,"imports":[{"module":"Test.Bogus","isMeta":false,"isExported":true,"importAll":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.
+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"
 
 # Cleanup
 rm -f produced.out
diff --git a/tests/pkg/setup/setup.json b/tests/pkg/setup/setup.json
index c079f5e6a8..d090a47ddb 100644
--- a/tests/pkg/setup/setup.json
+++ b/tests/pkg/setup/setup.json
@@ -9,10 +9,8 @@
       "isMeta": false
     }
   ],
-  "modules": {
-    "Dep": {
-      "olean": "Dep.olean"
-    }
+  "importArts": {
+    "Dep": ["Dep.olean"]
   },
   "dynlibs": [],
   "plugins": [],