From f89f6cb56cb61da3952cef8597d56bc5e66b375f Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Sun, 7 Aug 2022 14:01:43 +0200
Subject: [PATCH] chore: improve elab traces

---
 src/Lean/Elab/Term.lean | 7 ++++---
 src/Lean/Elab/Util.lean | 1 +
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Elab/Term.lean b/src/Lean/Elab/Term.lean
index 3e2120893d..1c2042c9fb 100644
--- a/src/Lean/Elab/Term.lean
+++ b/src/Lean/Elab/Term.lean
@@ -1329,11 +1329,10 @@ where
 private partial def elabTermAux (expectedType? : Option Expr) (catchExPostpone : Bool) (implicitLambda : Bool) : Syntax → TermElabM Expr
   | .missing => mkSyntheticSorryFor expectedType?
   | stx => withFreshMacroScope <| withIncRecDepth do
-    trace[Elab.step] "expected type: {expectedType?}, term\n{stx}"
+    withTraceNode `Elab.step (fun _ => return m!"expected type: {expectedType?}, term\n{stx}") do
     checkMaxHeartbeats "elaborator"
-    withNestedTraces do
     let env ← getEnv
-    match (← liftMacroM (expandMacroImpl? env stx)) with
+    let result ← match (← liftMacroM (expandMacroImpl? env stx)) with
     | some (decl, stxNew?) =>
       let stxNew ← liftMacroM <| liftExcept stxNew?
       withInfoContext' stx (mkInfo := mkTermInfo decl (expectedType? := expectedType?) stx) <|
@@ -1345,6 +1344,8 @@ private partial def elabTermAux (expectedType? : Option Expr) (catchExPostpone :
       match implicit? with
       | some expectedType => elabImplicitLambda stx catchExPostpone expectedType
       | none              => elabUsingElabFns stx expectedType? catchExPostpone
+    trace[Elab.step.result] result
+    pure result
 
 /-- Store in the `InfoTree` that `e` is a "dot"-completion target. -/
 def addDotCompletionInfo (stx : Syntax) (e : Expr) (expectedType? : Option Expr) (field? : Option Syntax := none) : TermElabM Unit := do
diff --git a/src/Lean/Elab/Util.lean b/src/Lean/Elab/Util.lean
index 6a80103abf..c78e9d2cc6 100644
--- a/src/Lean/Elab/Util.lean
+++ b/src/Lean/Elab/Util.lean
@@ -237,5 +237,6 @@ def throwErrorWithNestedErrors [MonadError m] [Monad m] [MonadLog m] (msg : Mess
 builtin_initialize
   registerTraceClass `Elab
   registerTraceClass `Elab.step
+  registerTraceClass `Elab.step.result (inherited := true)
 
 end Lean.Elab