feat: add set_option diag true for diagnostic counters

src/Lean/Compiler/ExternAttr.lean

@@ -155,7 +155,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 } { env := env }
+    let (arity, _) ← (getExternConstArity declName).toIO { fileName := "<compiler>", fileMap := default, diag := false } { env := env }
     return some arity
   catch
    | IO.Error.userError _   => return none

src/Lean/CoreM.lean

@@ -15,6 +15,12 @@ import Lean.MonadEnv
 namespace Lean
 namespace Core
 
+register_builtin_option diag : Bool := {
+  defValue := false
+  group    := "diagnostics"
+  descr    := "collect diagnostic information"
+}
+
 register_builtin_option maxHeartbeats : Nat := {
   defValue := 200000
   descr := "maximum amount of heartbeats per command. A heartbeat is number of (small) memory allocations (in thousands), 0 means no limit"
@@ -72,6 +78,11 @@ structure Context where
   Recall that runtime exceptions are `maxRecDepth` or `maxHeartbeats`.
   -/
   catchRuntimeEx : Bool := false
+  /--
+  If `diag := true`, different parts of the system collect diagnostics.
+  Use the `set_option diag true` to set it to true.
+  -/
+  diag           : Bool := false
   deriving Nonempty
 
 /-- CoreM is a monad for manipulating the Lean environment.
@@ -206,7 +217,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 } { env := env }
+    let (a, s) ← x.toIO { maxRecDepth := maxRecDepth.get opts, options := opts, fileName := "<CoreM>", fileMap := default, diag := diag.get opts } { env := env }
     MetaEval.eval s.env opts a (hideUnit := true)
 
 -- withIncRecDepth for a monad `m` such that `[MonadControlT CoreM n]`
@@ -372,9 +383,17 @@ def addAndCompile (decl : Declaration) : CoreM Unit := do
   addDecl decl;
   compileDecl decl
 
+def getDiag (opts : Options) : Bool :=
+  Core.diag.get opts
+
+/-- Return `true` if diagnostic information collection is enabled. -/
+def isDiagnosticsEnabled : CoreM Bool :=
+  return (← read).diag
+
 def ImportM.runCoreM (x : CoreM α) : ImportM α := do
   let ctx ← read
-  let (a, _) ← x.toIO { options := ctx.opts, fileName := "<ImportM>", fileMap := default } { env := ctx.env }
+  let opts := ctx.opts
+  let (a, _) ← x.toIO { options := opts, fileName := "<ImportM>", fileMap := default, diag := getDiag opts } { env := ctx.env }
   return a
 
 /-- Return `true` if the exception was generated by one our resource limits. -/

src/Lean/Elab/BuiltinTerm.lean

@@ -313,7 +313,7 @@ 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 }) do
+  withTheReader Core.Context (fun ctx => { ctx with maxRecDepth := maxRecDepth.get options, options := options, diag := getDiag options }) do
     elabTerm stx[5] expectedType?
 
 @[builtin_term_elab withAnnotateTerm] def elabWithAnnotateTerm : TermElab := fun stx expectedType? => do

src/Lean/Elab/Command.lean

@@ -138,7 +138,8 @@ private def mkCoreContext (ctx : Context) (s : State) (heartbeats : Nat) : Core.
     currNamespace  := scope.currNamespace
     openDecls      := scope.openDecls
     initHeartbeats := heartbeats
-    currMacroScope := ctx.currMacroScope }
+    currMacroScope := ctx.currMacroScope
+    diag           := getDiag scope.opts }
 
 private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceState : TraceState) : MessageLog := Id.run do
   if traceState.traces.isEmpty then return log
@@ -411,7 +412,7 @@ def liftTermElabM (x : TermElabM α) : CommandElabM α := do
   -- 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 : MetaM _      := (observing (try x finally Meta.reportDiag)).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 }
   let (((ea, _), _), coreS) ← liftEIO x

src/Lean/Elab/InfoTree/Main.lean

@@ -102,7 +102,8 @@ 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 }
+    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 }
 
 def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do

src/Lean/Elab/Tactic/BuiltinTactic.lean

