Collapsible traces with messages

src/Lean/Elab/Attributes.lean

@@ -41,7 +41,7 @@ def toAttributeKind (attrKindStx : Syntax) : MacroM AttributeKind := do
 def mkAttrKindGlobal : Syntax :=
   mkNode ``Lean.Parser.Term.attrKind #[mkNullNode]
 
-def elabAttr [Monad m] [MonadEnv m] [MonadResolveName m] [MonadError m] [MonadMacroAdapter m] [MonadRecDepth m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadInfoTree m] (attrInstance : Syntax) : m Attribute := do
+def elabAttr [Monad m] [MonadEnv m] [MonadResolveName m] [MonadError m] [MonadMacroAdapter m] [MonadRecDepth m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadInfoTree m] [MonadLiftT IO m] (attrInstance : Syntax) : m Attribute := do
   /- attrInstance     := ppGroup $ leading_parser attrKind >> attrParser -/
   let attrKind ← liftMacroM <| toAttributeKind attrInstance[0]
   let attr := attrInstance[1]

src/Lean/Elab/Command.lean

@@ -128,12 +128,19 @@ private def mkCoreContext (ctx : Context) (s : State) (heartbeats : Nat) : Core.
     initHeartbeats := heartbeats
     currMacroScope := ctx.currMacroScope }
 
-private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceState : TraceState) : MessageLog :=
-  traceState.traces.foldl (init := log) fun (log : MessageLog) traceElem =>
+private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceState : TraceState) : MessageLog := Id.run do
+  if traceState.traces.isEmpty then return log
+  let mut traces : Std.HashMap (String.Pos × String.Pos) (Array MessageData) := ∅
+  for traceElem in traceState.traces do
     let ref := replaceRef traceElem.ref ctx.ref
     let pos := ref.getPos?.getD 0
     let endPos := ref.getTailPos?.getD pos
-    log.add (mkMessageCore ctx.fileName ctx.fileMap traceElem.msg MessageSeverity.information pos endPos)
+    traces := traces.insert (pos, endPos) <| traces.findD (pos, endPos) #[] |>.push traceElem.msg
+  let mut log := log
+  let traces' := traces.toArray.qsort fun ((a, _), _) ((b, _), _) => a < b
+  for ((pos, endPos), traceMsg) in traces' do
+    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap (.joinSep traceMsg.toList "\n") .information pos endPos
+  return log
 
 private def addTraceAsMessages : CommandElabM Unit := do
   let ctx ← read

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

src/Lean/Elab/Util.lean

@@ -158,7 +158,7 @@ instance (m n) [MonadLift m n] [MonadMacroAdapter m] : MonadMacroAdapter n := {
   setNextMacroScope := fun s => liftM (MonadMacroAdapter.setNextMacroScope s : m _)
 }
 
-def liftMacroM {α} {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m] [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] (x : MacroM α) : m α := do
+def liftMacroM {α} {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m] [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadLiftT IO m] (x : MacroM α) : m α := do
   let env  ← getEnv
   let currNamespace ← getCurrNamespace
   let openDecls ← getOpenDecls
