From 5f1c4df07db4adcf4a72f9f7bf00c2a0a6939a61 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 2 May 2024 00:47:57 +0200
Subject: [PATCH] feat: display diagnostic information at term and tactic
 `set_option diagnostics true` (#4048)

We don't need to include reduction info at `simp` diagnostic
information.
---
 src/Lean/Compiler/ExternAttr.lean          |  2 +-
 src/Lean/CoreM.lean                        | 15 ++++--
 src/Lean/Elab/BuiltinTerm.lean             |  9 +++-
 src/Lean/Elab/InfoTree/Main.lean           |  7 +--
 src/Lean/Elab/Tactic/BuiltinTactic.lean    |  9 +++-
 src/Lean/Elab/Tactic/Simp.lean             |  4 +-
 src/Lean/Meta/Tactic/Simp/Diagnostics.lean | 11 +----
 src/Lean/Meta/Tactic/Simp/Main.lean        |  6 +--
 src/Lean/PrettyPrinter.lean                |  6 +--
 src/Lean/Replay.lean                       |  2 +-
 src/lake/Lake/CLI/Translate.lean           |  2 +-
 src/lake/Lake/Toml/Load.lean               |  2 +-
 tests/lean/run/diagnostics.lean            | 53 ++++++++++++++++++++++
 tests/lean/run/simpDiag.lean               | 12 +++--
 14 files changed, 103 insertions(+), 37 deletions(-)
 create mode 100644 tests/lean/run/diagnostics.lean

diff --git a/src/Lean/Compiler/ExternAttr.lean b/src/Lean/Compiler/ExternAttr.lean
index 4f6b5efaefc9..b7eb708e3566 100644
--- a/src/Lean/Compiler/ExternAttr.lean
+++ b/src/Lean/Compiler/ExternAttr.lean
@@ -156,7 +156,7 @@ private def getExternConstArity (declName : Name) : CoreM Nat := do
 @[export lean_get_extern_const_arity]
 def getExternConstArityExport (env : Environment) (declName : Name) : IO (Option Nat) := do
   try
-    let (arity, _) ← (getExternConstArity declName).toIO { fileName := "<compiler>", fileMap := default, diag := false } { env := env }
+    let (arity, _) ← (getExternConstArity declName).toIO { fileName := "<compiler>", fileMap := default } { env := env }
     return some arity
   catch
    | IO.Error.userError _   => return none
diff --git a/src/Lean/CoreM.lean b/src/Lean/CoreM.lean
index 0f0e2a10c1e5..76f9c36315f7 100644
--- a/src/Lean/CoreM.lean
+++ b/src/Lean/CoreM.lean
@@ -121,7 +121,15 @@ instance : MonadOptions CoreM where
   getOptions := return (← read).options
 
 instance : MonadWithOptions CoreM where
-  withOptions f x := withReader (fun ctx => { ctx with options := f ctx.options }) x
+  withOptions f x :=
+    withReader
+      (fun ctx =>
+        let options := f ctx.options
+        { ctx with
+          options
+          diag := diagnostics.get options
+          maxRecDepth := maxRecDepth.get options })
+      x
 
 instance : AddMessageContext CoreM where
   addMessageContext := addMessageContextPartial
@@ -223,7 +231,7 @@ def mkFreshUserName (n : Name) : CoreM Name :=
 instance [MetaEval α] : MetaEval (CoreM α) where
   eval env opts x _ := do
     let x : CoreM α := do try x finally printTraces
-    let (a, s) ← x.toIO { maxRecDepth := maxRecDepth.get opts, options := opts, fileName := "<CoreM>", fileMap := default, diag := diagnostics.get opts } { env := env }
+    let (a, s) ← (withOptions (fun _ => opts) x).toIO { fileName := "<CoreM>", fileMap := default } { env := env }
     MetaEval.eval s.env opts a (hideUnit := true)
 
 -- withIncRecDepth for a monad `m` such that `[MonadControlT CoreM n]`
@@ -398,8 +406,7 @@ def isDiagnosticsEnabled : CoreM Bool :=
 
 def ImportM.runCoreM (x : CoreM α) : ImportM α := do
   let ctx ← read
-  let opts := ctx.opts
-  let (a, _) ← x.toIO { options := opts, fileName := "<ImportM>", fileMap := default, diag := getDiag opts } { env := ctx.env }
+  let (a, _) ← (withOptions (fun _ => ctx.opts) x).toIO { fileName := "<ImportM>", fileMap := default } { env := ctx.env }
   return a
 
 /-- Return `true` if the exception was generated by one our resource limits. -/
diff --git a/src/Lean/Elab/BuiltinTerm.lean b/src/Lean/Elab/BuiltinTerm.lean
index 52b6654ae088..6e097ba8f272 100644
--- a/src/Lean/Elab/BuiltinTerm.lean
+++ b/src/Lean/Elab/BuiltinTerm.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Lean.Meta.Diagnostics
 import Lean.Elab.Open
 import Lean.Elab.SetOption
 import Lean.Elab.Eval
@@ -313,8 +314,12 @@ private def mkSilentAnnotationIfHole (e : Expr) : TermElabM Expr := do
 
 @[builtin_term_elab «set_option»] def elabSetOption : TermElab := fun stx expectedType? => do
   let options ← Elab.elabSetOption stx[1] stx[3]
-  withTheReader Core.Context (fun ctx => { ctx with maxRecDepth := maxRecDepth.get options, options := options, diag := getDiag options }) do
-    elabTerm stx[5] expectedType?
+  withOptions (fun _ => options) do
+    try
+      elabTerm stx[5] expectedType?
+    finally
+      if stx[1].getId == `diagnostics then
+        reportDiag
 
 @[builtin_term_elab withAnnotateTerm] def elabWithAnnotateTerm : TermElab := fun stx expectedType? => do
   match stx with
diff --git a/src/Lean/Elab/InfoTree/Main.lean b/src/Lean/Elab/InfoTree/Main.lean
index fafda87b125f..813ffa3185ce 100644
--- a/src/Lean/Elab/InfoTree/Main.lean
+++ b/src/Lean/Elab/InfoTree/Main.lean
@@ -102,9 +102,10 @@ def ContextInfo.runCoreM (info : ContextInfo) (x : CoreM α) : IO α := do
     have been used in `lctx` and `info.mctx`.
   -/
   (·.1) <$>
-    x.toIO { options := info.options, currNamespace := info.currNamespace, openDecls := info.openDecls
-             fileName := "<InfoTree>", fileMap := default, diag := getDiag info.options }
-           { env := info.env, ngen := info.ngen }
+    (withOptions (fun _ => info.options) x).toIO
+      { currNamespace := info.currNamespace, openDecls := info.openDecls
+        fileName := "<InfoTree>", fileMap := default }
+      { env := info.env, ngen := info.ngen }
 
 def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do
   (·.1) <$> info.runCoreM (x.run { lctx := lctx } { mctx := info.mctx })
diff --git a/src/Lean/Elab/Tactic/BuiltinTactic.lean b/src/Lean/Elab/Tactic/BuiltinTactic.lean
index 41bffa6677f2..ea40cb6f0fd3 100644
--- a/src/Lean/Elab/Tactic/BuiltinTactic.lean
+++ b/src/Lean/Elab/Tactic/BuiltinTactic.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Lean.Meta.Diagnostics
 import Lean.Meta.Tactic.Apply
 import Lean.Meta.Tactic.Assumption
 import Lean.Meta.Tactic.Contradiction
@@ -163,8 +164,12 @@ private def getOptRotation (stx : Syntax) : Nat :=
 
 @[builtin_tactic Parser.Tactic.set_option] def elabSetOption : Tactic := fun stx => do
   let options ← Elab.elabSetOption stx[1] stx[3]
-  withTheReader Core.Context (fun ctx => { ctx with maxRecDepth := maxRecDepth.get options, options := options, diag := getDiag options }) do
-    evalTactic stx[5]
+  withOptions (fun _ => options) do
+    try
+      evalTactic stx[5]
+    finally
+      if stx[1].getId == `diagnostics then
+        reportDiag
 
 @[builtin_tactic Parser.Tactic.allGoals] def evalAllGoals : Tactic := fun stx => do
   let mvarIds ← getGoals
diff --git a/src/Lean/Elab/Tactic/Simp.lean b/src/Lean/Elab/Tactic/Simp.lean
index 6477a18334a8..6c3a89d4ed27 100644
--- a/src/Lean/Elab/Tactic/Simp.lean
+++ b/src/Lean/Elab/Tactic/Simp.lean
@@ -414,10 +414,8 @@ where
     return stats
 
 def withSimpDiagnostics (x : TacticM Simp.Diagnostics) : TacticM Unit := do
-  let origDiag := (← getThe Meta.State).diag
   let stats ← x
-  Simp.reportDiag stats origDiag
-  return ()
+  Simp.reportDiag stats
 
 /-
   "simp" (config)? (discharger)? (" only")? (" [" ((simpStar <|> simpErase <|> simpLemma),*,?) "]")?
diff --git a/src/Lean/Meta/Tactic/Simp/Diagnostics.lean b/src/Lean/Meta/Tactic/Simp/Diagnostics.lean
index 766c8b8fc2ce..df2ff3261d8a 100644
--- a/src/Lean/Meta/Tactic/Simp/Diagnostics.lean
+++ b/src/Lean/Meta/Tactic/Simp/Diagnostics.lean
@@ -32,23 +32,16 @@ def mkSimpDiagSummary (counters : PHashMap Origin Nat) (usedCounters? : Option (
       data := data.push m!"{if data.isEmpty then "  " else "\n"}{key} ↦ {counter}{usedMsg}"
     return { data, max := entries[0]!.2 }
 
-def reportDiag (diag : Simp.Diagnostics) (diagOrig : Meta.Diagnostics) : MetaM Unit := do
+def reportDiag (diag : Simp.Diagnostics) : MetaM Unit := do
   if (← isDiagnosticsEnabled) then
     let used ← mkSimpDiagSummary diag.usedThmCounter
     let tried ← mkSimpDiagSummary diag.triedThmCounter diag.usedThmCounter
     let congr ← mkDiagSummary diag.congrThmCounter
-    let unfoldCounter := subCounters (← get).diag.unfoldCounter diagOrig.unfoldCounter
-    let unfoldDefault ← mkDiagSummaryForUnfolded unfoldCounter
-    let unfoldInstance ← mkDiagSummaryForUnfolded unfoldCounter (instances := true)
-    let unfoldReducible ← mkDiagSummaryForUnfoldedReducible unfoldCounter
-    unless used.isEmpty && tried.isEmpty && unfoldDefault.isEmpty && unfoldInstance.isEmpty && unfoldReducible.isEmpty do
+    unless used.isEmpty && tried.isEmpty && congr.isEmpty do
       let m := MessageData.nil
       let m := appendSection m `simp "used theorems" used
       let m := appendSection m `simp "tried theorems" tried
       let m := appendSection m `simp "tried congruence theorems" congr
-      let m := appendSection m `reduction "unfolded declarations" unfoldDefault
-      let m := appendSection m `reduction "unfolded instances" unfoldInstance
-      let m := appendSection m `reduction "unfolded reducible declarations" unfoldReducible
       let m := m ++ "use `set_option diagnostics.threshold <num>` to control threshold for reporting counters"
       logInfo m
 
diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean
index 99d5d5ba4239..a553e22d86c6 100644
--- a/src/Lean/Meta/Tactic/Simp/Main.lean
+++ b/src/Lean/Meta/Tactic/Simp/Main.lean
@@ -663,11 +663,10 @@ def main (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods :=
     return (r, { s with })
 where
   simpMain (e : Expr) : SimpM Result := withCatchingRuntimeEx do
-    let origDiag := (← getThe Meta.State).diag
     try
       withoutCatchingRuntimeEx <| simp e
     catch ex =>
-      reportDiag (← get).diag origDiag
+      reportDiag (← get).diag
       if ex.isRuntime then
         throwNestedTacticEx `simp ex
       else
@@ -679,11 +678,10 @@ def dsimpMain (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Method
     pure (r, { s with })
 where
   dsimpMain (e : Expr) : SimpM Expr := withCatchingRuntimeEx do
-    let origDiag := (← getThe Meta.State).diag
     try
       withoutCatchingRuntimeEx <| dsimp e
     catch ex =>
-      reportDiag (← get).diag origDiag
+      reportDiag (← get).diag
       if ex.isRuntime then
         throwNestedTacticEx `simp ex
       else
diff --git a/src/Lean/PrettyPrinter.lean b/src/Lean/PrettyPrinter.lean
index 561b7b8d1ecd..7ea1f5d428c1 100644
--- a/src/Lean/PrettyPrinter.lean
+++ b/src/Lean/PrettyPrinter.lean
@@ -51,9 +51,9 @@ def ppExprWithInfos (e : Expr) (optsPerPos : Delaborator.OptionsPerPos := {}) (d
 
 @[export lean_pp_expr]
 def ppExprLegacy (env : Environment) (mctx : MetavarContext) (lctx : LocalContext) (opts : Options) (e : Expr) : IO Format :=
-  Prod.fst <$> ((ppExpr e).run' { lctx := lctx } { mctx := mctx }).toIO {
-    options := opts, fileName := "<PrettyPrinter>", fileMap := default, diag := getDiag opts
-  } { env := env }
+  Prod.fst <$> ((withOptions (fun _ => opts) <| ppExpr e).run' { lctx := lctx } { mctx := mctx }).toIO
+    { fileName := "<PrettyPrinter>", fileMap := default }
+    { env := env }
 
 def ppTactic (stx : TSyntax `tactic) : CoreM Format := ppCategory `tactic stx
 
diff --git a/src/Lean/Replay.lean b/src/Lean/Replay.lean
index a05698de6bd6..d44392191ec6 100644
--- a/src/Lean/Replay.lean
+++ b/src/Lean/Replay.lean
@@ -50,7 +50,7 @@ def isTodo (name : Name) : M Bool := do
 
 /-- Use the current `Environment` to throw a `KernelException`. -/
 def throwKernelException (ex : KernelException) : M Unit := do
-    let ctx := { fileName := "", options := ({} : KVMap), fileMap := default, diag := false }
+    let ctx := { fileName := "", options := ({} : KVMap), fileMap := default }
     let state := { env := (← get).env }
     Prod.fst <$> (Lean.Core.CoreM.toIO · ctx state) do Lean.throwKernelException ex
 
diff --git a/src/lake/Lake/CLI/Translate.lean b/src/lake/Lake/CLI/Translate.lean
index f30be562f013..41175786bf35 100644
--- a/src/lake/Lake/CLI/Translate.lean
+++ b/src/lake/Lake/CLI/Translate.lean
@@ -28,7 +28,7 @@ def Package.mkConfigString (pkg : Package) (lang : ConfigLang) : LogIO String :=
   | .lean => do
     let env ← importModulesUsingCache #[`Lake] {} 1024
     let pp := ppModule <| descopeTSyntax <| pkg.mkLeanConfig
-    match (← pp.toIO {fileName := "", fileMap := default, diag := false} {env} |>.toBaseIO) with
+    match (← pp.toIO {fileName := "", fileMap := default} {env} |>.toBaseIO) with
     | .ok (fmt, _) => pure <| (toString fmt).trim ++ "\n"
     | .error ex =>
       error s!"(internal) failed to pretty print Lean configuration: {ex.toString}"
diff --git a/src/lake/Lake/Toml/Load.lean b/src/lake/Lake/Toml/Load.lean
index c7201ea29865..919ff4e6941b 100644
--- a/src/lake/Lake/Toml/Load.lean
+++ b/src/lake/Lake/Toml/Load.lean
@@ -22,7 +22,7 @@ def loadToml (ictx : InputContext) : EIO MessageLog Table := do
     throw <|  MessageLog.empty.add <| mkParserErrorMessage ictx s errorMsg
   else if ictx.input.atEnd s.pos then
     let act := elabToml ⟨s.stxStack.back⟩
-    match (← act.run {fileName := ictx.fileName, fileMap := ictx.fileMap, diag := false} {env} |>.toBaseIO) with
+    match (← act.run {fileName := ictx.fileName, fileMap := ictx.fileMap} {env} |>.toBaseIO) with
     | .ok (t, s) =>
       if s.messages.hasErrors then
         throw s.messages
diff --git a/tests/lean/run/diagnostics.lean b/tests/lean/run/diagnostics.lean
new file mode 100644
index 000000000000..abdc261f3c53
--- /dev/null
+++ b/tests/lean/run/diagnostics.lean
@@ -0,0 +1,53 @@
+opaque q : Nat → Nat
+def f (x : Nat) : Nat :=
+  match x with
+  | 0 => 1
+  | x+1 => q (f x)
+
+theorem f_eq : f (x + 1) = q (f x) := rfl
+
+set_option trace.Meta.debug true
+
+/--
+info: [reduction] unfolded declarations (max: 15, num: 6):
+    Nat.rec ↦ 15
+  ⏎
+  Add.add ↦ 10
+  ⏎
+  HAdd.hAdd ↦ 10
+  ⏎
+  Nat.add ↦ 10
+  ⏎
+  f ↦ 5
+  ⏎
+  OfNat.ofNat ↦ 5[reduction] unfolded instances (max: 5, num: 1):
+    instOfNatNat ↦ 5[reduction] unfolded reducible declarations (max: 15, num: 1):
+    Nat.casesOn ↦ 15use `set_option diagnostics.threshold <num>` to control threshold for reporting counters
+-/
+#guard_msgs in
+example : f (x + 5) = q (q (q (q (q (f x))))) :=
+  set_option diagnostics.threshold 4 in
+  set_option diagnostics true in
+  rfl
+
+/--
+info: [reduction] unfolded declarations (max: 15, num: 6):
+    Nat.rec ↦ 15
+  ⏎
+  Add.add ↦ 10
+  ⏎
+  HAdd.hAdd ↦ 10
+  ⏎
+  Nat.add ↦ 10
+  ⏎
+  f ↦ 5
+  ⏎
+  OfNat.ofNat ↦ 5[reduction] unfolded instances (max: 5, num: 1):
+    instOfNatNat ↦ 5[reduction] unfolded reducible declarations (max: 15, num: 1):
+    Nat.casesOn ↦ 15use `set_option diagnostics.threshold <num>` to control threshold for reporting counters
+-/
+#guard_msgs in
+example : f (x + 5) = q (q (q (q (q (f x))))) := by
+  set_option diagnostics.threshold 4 in
+  set_option diagnostics true in
+  rfl
diff --git a/tests/lean/run/simpDiag.lean b/tests/lean/run/simpDiag.lean
index 1fb98d33c0b7..4c12eeeaf7c3 100644
--- a/tests/lean/run/simpDiag.lean
+++ b/tests/lean/run/simpDiag.lean
@@ -98,11 +98,12 @@ opaque q1 : Nat → Nat → Prop
 /--
 info: [simp] used theorems (max: 1, num: 1):
     q1_ax ↦ 1[simp] tried theorems (max: 1, num: 1):
-    q1_ax ↦ 1, succeeded: 1[reduction] unfolded declarations (max: 246, num: 2):
+    q1_ax ↦ 1, succeeded: 1use `set_option diagnostics.threshold <num>` to control threshold for reporting counters
+---
+info: [reduction] unfolded declarations (max: 246, num: 2):
     Nat.rec ↦ 246
   ⏎
-  OfNat.ofNat ↦ 24[reduction] unfolded instances (max: 12, num: 1):
-    instOfNatNat ↦ 12[reduction] unfolded reducible declarations (max: 246, num: 2):
+  OfNat.ofNat ↦ 24[reduction] unfolded reducible declarations (max: 246, num: 2):
     h ↦ 246
    Nat.casesOn ↦ 246use `set_option diagnostics.threshold <num>` to control threshold for reporting counters
 -/
@@ -111,3 +112,8 @@ example : q1 x (h 40) := by
   set_option diagnostics true in
   set_option diagnostics.threshold 0 in
   simp
+
+example : q1 x (h 40) := by
+  set_option diagnostics true in
+  set_option diagnostics.threshold 0 in
+  simp