@@ -163,7 +163,7 @@ 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 }) do
+  withTheReader Core.Context (fun ctx => { ctx with maxRecDepth := maxRecDepth.get options, options := options, diag := getDiag options }) do
     evalTactic stx[5]
 
 @[builtin_tactic Parser.Tactic.allGoals] def evalAllGoals : Tactic := fun stx => do

src/Lean/Meta/Basic.lean

@@ -27,6 +27,12 @@ namespace Lean.Meta
 
 builtin_initialize isDefEqStuckExceptionId : InternalExceptionId ← registerInternalExceptionId `isDefEqStuck
 
+register_builtin_option diag.isDefEq.threshould : Nat := {
+  defValue := 20
+  group    := "diagnostics"
+  descr    := "only diagnostic counters above this threshold are reported by the definitional equality"
+}
+
 /--
 Configuration flags for the `MetaM` monad.
 Many of them are used to control the `isDefEq` function that checks whether two terms are definitionally equal or not.
@@ -266,6 +272,13 @@ structure PostponedEntry where
   ctx? : Option DefEqContext
   deriving Inhabited
 
+structure Diagnostics where
+  /-- Number of times each declaration has been unfolded -/
+  unfoldCounter : PHashMap Name Nat := {}
+  /-- Number of times `f a =?= f b` heuristic has been used per function `f`. -/
+  heuristicCounter : PHashMap Name Nat := {}
+  deriving Inhabited
+
 /--
   `MetaM` monad state.
 -/
@@ -276,6 +289,7 @@ structure State where
   zetaDeltaFVarIds : FVarIdSet := {}
   /-- Array of postponed universe level constraints -/
   postponed        : PersistentArray PostponedEntry := {}
+  diag             : Diagnostics := {}
   deriving Inhabited
 
 /--
@@ -440,7 +454,7 @@ section Methods
 variable [MonadControlT MetaM n] [Monad n]
 
 @[inline] def modifyCache (f : Cache → Cache) : MetaM Unit :=
-  modify fun { mctx, cache, zetaDeltaFVarIds, postponed } => { mctx, cache := f cache, zetaDeltaFVarIds, postponed }
+  modify fun { mctx, cache, zetaDeltaFVarIds, postponed, diag } => { mctx, cache := f cache, zetaDeltaFVarIds, postponed, diag }
 
 @[inline] def modifyInferTypeCache (f : InferTypeCache → InferTypeCache) : MetaM Unit :=
   modifyCache fun ⟨ic, c1, c2, c3, c4, c5, c6⟩ => ⟨f ic, c1, c2, c3, c4, c5, c6⟩
@@ -454,6 +468,55 @@ variable [MonadControlT MetaM n] [Monad n]
 @[inline] def resetDefEqPermCaches : MetaM Unit :=
   modifyDefEqPermCache fun _ => {}
 
+@[inline] def modifyDiag (f : Diagnostics → Diagnostics) : MetaM Unit := do
+  if (← isDiagnosticsEnabled) then
+    modify fun { mctx, cache, zetaDeltaFVarIds, postponed, diag } => { mctx, cache, zetaDeltaFVarIds, postponed, diag := f diag }
+
+/-- If diagnostics are enabled, record that `declName` has been unfolded. -/
+def recordUnfold (declName : Name) : MetaM Unit := do
+  modifyDiag fun { unfoldCounter, heuristicCounter } =>
+    let newC := if let some c := unfoldCounter.find? declName then c + 1 else 1
+    { unfoldCounter := unfoldCounter.insert declName newC, heuristicCounter }
+
+/-- If diagnostics are enabled, record that heuristic for solving `f a =?= f b` has been used. -/
+def recordDefEqHeuristic (declName : Name) : MetaM Unit := do
+  modifyDiag fun { unfoldCounter, heuristicCounter } =>
+    let newC := if let some c := heuristicCounter.find? declName then c + 1 else 1
+    { unfoldCounter, heuristicCounter := heuristicCounter.insert declName newC }
+
+def collectAboveThreshold (counters : PHashMap Name Nat) (threshold : Nat) : Array (Name × Nat) := Id.run do
+  let mut r := #[]
+  for (declName, counter) in counters do
+    if counter > threshold then
+      r := r.push (declName, counter)
+  return r.qsort fun (d₁, c₁) (d₂, c₂) => if c₁ == c₂ then d₁.lt d₂ else c₁ > c₂
+
+def mkMessageBodyFor? (counters : PHashMap Name Nat) (threshold : Nat) : Option MessageData := Id.run do
+  let entries := collectAboveThreshold counters threshold
+  if entries.isEmpty then
+    return none
+  else
+    let mut m := MessageData.nil
+    for (declName, counter) in entries do
+      m := m ++ m!"{declName} ↦ {counter}\n"
+    return some m
+
+def appendOptMessageData (m : MessageData) (header : String) (m? : Option MessageData) : MessageData :=
+  if let some m' := m? then
+    m ++ header ++ indentD m'
+  else
+    m
+
+def reportDiag : MetaM Unit := do
+  if (← isDiagnosticsEnabled) then
+    let threshould := diag.isDefEq.threshould.get (← getOptions)
+    let unfold? := mkMessageBodyFor? (← get).diag.unfoldCounter threshould
+    let heu?    := mkMessageBodyFor? (← get).diag.heuristicCounter threshould
+    if unfold?.isSome || heu?.isSome then
+      let m := appendOptMessageData MessageData.nil "unfolded declarations:" unfold?
+      let m := appendOptMessageData m "`isDefEq` heuristic:" heu?
+      logInfo m
+
 def getLocalInstances : MetaM LocalInstances :=
   return (← read).localInstances
 

src/Lean/Meta/ExprDefEq.lean

@@ -1251,6 +1251,7 @@ private def tryHeuristic (t s : Expr) : MetaM Bool := do
     unless t.hasExprMVar || s.hasExprMVar do
       return false
   withTraceNodeBefore `Meta.isDefEq.delta (return m!"{t} =?= {s}") do
