fix: broken trace tree on elab runtime exception (#3371)

RELEASES.md

@@ -66,6 +66,11 @@ v4.7.0 (development in progress)
     rw [h]
   ```
 
+Breaking changes:
+* `Lean.withTraceNode` and variants got a stronger `MonadAlwaysExcept` assumption to
+  fix trace trees not being built on elaboration runtime exceptions. Instances for most elaboration
+  monads built on `EIO Exception` should be synthesized automatically.
+
 v4.6.0
 ---------
 

src/Lean/CoreM.lean

@@ -363,7 +363,7 @@ def Exception.isRuntime (ex : Exception) : Bool :=
 
 /--
 Custom `try-catch` for all monads based on `CoreM`. We don't want to catch "runtime exceptions"
-in these monads, but on `CommandElabM`. See issues #2775 and #2744
+in these monads, but on `CommandElabM`. See issues #2775 and #2744 as well as `MonadAlwayExcept`.
 -/
 @[inline] protected def Core.tryCatch (x : CoreM α) (h : Exception → CoreM α) : CoreM α := do
   try

src/Lean/Elab/Command.lean

@@ -310,28 +310,29 @@ macro expansion etc.
 def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do profileitM Exception "elaboration" (← getOptions) do
   let initMsgs ← modifyGet fun st => (st.messages, { st with messages := {} })
   let initInfoTrees ← getResetInfoTrees
-  -- We should *not* factor out `elabCommand`'s `withLogging` to here since it would make its error
-  -- recovery more coarse. In particular, If `c` in `set_option ... in $c` fails, the remaining
-  -- `end` command of the `in` macro would be skipped and the option would be leaked to the outside!
-  elabCommand stx
-  withLogging do
-    runLinters stx
-
-  -- 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
-  -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
-  if !showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors && stx.hasMissing then
-    -- discard elaboration errors, except for a few important and unlikely misleading ones, on parse error
-    msgs := ⟨msgs.msgs.filter fun msg =>
-      msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder || tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
-  for tree in (← getInfoTrees) do
-    trace[Elab.info] (← tree.format)
-  modify fun st => { st with
-    messages := initMsgs ++ msgs
-    infoState := { st.infoState with trees := initInfoTrees ++ st.infoState.trees }
-  }
-  addTraceAsMessages
+  try
+    -- We should *not* factor out `elabCommand`'s `withLogging` to here since it would make its error
+    -- recovery more coarse. In particular, If `c` in `set_option ... in $c` fails, the remaining
+    -- `end` command of the `in` macro would be skipped and the option would be leaked to the outside!
+    elabCommand stx
+    withLogging do
+      runLinters stx
+  finally
+    -- 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
+    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
+    if !showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors && stx.hasMissing then
+      -- discard elaboration errors, except for a few important and unlikely misleading ones, on parse error
+      msgs := ⟨msgs.msgs.filter fun msg =>
+        msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder || tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
+    for tree in (← getInfoTrees) do
+      trace[Elab.info] (← tree.format)
+    modify fun st => { st with
+      messages := initMsgs ++ msgs
+      infoState := { st.infoState with trees := initInfoTrees ++ st.infoState.trees }
+    }
+    addTraceAsMessages
 
 /-- Adapt a syntax transformation to a regular, command-producing elaborator. -/
 def adaptExpander (exp : Syntax → CommandElabM Syntax) : CommandElab := fun stx => do
@@ -400,6 +401,9 @@ def liftTermElabM (x : TermElabM α) : CommandElabM α := do
   -- We execute `x` with an empty message log. Thus, `x` cannot modify/view messages produced by previous commands.
   -- This is useful for implementing `runTermElabM` where we use `Term.resetMessageLog`
   let x : TermElabM _  := withSaveInfoContext x
+  -- make sure `observing` below also catches runtime exceptions (like we do by default in
+  -- `CommandElabM`)
+  let _ := MonadAlwaysExcept.except (m := TermElabM)
   let x : MetaM _      := (observing x).run (mkTermContext ctx s) { levelNames := scope.levelNames }
   let x : CoreM _      := x.run mkMetaContext {}
   let x : EIO _ _      := x.run (mkCoreContext ctx s heartbeats) { env := s.env, ngen := s.ngen, nextMacroScope := s.nextMacroScope, infoState.enabled := s.infoState.enabled, traceState := s.traceState }
@@ -412,9 +416,7 @@ def liftTermElabM (x : TermElabM α) : CommandElabM α := do
     traceState.traces := coreS.traceState.traces.map fun t => { t with ref := replaceRef t.ref ctx.ref }
     messages          := s.messages ++ coreS.messages
   }
-  match ea with
-  | Except.ok a     => pure a
-  | Except.error ex => throw ex
+  MonadExcept.ofExcept ea
 
 /--
 Execute the monadic action `elabFn xs` as a `CommandElabM` monadic action, where `xs` are free variables

src/Lean/Util/Trace.lean

@@ -181,8 +181,34 @@ def shouldProfile : m Bool := do
 def shouldEnableNestedTrace (cls : Name) (secs : Float) : m Bool := do
   return (← isTracingEnabledFor cls) || secs < trace.profiler.threshold.getSecs (← getOptions)
 
-def withTraceNode [MonadExcept ε m] [MonadLiftT BaseIO m] (cls : Name) (msg : Except ε α → m MessageData) (k : m α)
-    (collapsed := true) : m α := do
+/--
+`MonadExcept` variant that is expected to catch all exceptions of the given type in case the
+standard instance doesn't.
+
+In most circumstances, we want to let runtime exceptions during term elaboration bubble up to the
+command elaborator (see `Core.tryCatch`). However, in a few cases like building the trace tree, we
+really need to handle (and then re-throw) every exception lest we end up with a broken tree.
+-/
+class MonadAlwaysExcept (ε : outParam (Type u)) (m : Type u → Type v) where
+  except : MonadExceptOf ε m
+
+-- instance sufficient for inferring `MonadAlwaysExcept` for the elaboration monads
+
+instance (ε) : MonadAlwaysExcept ε (EIO ε) where
+  except := inferInstance
+
+instance (ε) [always : MonadAlwaysExcept ε m] : MonadAlwaysExcept ε (StateT σ m) where
+  except := let _ := always.except; inferInstance
+
+instance (ε) [always : MonadAlwaysExcept ε m] : MonadAlwaysExcept ε (StateRefT' ω σ m) where
+  except := let _ := always.except; inferInstance
+
+instance (ε) [always : MonadAlwaysExcept ε m] : MonadAlwaysExcept ε (ReaderT ρ m) where
+  except := let _ := always.except; inferInstance
+
+def withTraceNode [always : MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] (cls : Name)
+    (msg : Except ε α → m MessageData) (k : m α) (collapsed := true) : m α := do
+  let _ := always.except
   let opts ← getOptions
   let clsEnabled ← isTracingEnabledFor cls
   unless clsEnabled || trace.profiler.get opts do
@@ -200,7 +226,8 @@ def withTraceNode [MonadExcept ε m] [MonadLiftT BaseIO m] (cls : Name) (msg : E
   addTraceNode oldTraces cls ref m collapsed
   MonadExcept.ofExcept res
 
-def withTraceNode' [MonadExcept Exception m] [MonadLiftT BaseIO m] (cls : Name) (k : m (α × MessageData)) (collapsed := true) : m α :=
+def withTraceNode' [MonadAlwaysExcept Exception m] [MonadLiftT BaseIO m] (cls : Name)
+    (k : m (α × MessageData)) (collapsed := true) : m α :=
   let msg := fun
     | .ok (_, msg) => return msg
     | .error err => return err.toMessageData
@@ -265,7 +292,10 @@ the result produced by `k` into an emoji (e.g., `💥`, `✅`, `❌`).
 
 TODO: find better name for this function.
 -/
-def withTraceNodeBefore [MonadRef m] [AddMessageContext m] [MonadOptions m] [MonadExcept ε m] [MonadLiftT BaseIO m] [ExceptToEmoji ε α] (cls : Name) (msg : m MessageData) (k : m α) (collapsed := true) : m α := do
+def withTraceNodeBefore [MonadRef m] [AddMessageContext m] [MonadOptions m]
+    [always : MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] [ExceptToEmoji ε α] (cls : Name)
+    (msg : m MessageData) (k : m α) (collapsed := true) : m α := do
+  let _ := always.except
   let opts ← getOptions
   let clsEnabled ← isTracingEnabledFor cls
   unless clsEnabled || trace.profiler.get opts do

tests/lean/exceptionalTrace.lean

@@ -0,0 +1,7 @@
+import Lean.Elab.Term
+
+elab "rt" : term => Lean.throwMaxRecDepthAt .missing
+
+set_option trace.Elab.step true
+
+#check rt + 1

tests/lean/exceptionalTrace.lean.expected.out

@@ -0,0 +1,6 @@
+exceptionalTrace.lean:7:0-7:13: error: maximum recursion depth has been reached (use `set_option maxRecDepth <num>` to increase limit)
+[Elab.step] expected type: <not-available>, term
+    rt + 1
+  [Elab.step] expected type: <not-available>, term
+      binop% HAdd.hAdd✝ rt 1
+    [Elab.step] expected type: <not-available>, term rt