From 00718c3959d93972b43ee30ac008e9d655d9f151 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Tue, 3 Dec 2024 13:42:02 +0100
Subject: [PATCH] chore: clean up `Elab.async` handling (#6299)

* Make sure metaprogramming users cannot be surprised by its
introduction
* Make `#guard_msgs` compatible with its use
---
 src/Lean/CoreM.lean          | 12 ++++++++++--
 src/Lean/Elab/Command.lean   |  7 ++++++-
 src/Lean/Elab/GuardMsgs.lean | 16 +++++++++-------
 src/Lean/Language/Lean.lean  |  2 ++
 4 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/Lean/CoreM.lean b/src/Lean/CoreM.lean
index 7b3037d05894..ff774004492f 100644
--- a/src/Lean/CoreM.lean
+++ b/src/Lean/CoreM.lean
@@ -32,8 +32,16 @@ register_builtin_option maxHeartbeats : Nat := {
 }
 
 register_builtin_option Elab.async : Bool := {
-  defValue := true
-  descr := "perform elaboration using multiple threads where possible"
+  defValue := false
+  descr := "perform elaboration using multiple threads where possible\
+    \n\
+    \nThis option defaults to `false` but (when not explicitly set) is overridden to `true` in \
+      `Lean.Language.Lean.process` as used by the cmdline driver and language server. \
+      Metaprogramming users driving elaboration directly via e.g. \
+      `Lean.Elab.Command.elabCommandTopLevel` can opt into asynchronous elaboration by setting \
+      this option but then are responsible for processing messages and other data not only in the \
+      resulting command state but also from async tasks in `Lean.Command.Context.snap?` and \
+      `Lean.Command.State.snapshotTasks`."
 }
 
 /--
diff --git a/src/Lean/Elab/Command.lean b/src/Lean/Elab/Command.lean
index 1f99971162a9..4ed5181437e2 100644
--- a/src/Lean/Elab/Command.lean
+++ b/src/Lean/Elab/Command.lean
@@ -101,7 +101,7 @@ structure Context where
   (mutual) defs and contained tactics, in which case the `DynamicSnapshot` is a
   `HeadersParsedSnapshot`.
 
-  Definitely resolved in `Language.Lean.process.doElab`.
+  Definitely resolved in `Lean.Elab.Command.elabCommandTopLevel`.
 
   Invariant: if the bundle's `old?` is set, the context and state at the beginning of current and
   old elaboration are identical.
@@ -562,6 +562,11 @@ def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do pro
       withLogging do
         runLintersAsync stx
   finally
+    -- Make sure `snap?` is definitely resolved; we do not use it for reporting as `#guard_msgs` may
+    -- be the caller of this function and add new messages and info trees
+    if let some snap := (← read).snap? then
+      snap.new.resolve default
+
     -- note the order: first process current messages & info trees, then add back old messages & trees,
     -- then convert new traces to messages
     let mut msgs := (← get).messages
diff --git a/src/Lean/Elab/GuardMsgs.lean b/src/Lean/Elab/GuardMsgs.lean
index f48641fafbfb..0bd33fcc5a9b 100644
--- a/src/Lean/Elab/GuardMsgs.lean
+++ b/src/Lean/Elab/GuardMsgs.lean
@@ -140,13 +140,15 @@ def MessageOrdering.apply (mode : MessageOrdering) (msgs : List String) : List S
         |>.trim |> removeTrailingWhitespaceMarker
     let (whitespace, ordering, specFn) ← parseGuardMsgsSpec spec?
     let initMsgs ← modifyGet fun st => (st.messages, { st with messages := {} })
-    -- disable async elaboration for `cmd` so `msgs` will contain all messages
-    let async := Elab.async.get (← getOptions)
-    modifyScope fun scope => { scope with opts := Elab.async.set scope.opts false }
-    -- The `#guard_msgs` command is special-cased in `elabCommandTopLevel` to ensure linters only run once.
-    elabCommandTopLevel cmd
-    modifyScope fun scope => { scope with opts := Elab.async.set scope.opts async }
-    let msgs := (← get).messages
+    -- do not forward snapshot as we don't want messages assigned to it to leak outside
+    withReader ({ · with snap? := none }) do
+      -- The `#guard_msgs` command is special-cased in `elabCommandTopLevel` to ensure linters only run once.
+      elabCommandTopLevel cmd
+    -- collect sync and async messages
+    let msgs := (← get).messages ++
+      (← get).snapshotTasks.foldl (· ++ ·.get.getAll.foldl (· ++ ·.diagnostics.msgLog) {}) {}
+    -- clear async messages as we don't want them to leak outside
+    modify ({ · with snapshotTasks := #[] })
     let mut toCheck : MessageLog := .empty
     let mut toPassthrough : MessageLog := .empty
     for msg in msgs.toList do
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 33a0807bee97..c1a467f4a68b 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -433,6 +433,8 @@ where
         }
       -- now that imports have been loaded, check options again
       let opts ← reparseOptions setup.opts
+      -- default to async elaboration; see also `Elab.async` docs
+      let opts := Elab.async.setIfNotSet opts true
       let cmdState := Elab.Command.mkState headerEnv msgLog opts
       let cmdState := { cmdState with
         infoState := {