@@ -192,7 +192,7 @@ def liftMacroM {α} {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEn
     s.traceMsgs.reverse.forM fun (clsName, msg) => trace clsName fun _ => msg
     return a
 
-@[inline] def adaptMacro {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m]  [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] (x : Macro) (stx : Syntax) : m Syntax :=
+@[inline] def adaptMacro {m : Type → Type} [Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m]  [MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadLiftT IO m] (x : Macro) (stx : Syntax) : m Syntax :=
   liftMacroM (x stx)
 
 partial def mkUnusedBaseName (baseName : Name) : MacroM Name := do
@@ -220,13 +220,6 @@ def withLogging [Monad m] [MonadLog m] [MonadExcept Exception m] [AddMessageCont
     (x : m Unit) : m Unit := do
   try x catch ex => logException ex
 
-@[inline] def trace [Monad m] [MonadLog m] [AddMessageContext m] [MonadOptions m] (cls : Name) (msg : Unit → MessageData) : m Unit := do
-  if checkTraceOption (← getOptions) cls then
-    logTrace cls (msg ())
-
-def logDbgTrace [Monad m] [MonadLog m] [AddMessageContext m] [MonadOptions m] (msg : MessageData) : m Unit := do
-  trace `Elab.debug fun _ => msg
-
 def nestedExceptionToMessageData [Monad m] [MonadLog m] (ex : Exception) : m MessageData := do
   let pos ← getRefPos
   match ex.getRef.getPos? with
@@ -244,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

src/Lean/Log.lean

@@ -95,10 +95,6 @@ def logWarning [MonadOptions m] (msgData : MessageData) : m Unit := do
 def logInfo (msgData : MessageData) : m Unit :=
   log msgData MessageSeverity.information
 
-/-- Log a trace message into the message log. -/
-def logTrace (cls : Name) (msgData : MessageData) : m Unit := do
-  logInfo (MessageData.tagged cls m!"[{cls}] {msgData}")
-
 /-- Log the error message "unknown declaration" -/
 def logUnknownDecl (declName : Name) : m Unit :=
   logError m!"unknown declaration '{declName}'"

src/Lean/Message.lean

@@ -54,7 +54,8 @@ inductive MessageData where
   /- Tagged sections. `Name` should be viewed as a "kind", and is used by `MessageData` inspector functions.
      Example: an inspector that tries to find "definitional equality failures" may look for the tag "DefEqFailure". -/
   | tagged            : Name → MessageData → MessageData
-  | node              : Array MessageData → MessageData
+  | trace (cls : Name) (msg : MessageData) (children : Array MessageData)
+    (collapsed : Bool := false)
   deriving Inhabited
 
 namespace MessageData
@@ -67,7 +68,6 @@ partial def isEmpty : MessageData → Bool
   | group m => m.isEmpty
   | compose m₁ m₂ => m₁.isEmpty && m₂.isEmpty
   | tagged _ m => m.isEmpty
-  | node ms => ms.all (·.isEmpty)
   | _ => false
 
 /- Instantiate metavariables occurring in nexted `ofExpr` constructors.
@@ -85,7 +85,7 @@ where
     | group msg                 => group <| visit msg mctx
     | compose msg₁ msg₂         => compose (visit msg₁ mctx) <| visit msg₂ mctx
     | tagged n msg              => tagged n <| visit msg mctx
-    | node msgs                 => node <| msgs.map (visit · mctx)
+    | trace cls msg msgs c      => trace cls (visit msg mctx) (msgs.map (visit · mctx)) c
     | _                         => msg
 
 variable (p : Name → Bool) in
@@ -96,7 +96,7 @@ partial def hasTag : MessageData → Bool
   | group msg               => hasTag msg
   | compose msg₁ msg₂       => hasTag msg₁ || hasTag msg₂
   | tagged n msg            => p n || hasTag msg
-  | node msgs               => msgs.any hasTag
+  | trace cls msg msgs _    => p cls || hasTag msg || msgs.any hasTag
   | _                       => false
 
 def nil : MessageData :=
@@ -127,18 +127,14 @@ partial def formatAux : NamingContext → Option MessageDataContext → MessageD
   | nCtx, some ctx,  ofGoal mvarId            => ppGoal (mkPPContext nCtx ctx) mvarId
   | nCtx, _,         withContext ctx d        => formatAux nCtx ctx d
   | _,    ctx,       withNamingContext nCtx d => formatAux nCtx ctx d
-  | nCtx, ctx,       tagged t d               =>
-    /- Messages starting a trace context have their tags postfixed with `_traceCtx` so that
-    we can detect them later. Here, we do so in order to print the trace context class. -/
-    if let .str cls "_traceCtx" := t then do
-      let d₁ ← formatAux nCtx ctx d
-      return f!"[{cls}] {d₁}"
-    else
-      formatAux nCtx ctx d
+  | nCtx, ctx,       tagged _ d               => formatAux nCtx ctx d
   | nCtx, ctx,       nest n d                 => Format.nest n <$> formatAux nCtx ctx d
   | nCtx, ctx,       compose d₁ d₂            => return (← formatAux nCtx ctx d₁) ++ (← formatAux nCtx ctx d₂)
   | nCtx, ctx,       group d                  => Format.group <$> formatAux nCtx ctx d
-  | nCtx, ctx,       node ds                  => Format.nest 2 <$> ds.foldlM (fun r d => return r ++ Format.line ++ (← formatAux nCtx ctx d)) Format.nil
+  | nCtx, ctx,       trace cls header children _ => do
+    let msg := f!"[{cls}] {(← formatAux nCtx ctx header).nest 2}"
+    let children ← children.mapM (formatAux nCtx ctx)
+    return .nest 2 (.joinSep (msg::children.toList) "\n")
 
 protected def format (msgData : MessageData) : IO Format :=
   formatAux { currNamespace := Name.anonymous, openDecls := [] } none msgData

src/Lean/Meta/AppBuilder.lean

@@ -254,6 +254,18 @@ private def mkFun (constName : Name) : MetaM (Expr × Expr) := do
   let fType ← instantiateTypeLevelParams cinfo us
   return (f, fType)
 
+private def withAppBuilderTrace [ToMessageData α] [ToMessageData β]
+    (f : α) (xs : β) (k : MetaM Expr) : MetaM Expr :=
+  let emoji | .ok .. => checkEmoji | .error .. => crossEmoji
+  withTraceNode `Meta.appBuilder (return m!"{emoji ·} f: {f}, xs: {xs}") do
+    try
+      let res ← k
+      trace[Meta.appBuilder.result] res
+      pure res
+    catch ex =>
+      trace[Meta.appBuilder.error] ex.toMessageData
+      throw ex
+
 /--
   Return the application `constName xs`.
   It tries to fill the implicit arguments before the last element in `xs`.
@@ -264,19 +276,15 @@ private def mkFun (constName : Name) : MetaM (Expr × Expr) := do
   Given a `x : (([Decidable p] → Bool) × Nat`, ``mkAppM `Prod.fst #[x]`` returns `@Prod.fst ([Decidable p] → Bool) Nat x`
 -/
 def mkAppM (constName : Name) (xs : Array Expr) : MetaM Expr := do
-  traceCtx `Meta.appBuilder <| withNewMCtxDepth do
+  withAppBuilderTrace constName xs do withNewMCtxDepth do
     let (f, fType) ← mkFun constName
-    let r ← mkAppMArgs f fType xs
-    trace[Meta.appBuilder] "constName: {constName}, xs: {xs}, result: {r}"
-    return r
+    mkAppMArgs f fType xs
 
 /-- Similar to `mkAppM`, but takes an `Expr` instead of a constant name. -/
 def mkAppM' (f : Expr) (xs : Array Expr) : MetaM Expr := do
   let fType ← inferType f
-  traceCtx `Meta.appBuilder <| withNewMCtxDepth do
-    let r ← mkAppMArgs f fType xs
-    trace[Meta.appBuilder] "f: {f}, xs: {xs}, result: {r}"
-    return r
+  withAppBuilderTrace f xs do withNewMCtxDepth do
+    mkAppMArgs f fType xs
 
 private partial def mkAppOptMAux (f : Expr) (xs : Array (Option Expr)) : Nat → Array Expr → Nat → Array MVarId → Expr → MetaM Expr
   | i, args, j, instMVars, Expr.forallE n d b bi => do
@@ -325,14 +333,14 @@ private partial def mkAppOptMAux (f : Expr) (xs : Array (Option Expr)) : Nat →
   fails because the only explicit argument `(a : α)` is not sufficient for inferring the remaining arguments,
   we would need the expected type. -/
 def mkAppOptM (constName : Name) (xs : Array (Option Expr)) : MetaM Expr := do
-  traceCtx `Meta.appBuilder <| withNewMCtxDepth do
+  withAppBuilderTrace constName xs do withNewMCtxDepth do
     let (f, fType) ← mkFun constName
     mkAppOptMAux f xs 0 #[] 0 #[] fType
 
 /-- Similar to `mkAppOptM`, but takes an `Expr` instead of a constant name -/
 def mkAppOptM' (f : Expr) (xs : Array (Option Expr)) : MetaM Expr := do
   let fType ← inferType f
-  traceCtx `Meta.appBuilder <| withNewMCtxDepth do
+  withAppBuilderTrace f xs do withNewMCtxDepth do
     mkAppOptMAux f xs 0 #[] 0 #[] fType
 
 def mkEqNDRec (motive h1 h2 : Expr) : MetaM Expr := do
@@ -581,6 +589,9 @@ def mkLE (a b : Expr) : MetaM Expr := mkBinaryRel ``LE ``LE.le a b
 /-- Return `a < b`. This method assumes `a` and `b` have the same type. -/
 def mkLT (a b : Expr) : MetaM Expr := mkBinaryRel ``LT ``LT.lt a b
 
-builtin_initialize registerTraceClass `Meta.appBuilder
+builtin_initialize do
+  registerTraceClass `Meta.appBuilder
+  registerTraceClass `Meta.appBuilder.result (inherited := true)
+  registerTraceClass `Meta.appBuilder.error (inherited := true)
 
 end Lean.Meta

src/Lean/Meta/Basic.lean

@@ -1577,29 +1577,29 @@ def mkLevelStuckErrorMessage (entry : PostponedEntry) : MetaM MessageData := do
 def mkLevelErrorMessage (entry : PostponedEntry) : MetaM MessageData := do
   mkLeveErrorMessageCore "failed to solve universe constraint" entry
 
-private def processPostponedStep (exceptionOnFailure : Bool) : MetaM Bool :=
-  traceCtx `Meta.isLevelDefEq.postponed.step do
-    let ps ← getResetPostponed
-    for p in ps do
-      unless (← withReader (fun ctx => { ctx with defEqCtx? := p.ctx? }) <| isLevelDefEqAux p.lhs p.rhs) do
-        if exceptionOnFailure then
-          withRef p.ref do
-            throwError (← mkLevelErrorMessage p)
-        else
-          return false
-    return true
+private def processPostponedStep (exceptionOnFailure : Bool) : MetaM Bool := do
+  let ps ← getResetPostponed
+  for p in ps do
+    unless (← withReader (fun ctx => { ctx with defEqCtx? := p.ctx? }) <| isLevelDefEqAux p.lhs p.rhs) do
+      if exceptionOnFailure then
+        withRef p.ref do
+          throwError (← mkLevelErrorMessage p)
+      else
+        return false
+  return true
 
 partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure := false) : MetaM Bool := do
   if (← getNumPostponed) == 0 then
     return true
   else
-    traceCtx `Meta.isLevelDefEq.postponed do
+    let numPostponedBegin ← getNumPostponed
+    withTraceNode `Meta.isLevelDefEq.postponed
+        (fun _ => return m!"processing #{numPostponedBegin} postponed is-def-eq level constraints") do
       let rec loop : MetaM Bool := do
         let numPostponed ← getNumPostponed
         if numPostponed == 0 then
           return true
         else
-          trace[Meta.isLevelDefEq.postponed] "processing #{numPostponed} postponed is-def-eq level constraints"
           if !(← processPostponedStep exceptionOnFailure) then
             return false
           else
@@ -1654,17 +1654,12 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
   Determines whether two universe level expressions are definitionally equal to each other.
 -/
 def isLevelDefEq (u v : Level) : MetaM Bool :=
-  traceCtx `Meta.isLevelDefEq do
-    let b ← checkpointDefEq (mayPostpone := true) <| Meta.isLevelDefEqAux u v
-    trace[Meta.isLevelDefEq] "{u} =?= {v} ... {if b then "success" else "failure"}"
-    return b
+  checkpointDefEq (mayPostpone := true) <| Meta.isLevelDefEqAux u v
 
 /-- See `isDefEq`. -/
 def isExprDefEq (t s : Expr) : MetaM Bool :=
-  traceCtx `Meta.isDefEq <| withReader (fun ctx => { ctx with defEqCtx? := some { lhs := t, rhs := s, lctx := ctx.lctx, localInstances := ctx.localInstances } }) do
-    let b ← checkpointDefEq (mayPostpone := true) <| Meta.isExprDefEqAux t s
-    trace[Meta.isDefEq] "{t} =?= {s} ... {if b then "success" else "failure"}"
-    return b
+  withReader (fun ctx => { ctx with defEqCtx? := some { lhs := t, rhs := s, lctx := ctx.lctx, localInstances := ctx.localInstances } }) do
+    checkpointDefEq (mayPostpone := true) <| Meta.isExprDefEqAux t s
 
 /--
   Determines whether two expressions are definitionally equal to each other.

src/Lean/Meta/Check.lean

@@ -187,8 +187,13 @@ where
 Throw an exception if `e` is not type correct.
 -/
 def check (e : Expr) : MetaM Unit :=
-  traceCtx `Meta.check do
-    withTransparency TransparencyMode.all $ checkAux e
+  withTraceNode `Meta.check (fun res =>
+      return m!"{if res.isOk then checkEmoji else crossEmoji} {e}") do
+    try
+      withTransparency TransparencyMode.all $ checkAux e
+    catch ex =>
+      trace[Meta.check] ex.toMessageData
+      throw ex
 
 /--
 Return true if `e` is type correct.
@@ -197,12 +202,10 @@ def isTypeCorrect (e : Expr) : MetaM Bool := do
   try
     check e
     pure true
-  catch ex =>
-    trace[Meta.typeError] ex.toMessageData
+  catch _ =>
     pure false
 
 builtin_initialize
   registerTraceClass `Meta.check
-  registerTraceClass `Meta.typeError
 
 end Lean.Meta