fix: Runtime.markPersistent is unsafe (#6209)

This PR documents under which conditions `Runtime.markPersistent` is
unsafe and adjusts the elaborator accordingly

src/Init/Util.lean

@@ -90,10 +90,14 @@ def withPtrAddr {α : Type u} {β : Type v} (a : α) (k : USize → β) (h : ∀
 def Runtime.markMultiThreaded (a : α) : α := a
 
 /--
-  Marks given value and its object graph closure as persistent. This will remove
-  reference counter updates but prevent the closure from being deallocated until
-  the end of the process! It can still be useful to do eagerly when the value
-  will be marked persistent later anyway and there is available time budget to
-  mark it now or it would be unnecessarily marked multi-threaded in between. -/
+Marks given value and its object graph closure as persistent. This will remove
+reference counter updates but prevent the closure from being deallocated until
+the end of the process! It can still be useful to do eagerly when the value
+will be marked persistent later anyway and there is available time budget to
+mark it now or it would be unnecessarily marked multi-threaded in between.
+
+This function is only safe to use on objects (in the full closure) which are
+not used concurrently or which are already persistent.
+-/
 @[extern "lean_runtime_mark_persistent"]
-def Runtime.markPersistent (a : α) : α := a
+unsafe def Runtime.markPersistent (a : α) : α := a

src/Lean/CoreM.lean

@@ -31,6 +31,11 @@ register_builtin_option maxHeartbeats : Nat := {
   descr := "maximum amount of heartbeats per command. A heartbeat is number of (small) memory allocations (in thousands), 0 means no limit"
 }
 
+register_builtin_option Elab.async : Bool := {
+  defValue := false
+  descr := "perform elaboration using multiple threads where possible"
+}
+
 /--
 If the `diagnostics` option is not already set, gives a message explaining this option.
 Begins with a `\n`, so an error message can look like `m!"some error occurred{useDiagnosticMsg}"`.

src/Lean/Environment.lean

@@ -897,13 +897,18 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
        initialized constant. We have seen significant savings in `open Mathlib`
        timings, where we have both a big environment and interpreted environment
        extensions, from this. There is no significant extra cost to calling
-       `markPersistent` multiple times like this. -/
-    env := Runtime.markPersistent env
+       `markPersistent` multiple times like this.
+
+       Safety: There are no concurrent accesses to `env` at this point. -/
+    env := unsafe Runtime.markPersistent env
   env ← finalizePersistentExtensions env s.moduleData opts
   if leakEnv then
     /- Ensure the final environment including environment extension states is
-       marked persistent as documented. -/
-    env := Runtime.markPersistent env
+       marked persistent as documented.
+
+       Safety: There are no concurrent accesses to `env` at this point, assuming
+       extensions' `addImportFn`s did not spawn any unbound tasks. -/
+    env := unsafe Runtime.markPersistent env
   pure env
 
 @[export lean_import_modules]

src/Lean/Language/Lean.lean

@@ -454,11 +454,6 @@ where
         traceState
       }
       let prom ← IO.Promise.new
-      -- The speedup of these `markPersistent`s is negligible but they help in making unexpected
-      -- `inc_ref_cold`s more visible
-      let parserState := Runtime.markPersistent parserState
-      let cmdState := Runtime.markPersistent cmdState
-      let ctx := Runtime.markPersistent ctx
       parseCmd none parserState cmdState prom (sync := true) ctx
       return {
         diagnostics
@@ -648,16 +643,24 @@ where
     -- definitely resolve eventually
     snap.new.resolve <| .ofTyped { diagnostics := .empty : SnapshotLeaf }
 
-    let mut infoTree := cmdState.infoState.trees[0]!
+    let mut infoTree : InfoTree := cmdState.infoState.trees[0]!
     let cmdline := internal.cmdlineSnapshots.get scope.opts && !Parser.isTerminalCommand stx
-    if cmdline then
-      infoTree := Runtime.markPersistent infoTree
+    if cmdline && !Elab.async.get scope.opts then
+      /-
+      Safety: `infoTree` was created by `elabCommandTopLevel`. Thus it
+      should not have any concurrent accesses if we are on the cmdline and
+      async elaboration is disabled.
+      -/
+      -- TODO: we should likely remove this call when `Elab.async` is turned on
+      -- by default
+      infoTree := unsafe Runtime.markPersistent infoTree
     finishedPromise.resolve {
       diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
       infoTree? := infoTree
       traces := cmdState.traceState
       cmdState := if cmdline then {
-        env := Runtime.markPersistent cmdState.env
+        /- Safety: as above -/
+        env := unsafe Runtime.markPersistent cmdState.env
         maxRecDepth := 0
       } else cmdState
     }