fix: mainModuleName should use srcSearchPath (#4066)

As [reported on
Zulip](https://leanprover.zulipchat.com/#narrow/stream/341532-lean4-dev/topic/Find.20references.20broken.20in.20lean.20core/near/437051935).
The `mainModuleName` was being set incorrectly when browsing lean core
sources, resulting in failure of cross-file server requests like "Find
References". Because the `srcSearchPath` is generated asynchronously, we
store it as a `Task Name` which is resolved some time before the header
is finished parsing. (I don't think the `.get` here will ever block,
because the srcSearchPath will be ready by the time the initial command
snap is requested.)

---------

Co-authored-by: Sebastian Ullrich <sebasti@nullri.ch>

src/Lean/Elab/Frontend.lean

@@ -154,9 +154,9 @@ def runFrontend
 
     return (s.commandState.env, !s.commandState.messages.hasErrors)
 
-  let ctx := { inputCtx with mainModuleName, opts, trustLevel }
+  let ctx := { inputCtx with }
   let processor := Language.Lean.process
-  let snap ← processor none ctx
+  let snap ← processor (fun _ => pure <| .ok { mainModuleName, opts, trustLevel }) none ctx
   let snaps := Language.toSnapshotTree snap
   snaps.runAndReport opts jsonOutput
   if let some ileanFileName := ileanFileName? then

src/Lean/Language/Basic.lean

@@ -245,17 +245,8 @@ def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) (json := false
 def SnapshotTree.getAll (s : SnapshotTree) : Array Snapshot :=
   s.forM (m := StateM _) (fun s => modify (·.push s)) |>.run #[] |>.2
 
-/-- Metadata that does not change during the lifetime of the language processing process. -/
-structure ModuleProcessingContext where
-  /-- Module name of the file being processed. -/
-  mainModuleName : Name
-  /-- Options provided outside of the file content, e.g. on the cmdline or in the lakefile. -/
-  opts : Options
-  /-- Kernel trust level. -/
-  trustLevel : UInt32 := 0
-
 /-- Context of an input processing invocation. -/
-structure ProcessingContext extends ModuleProcessingContext, Parser.InputContext
+structure ProcessingContext extends Parser.InputContext
 
 /-- Monad transformer holding all relevant data for processing. -/
 abbrev ProcessingT m := ReaderT ProcessingContext m
@@ -296,10 +287,10 @@ end Language
 /--
   Builds a function for processing a language using incremental snapshots by passing the previous
   snapshot to `Language.process` on subsequent invocations. -/
-def Language.mkIncrementalProcessor (process : Option InitSnap → ProcessingM InitSnap)
-    (ctx : ModuleProcessingContext) : BaseIO (Parser.InputContext → BaseIO InitSnap) := do
+def Language.mkIncrementalProcessor (process : Option InitSnap → ProcessingM InitSnap) :
+    BaseIO (Parser.InputContext → BaseIO InitSnap) := do
   let oldRef ← IO.mkRef none
   return fun ictx => do
-    let snap ← process (← oldRef.get) { ctx, ictx with }
+    let snap ← process (← oldRef.get) { ictx with }
     oldRef.set (some snap)
     return snap

src/Lean/Language/Lean.lean

@@ -263,7 +263,28 @@ private def withHeaderExceptions (ex : Snapshot → α) (act : LeanProcessingT I
   | .error e => return ex { diagnostics := (← diagnosticsOfHeaderError e.toString) }
   | .ok a => return a
 
-/-- Entry point of the Lean language processor. -/
+/--
+Result of retrieving additional metadata about the current file after parsing imports. In the
+language server, these are derived from the `lake setup-file` result. On the cmdline and for similar
+simple uses, these can be computed eagerly without looking at the imports.
+-/
+structure SetupImportsResult where
+  /-- Module name of the file being processed. -/
+  mainModuleName : Name
+  /-- Options provided outside of the file content, e.g. on the cmdline or in the lakefile. -/
+  opts : Options
+  /-- Kernel trust level. -/
+  trustLevel : UInt32 := 0
+
+/--
+Entry point of the Lean language processor.
+
+The `setupImports` function is called after the header has been parsed and before the first command
+is parsed in order to supply additional file metadata (or abort with a given terminal snapshot); see
+`SetupImportsResult`.
+
+`old?` is a previous resulting snapshot, if any, to be reused for incremental processing.
+-/
 /-
 General notes:
 * For each processing function we pass in the previous state, if any, in order to reuse still-valid
@@ -277,8 +298,7 @@ General notes:
   fast-forwarded snapshots without having to wait on tasks.
 -/
 partial def process
-    (setupImports : Syntax → ProcessingT IO (Except HeaderProcessedSnapshot Options) :=
-      fun _ => pure <| .ok {})
+    (setupImports : Syntax → ProcessingT IO (Except HeaderProcessedSnapshot SetupImportsResult))
     (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
   -- compute position of syntactic change once
   let firstDiffPos? := old?.map (·.ictx.input.firstDiffPos (← read).input)
@@ -354,20 +374,18 @@ where
     SnapshotTask.ofIO (some ⟨0, ctx.input.endPos⟩) <|
     ReaderT.run (r := ctx) <|  -- re-enter reader in new task
     withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none }) do
-      let opts ← match (← setupImports stx) with
-        | .ok opts => pure opts
+      let setup ← match (← setupImports stx) with
+        | .ok setup => pure setup
         | .error snap => return snap
-      -- override context options with file options
-      let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) opts
       -- allows `headerEnv` to be leaked, which would live until the end of the process anyway
-      let (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx opts .empty
-        ctx.toInputContext ctx.trustLevel
+      let (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx setup.opts .empty
+        ctx.toInputContext setup.trustLevel
       let diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
       if msgLog.hasErrors then
         return { diagnostics, result? := none }
 
-      let headerEnv := headerEnv.setMainModule ctx.mainModuleName
-      let cmdState := Elab.Command.mkState headerEnv msgLog opts
+      let headerEnv := headerEnv.setMainModule setup.mainModuleName
+      let cmdState := Elab.Command.mkState headerEnv msgLog setup.opts
       let cmdState := { cmdState with infoState := {
         enabled := true
         trees := #[Elab.InfoTree.context (.commandCtx {

src/Lean/Server/FileWorker.lean

@@ -265,9 +265,9 @@ open Language Lean in
 Callback from Lean language processor after parsing imports that requests necessary information from
 Lake for processing imports.
 -/
-def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
+def setupImports (meta : DocumentMeta) (cmdlineOpts : Options) (chanOut : Channel JsonRpc.Message)
     (srcSearchPathPromise : Promise SearchPath) (stx : Syntax) :
-    Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot Options) := do
+    Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot SetupImportsResult) := do
   let importsAlreadyLoaded ← importsLoadedRef.modifyGet ((·, true))
   if importsAlreadyLoaded then
     -- As we never unload imports in the server, we should not run the code below twice in the
@@ -306,27 +306,38 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
   | _ => pure ()
 
   srcSearchPathPromise.resolve fileSetupResult.srcSearchPath
-  return .ok fileSetupResult.fileOptions
+
+  let mainModuleName ← if let some path := System.Uri.fileUriToPath? meta.uri then
+    EIO.catchExceptions (h := fun _ => pure Name.anonymous) do
+      if let some mod ← searchModuleNameOfFileName path fileSetupResult.srcSearchPath then
+        pure mod
+      else
+        moduleNameOfFileName path none
+  else
+    pure Name.anonymous
+
+  -- override cmdline options with file options
+  let opts := cmdlineOpts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
+
+  return .ok {
+    mainModuleName
+    opts
+  }
 
 /- Worker initialization sequence. -/
 section Initialization
   def initializeWorker (meta : DocumentMeta) (o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
       : IO (WorkerContext × WorkerState) := do
     let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
-    let mut mainModuleName := Name.anonymous
-    try
-      if let some path := System.Uri.fileUriToPath? meta.uri then
-        mainModuleName ← moduleNameOfFileName path none
-    catch _ => pure ()
     let maxDocVersionRef ← IO.mkRef 0
-    let freshRequestIdRef ← IO.mkRef 0
+    let freshRequestIdRef ← IO.mkRef (0 : Int)
     let chanIsProcessing ← IO.Channel.new
     let stickyDiagnosticsRef ← IO.mkRef ∅
     let chanOut ← mkLspOutputChannel maxDocVersionRef chanIsProcessing
     let srcSearchPathPromise ← IO.Promise.new
 
-    let processor := Language.Lean.process (setupImports meta chanOut srcSearchPathPromise)
-    let processor ← Language.mkIncrementalProcessor processor { opts, mainModuleName }
+    let processor := Language.Lean.process (setupImports meta opts chanOut srcSearchPathPromise)
+    let processor ← Language.mkIncrementalProcessor processor
     let initSnap ← processor meta.mkInputContext
     let _ ← IO.mapTask (t := srcSearchPathPromise.result) fun srcSearchPath => do
       let importClosure := getImportClosure? initSnap