fix: report replay kernel errors as standard diagnostics (#7471)

Avoids panics from follow-up cancellation errors

Fixes #7462

src/Lean/AddDecl.lean

@@ -75,7 +75,7 @@ def addDecl (decl : Declaration) : CoreM Unit := do
     setEnv async.asyncEnv
     doAdd
     async.commitCheckEnv (← getEnv)
-  let t ← BaseIO.mapTask (fun _ => checkAct) env.checked
+  let t ← BaseIO.mapTask checkAct env.checked
   let endRange? := (← getRef).getTailPos?.map fun pos => ⟨pos, pos⟩
   Core.logSnapshotTask { stx? := none, reportingRange? := endRange?, task := t, cancelTk? := cancelTk }
 where doAdd := do

src/Lean/CoreM.lean

@@ -410,17 +410,17 @@ def logSnapshotTask (task : Language.SnapshotTask Language.SnapshotTree) : CoreM
   modify fun s => { s with snapshotTasks := s.snapshotTasks.push task }
 
 /-- Wraps the given action for use in `EIO.asTask` etc., discarding its final monadic state. -/
-def wrapAsync (act : Unit → CoreM α) (cancelTk? : Option IO.CancelToken) :
-    CoreM (EIO Exception α) := do
+def wrapAsync {α : Type} (act : α → CoreM β) (cancelTk? : Option IO.CancelToken) :
+    CoreM (α → EIO Exception β) := do
   let st ← get
   let ctx ← read
   let ctx := { ctx with cancelTk? }
   let heartbeats := (← IO.getNumHeartbeats) - ctx.initHeartbeats
-  return withCurrHeartbeats (do
+  return fun a => withCurrHeartbeats (do
       -- include heartbeats since start of elaboration in new thread as well such that forking off
       -- an action doesn't suddenly allow it to succeed from a lower heartbeat count
       IO.addHeartbeats heartbeats
-      act () : CoreM _)
+      act a : CoreM _)
     |>.run' ctx st
 
 /-- Option for capturing output to stderr during elaboration. -/
@@ -457,9 +457,9 @@ Wraps the given action for use in `BaseIO.asTask` etc., discarding its final sta
 token, if any, should be stored in a `SnapshotTask` for the server to trigger it when the result is
 no longer needed.
 -/
-def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (cancelTk? : Option IO.CancelToken)
-    (desc : String := by exact decl_name%.toString) : CoreM (BaseIO SnapshotTree) := do
-  let t ← wrapAsync (cancelTk? := cancelTk?) fun _ => do
+def wrapAsyncAsSnapshot {α : Type} (act : α → CoreM Unit) (cancelTk? : Option IO.CancelToken)
+    (desc : String := by exact decl_name%.toString) : CoreM (α → BaseIO SnapshotTree) := do
+  let f ← wrapAsync (cancelTk? := cancelTk?) fun a => do
     IO.FS.withIsolatedStreams (isolateStderr := stderrAsMessages.get (← getOptions)) do
       let tid ← IO.getTID
       -- reset trace/info state and message log so as not to report them twice
@@ -471,7 +471,7 @@ def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (cancelTk? : Option IO.Cance
       }
       try
         withTraceNode `Elab.async (fun _ => return desc) do
-          act ()
+          act a
       catch e =>
         unless e.isInterrupt do
           logError e.toMessageData
@@ -479,8 +479,8 @@ def wrapAsyncAsSnapshot (act : Unit → CoreM Unit) (cancelTk? : Option IO.Cance
         addTraceAsMessages
       get
   let ctx ← readThe Core.Context
-  return do
-    match (← t.toBaseIO) with
+  return fun a => do
+    match (← (f a).toBaseIO) with
     | .ok (output, st) => mkSnapshot output ctx st desc
     -- interrupt or abort exception as `try catch` above should have caught any others
     | .error _ => default
@@ -568,7 +568,7 @@ partial def compileDecls (decls : List Name) (ref? : Option Declaration := none)
       doCompile
     finally
       res.commitChecked (← getEnv)
-  let t ← BaseIO.mapTask (fun _ => checkAct) env.checked
+  let t ← BaseIO.mapTask checkAct env.checked
   let endRange? := (← getRef).getTailPos?.map fun pos => ⟨pos, pos⟩
   Core.logSnapshotTask { stx? := none, reportingRange? := endRange?, task := t, cancelTk? := cancelTk }
 where doCompile := do

src/Lean/Elab/MutualDef.lean

@@ -1115,8 +1115,7 @@ private def logGoalsAccomplishedSnapshotTask (views : Array DefView)
           continue
         logGoalsAccomplished
       | _ => continue
-  let logGoalsAccomplishedTask ← BaseIO.mapTask (t := ← tree.waitAll) fun _ =>
-    logGoalsAccomplishedAct
+  let logGoalsAccomplishedTask ← BaseIO.mapTask (t := ← tree.waitAll) logGoalsAccomplishedAct
   Core.logSnapshotTask {
     stx? := none
     -- Use first line of the mutual block to avoid covering the progress of the whole mutual block

src/Lean/Elab/Term.lean

@@ -484,15 +484,15 @@ def withoutTacticReuse [Monad m] [MonadWithReaderOf Context m] [MonadOptions m]
   }) act
 
 @[inherit_doc Core.wrapAsyncAsSnapshot]
-def wrapAsyncAsSnapshot (act : Unit → TermElabM Unit) (cancelTk? : Option IO.CancelToken)
+def wrapAsyncAsSnapshot {α : Type} (act : α → TermElabM Unit) (cancelTk? : Option IO.CancelToken)
     (desc : String := by exact decl_name%.toString) :
-    TermElabM (BaseIO Language.SnapshotTree) := do
+    TermElabM (α → BaseIO Language.SnapshotTree) := do
   let ctx ← read
   let st ← get
   let metaCtx ← readThe Meta.Context
   let metaSt ← getThe Meta.State
-  Core.wrapAsyncAsSnapshot (cancelTk? := cancelTk?) (desc := desc) fun _ =>
-    act () |>.run ctx |>.run' st |>.run' metaCtx metaSt
+  Core.wrapAsyncAsSnapshot (cancelTk? := cancelTk?) (desc := desc) fun a =>
+    act a |>.run ctx |>.run' st |>.run' metaCtx metaSt
 
 abbrev TermElabResult (α : Type) := EStateM.Result Exception SavedState α
 

src/Lean/Environment.lean

@@ -262,27 +262,6 @@ inductive Exception where
   | interrupted
 deriving Nonempty
 
-/-- Basic `Exception` formatting without `MessageData` dependency. -/
-private def Exception.toRawString : Kernel.Exception → String
-  | unknownConstant _ constName       => s!"(kernel) unknown constant '{constName}'"
-  | alreadyDeclared _ constName       => s!"(kernel) constant has already been declared '{constName}'"
-  | declTypeMismatch _ _ _ => s!"(kernel) declaration type mismatch"
-  | declHasMVars _ constName _        => s!"(kernel) declaration has metavariables '{constName}'"
-  | declHasFVars _ constName _        => s!"(kernel) declaration has free variables '{constName}'"
-  | funExpected _ _ e                 => s!"(kernel) function expected: {e}"
-  | typeExpected _ _ e                => s!"(kernel) type expected: {e}"
-  | letTypeMismatch  _ _ n _ _        => s!"(kernel) let-declaration type mismatch '{n}'"
-  | exprTypeMismatch _ _ e _          => s!"(kernel) type mismatch at {e}"
-  | appTypeMismatch  _ _ e fnType argType =>
-    s!"application type mismatch: {e}\nargument has type {argType}\nbut function has type {fnType}"
-  | invalidProj _ _ e                 => s!"(kernel) invalid projection {e}"
-  | thmTypeIsNotProp _ constName type => s!"(kernel) type of theorem '{constName}' is not a proposition: {type}"
-  | other msg                         => s!"(kernel) {msg}"
-  | deterministicTimeout              => "(kernel) deterministic timeout"
-  | excessiveMemory                   => "(kernel) excessive memory consumption detected"
-  | deepRecursion                     => "(kernel) deep recursion detected"
-  | interrupted                       => "(kernel) interrupted"
-
 namespace Environment
 
 @[export lean_environment_find]
@@ -489,7 +468,7 @@ private structure RealizationContext where
   changes onto a derived kernel environment, and auxiliary data (always `SnapshotTree` in builtin
   uses, but untyped to avoid cyclic module references).
   -/
-  constsRef : IO.Ref (NameMap (Task (List AsyncConst × (Kernel.Environment → Kernel.Environment) × Dynamic)))
+  constsRef : IO.Ref (NameMap (Task (List AsyncConst × (Kernel.Environment → Except Kernel.Exception Kernel.Environment) × Dynamic)))
 
 /--
 Elaboration-specific extension of `Kernel.Environment` that adds tracking of asynchronously
@@ -1927,11 +1906,11 @@ def replayConsts (dest : Environment) (oldEnv newEnv : Environment) (skipExistin
         consts
       else
         consts.add c
-    checked := dest.checked.map (replayKernel exts consts)
+    checked := dest.checked.map fun kenv => replayKernel exts consts kenv |>.toOption.getD kenv
   }
 where
   replayKernel (exts : Array (EnvExtension EnvExtensionState)) (consts : List AsyncConst)
-      (kenv : Kernel.Environment) : Kernel.Environment := Id.run do
+      (kenv : Kernel.Environment) : Except Kernel.Exception Kernel.Environment := do
     let mut kenv := kenv
     for c in consts do
       if skipExisting && (kenv.find? c.constInfo.name).isSome then
@@ -1947,16 +1926,13 @@ where
       -- for panics
       let _ : Inhabited Kernel.Environment := ⟨kenv⟩
       let decl ← match info with
-        | .thmInfo thm   => .thmDecl thm
-        | .defnInfo defn => .defnDecl defn
+        | .thmInfo thm   => pure <| .thmDecl thm
+        | .defnInfo defn => pure <| .defnDecl defn
         | _              =>
           return panic! s!"{c.constInfo.name} must be definition/theorem"
       -- realized kernel additions cannot be interrupted - which would be bad anyway as they can be
       -- reused between snapshots
-      match kenv.addDeclCore 0 decl none with
-      | .ok kenv' => kenv := kenv'
-      | .error e =>
-        return panic! s!"failed to add {c.constInfo.name} to environment\n{e.toRawString}"
+      kenv ← ofExcept <| kenv.addDeclCore 0 decl none
     for ext in exts do
       if let some replay := ext.replay? then
         kenv := { kenv with
@@ -1992,7 +1968,7 @@ def hasUnsafe (env : Environment) (e : Expr) : Bool :=
 /-- Plumbing function for `Lean.Meta.realizeConst`; see documentation there. -/
 def realizeConst (env : Environment) (forConst : Name) (constName : Name)
     (realize : Environment → Options → BaseIO (Environment × Dynamic)) :
-    IO (Environment × Dynamic) := do
+    IO (Environment × Task (Option Kernel.Exception) × Dynamic) := do
   -- the following code is inherently non-deterministic in number of heartbeats, reset them at the
   -- end
   let heartbeats ← IO.getNumHeartbeats
@@ -2045,19 +2021,25 @@ def realizeConst (env : Environment) (forConst : Name) (constName : Name)
       let replay := replayConsts.replayKernel (skipExisting := true) realizeEnv realizeEnv' exts consts
       prom.resolve (consts, replay, dyn)
       pure (consts, replay, dyn)
+    let exPromise ← IO.Promise.new
     let env := { env with
       asyncConsts := consts.foldl (init := env.asyncConsts) fun consts c =>
         if consts.find? c.constInfo.name |>.isSome then
           consts
         else
           consts.add c
-      checked := env.checked.map replay
+      checked := (← BaseIO.mapTask (t := env.checked) fun kenv => do
+        match replay kenv with
+        | .ok kenv => return kenv
+        | .error e =>
+          exPromise.resolve e
+          return kenv)
       allRealizations := env.allRealizations.map (sync := true) fun allRealizations =>
         consts.foldl (init := allRealizations) fun allRealizations c =>
           allRealizations.insert c.constInfo.name c
     }
     IO.setNumHeartbeats heartbeats
-    return (env, dyn)
+    return (env, exPromise.result?, dyn)
 
 end Environment
 

src/Lean/Meta/Basic.lean

@@ -2278,7 +2278,15 @@ def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) :
       -- heartbeat limits inside `realizeAndReport` should be measured from this point on
       initHeartbeats := (← IO.getNumHeartbeats)
     }
-    let (env, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
+    let (env, exTask, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
+    let exAct ← Core.wrapAsyncAsSnapshot (cancelTk? := none) fun
+      | none => return
+      | some ex => do
+        logError <| ex.toMessageData (← getOptions)
+    Core.logSnapshotTask {
+      stx? := none
+      task := (← BaseIO.mapTask (t := exTask) exAct)
+    }
     if let some res := dyn.get? RealizeConstantResult then
       let mut snap := res.snap
       -- localize diagnostics

src/Lean/Server/Test/Cancel.lean

@@ -73,7 +73,7 @@ elab "wait_for_unblock_async" : tactic => do
     if (← cancelTk.isSet) then
       IO.eprintln "cancelled!"
       log "cancelled (should never be visible)"
-  let t ← BaseIO.asTask act
+  let t ← BaseIO.asTask (act ())
   Core.logSnapshotTask { stx? := none, task := t, cancelTk? := cancelTk }
 
   log "blocked"
@@ -105,7 +105,7 @@ elab_rules : tactic
     log "cancelled (should never be visible)"
     prom.resolve ()
     Core.checkInterrupted
-  let t ← BaseIO.asTask act
+  let t ← BaseIO.asTask (act ())
   Core.logSnapshotTask { stx? := none, task := t, cancelTk? := cancelTk }
 
   dbg_trace "blocked!"

tests/lean/run/6065.lean

@@ -10,7 +10,7 @@ def foo (r : Nat) : Nat :=
       | _ => 0
 
 example {r : Nat} : foo r = 0 := by
-  simp only [foo.eq_def]
+  simp only [foo]
   guard_target =ₛ
     (match r with
         | Nat.zero => 0

tests/lean/run/998Export.lean

@@ -75,4 +75,5 @@ def exportLevel (L : Level) : ExportM Nat := do
       let i ← alloc L; IO.println s!"{i} #UP {← exportName n}"; pure i
     | Level.mvar n => unreachable!
 
-attribute [simp] exportLevel
+-- TODO: this test has been broken for a while with a panic that was ignored by the test suite
+--attribute [simp] exportLevel