+    recordDefEqHeuristic tFn.constName!
     /-
       We process arguments before universe levels to reduce a source of brittleness in the TC procedure.
 

src/Lean/Meta/LazyDiscrTree.lean

@@ -978,12 +978,13 @@ def createImportedDiscrTree [Monad m] [MonadLog m] [AddMessageContext m] [MonadO
 
 /-- Creates the core context used for initializing a tree using the current context. -/
 private def createTreeCtx (ctx : Core.Context) : Core.Context := {
-    fileName := ctx.fileName,
-    fileMap := ctx.fileMap,
-    options := ctx.options,
-    maxRecDepth := ctx.maxRecDepth,
-    maxHeartbeats := 0,
-    ref := ctx.ref,
+    fileName := ctx.fileName
+    fileMap := ctx.fileMap
+    options := ctx.options
+    maxRecDepth := ctx.maxRecDepth
+    maxHeartbeats := 0
+    ref := ctx.ref
+    diag := getDiag ctx.options
   }
 
 def findImportMatches

src/Lean/Meta/WHNF.lean

@@ -706,7 +706,7 @@ mutual
         | some e =>
           match (← withReducibleAndInstances <| reduceProj? e.getAppFn) with
           | none   => return none
-          | some r => return mkAppN r e.getAppArgs |>.headBeta
+          | some r => recordUnfold declName; return mkAppN r e.getAppArgs |>.headBeta
       | _ => return none
     | _ => return none
 
@@ -729,8 +729,9 @@ mutual
         if fInfo.levelParams.length != fLvls.length then
           return none
         else
-          let unfoldDefault (_ : Unit) : MetaM (Option Expr) :=
+          let unfoldDefault (_ : Unit) : MetaM (Option Expr) := do
             if fInfo.hasValue then
+              recordUnfold fInfo.name
               deltaBetaDefinition fInfo fLvls e.getAppRevArgs (fun _ => pure none) (fun e => pure (some e))
             else
               return none
@@ -778,11 +779,13 @@ mutual
                   Thus, we should keep `return some r` until Mathlib has been ported to Lean 3.
                   Note that the `Vector` example above does not even work in Lean 3.
                 -/
-                let some recArgPos ← getStructuralRecArgPos? fInfo.name | return some r
+                let some recArgPos ← getStructuralRecArgPos? fInfo.name
+                  | recordUnfold fInfo.name; return some r
                 let numArgs := e.getAppNumArgs
                 if recArgPos >= numArgs then return none
                 let recArg := e.getArg! recArgPos numArgs
                 if !(← isConstructorApp (← whnfMatcher recArg)) then return none
+                recordUnfold fInfo.name
                 return some r
             | _ =>
               if (← getMatcherInfo? fInfo.name).isSome then
@@ -800,7 +803,7 @@ mutual
         unless cinfo.hasValue do return none
         deltaDefinition cinfo lvls
           (fun _ => pure none)
-          (fun e => pure (some e))
+          (fun e => do recordUnfold declName; pure (some e))
     | _ => return none
 end
 
@@ -833,11 +836,11 @@ def reduceRecMatcher? (e : Expr) : MetaM (Option Expr) := do
     | .reduced e => return e
     | _ => matchConstAux e.getAppFn (fun _ => pure none) fun cinfo lvls => do
       match cinfo with
-      | .recInfo «rec»  => reduceRec «rec» lvls e.getAppArgs (fun _ => pure none) (fun e => pure (some e))
-      | .quotInfo «rec» => reduceQuotRec «rec» e.getAppArgs (fun _ => pure none) (fun e => pure (some e))
+      | .recInfo «rec»  => reduceRec «rec» lvls e.getAppArgs (fun _ => pure none) (fun e => do recordUnfold cinfo.name; pure (some e))
+      | .quotInfo «rec» => reduceQuotRec «rec» e.getAppArgs (fun _ => pure none) (fun e => do recordUnfold cinfo.name; pure (some e))
       | c@(.defnInfo _) =>
         if (← isAuxDef c.name) then
-          deltaBetaDefinition c lvls e.getAppRevArgs (fun _ => pure none) (fun e => pure (some e))
+          deltaBetaDefinition c lvls e.getAppRevArgs (fun _ => pure none) (fun e => do recordUnfold c.name; pure (some e))
         else
           return none
       | _ => return none

src/Lean/PrettyPrinter.lean

@@ -13,7 +13,10 @@ import Lean.ParserCompiler
 namespace Lean
 
 def PPContext.runCoreM {α : Type} (ppCtx : PPContext) (x : CoreM α) : IO α :=
-  Prod.fst <$> x.toIO { options := ppCtx.opts, currNamespace := ppCtx.currNamespace, openDecls := ppCtx.openDecls, fileName := "<PrettyPrinter>", fileMap := default }
+  Prod.fst <$> x.toIO { options := ppCtx.opts, currNamespace := ppCtx.currNamespace
+                        openDecls := ppCtx.openDecls
+                        fileName := "<PrettyPrinter>", fileMap := default
+                        diag     := getDiag ppCtx.opts }
                       { env := ppCtx.env, ngen := { namePrefix := `_pp_uniq } }
 
 def PPContext.runMetaM {α : Type} (ppCtx : PPContext) (x : MetaM α) : IO α :=
@@ -48,7 +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 } { env := env }
+  Prod.fst <$> ((ppExpr e).run' { lctx := lctx } { mctx := mctx }).toIO {
+    options := opts, fileName := "<PrettyPrinter>", fileMap := default, diag := getDiag opts
+  } { env := env }
 
 def ppTactic (stx : TSyntax `tactic) : CoreM Format := ppCategory `tactic stx
 

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 }
+    let ctx := { fileName := "", options := ({} : KVMap), fileMap := default, diag := false }
     let state := { env := (← get).env }
     Prod.fst <$> (Lean.Core.CoreM.toIO · ctx state) do Lean.throwKernelException ex
 

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} {env} |>.toBaseIO) with
+    match (← pp.toIO {fileName := "", fileMap := default, diag := false} {env} |>.toBaseIO) with
     | .ok (fmt, _) => pure <| (toString fmt).trim ++ "\n"
     | .error ex =>
       error s!"(internal) failed to pretty print Lean configuration: {ex.toString}"

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} {env} |>.toBaseIO) with
+    match (← act.run {fileName := ictx.fileName, fileMap := ictx.fileMap, diag := false} {env} |>.toBaseIO) with
     | .ok (t, s) =>
       if s.messages.hasErrors then
         throw s.messages