feat: theorem diagnostics (#4924)

When `set_option diagnostics true`, for each theorem with size >
`diagnostics.threshold.proofSize`, display proof size, and the number of
applications for each constant symbol.

src/Lean/Elab/PreDefinition/Basic.lean

@@ -7,6 +7,9 @@ prelude
 import Init.ShareCommon
 import Lean.Compiler.NoncomputableAttr
 import Lean.Util.CollectLevelParams
+import Lean.Util.NumObjs
+import Lean.Util.NumApps
+import Lean.PrettyPrinter
 import Lean.Meta.AbstractNestedProofs
 import Lean.Meta.ForEachExpr
 import Lean.Elab.RecAppSyntax
@@ -17,7 +20,6 @@ namespace Lean.Elab
 open Meta
 open Term
 
-
 /--
   A (potentially recursive) definition.
   The elaborator converts it into Kernel definitions using many different strategies.
@@ -98,15 +100,33 @@ private def compileDecl (decl : Declaration) : TermElabM Bool := do
       throw ex
   return true
 
+register_builtin_option diagnostics.threshold.proofSize : Nat := {
+  defValue := 16384
+  group    := "diagnostics"
+  descr    := "only display proof statistics when proof has at least this number of terms"
+}
+
+private def reportTheoremDiag (d : TheoremVal) : TermElabM Unit := do
+  if (← isDiagnosticsEnabled) then
+    let proofSize ← d.value.numObjs
+    if proofSize > diagnostics.threshold.proofSize.get (← getOptions) then
+      let sizeMsg := MessageData.trace { cls := `size } m!"{proofSize}" #[]
+      let constOccs ← d.value.numApps (threshold := diagnostics.threshold.get (← getOptions))
+      let constOccsMsg ← constOccs.mapM fun (declName, numOccs) => return MessageData.trace { cls := `occs } m!"{MessageData.ofConst (← mkConstWithLevelParams declName)} ↦ {numOccs}" #[]
+      -- let info
+      logInfo <| MessageData.trace { cls := `theorem } m!"{d.name}" (#[sizeMsg] ++ constOccsMsg)
+
 private def addNonRecAux (preDef : PreDefinition) (compile : Bool) (all : List Name) (applyAttrAfterCompilation := true) : TermElabM Unit :=
   withRef preDef.ref do
     let preDef ← abstractNestedProofs preDef
     let decl ←
       match preDef.kind with
       | DefKind.«theorem» =>
-        pure <| Declaration.thmDecl {
+        let d := {
           name := preDef.declName, levelParams := preDef.levelParams, type := preDef.type, value := preDef.value, all
         }
+        reportTheoremDiag d
+        pure <| Declaration.thmDecl d
       | DefKind.«opaque»  =>
         pure <| Declaration.opaqueDecl {
           name := preDef.declName, levelParams := preDef.levelParams, type := preDef.type, value := preDef.value

src/Lean/Util.lean

@@ -32,3 +32,4 @@ import Lean.Util.Heartbeats
 import Lean.Util.SearchPath
 import Lean.Util.SafeExponentiation
 import Lean.Util.NumObjs
+import Lean.Util.NumApps

src/Lean/Util/NumApps.lean

@@ -0,0 +1,57 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Expr
+import Lean.Util.PtrSet
+
+namespace Lean.Expr
+namespace NumApps
+
+unsafe structure State where
+ visited  : PtrSet Expr := mkPtrSet
+ counters : NameMap Nat := {}
+
+unsafe abbrev M := StateM State
+
+unsafe def visit (e : Expr) : M Unit :=
+  unless (← get).visited.contains e do
+    modify fun s => { s with visited := s.visited.insert e }
+    match e with
+      | .forallE _ d b _ => visit d; visit b
+      | .lam _ d b _     => visit d; visit b
+      | .mdata _ b       => visit b
+      | .letE _ t v b _  => visit t; visit v; visit b
+      | .app ..          => e.withApp fun f args => do
+        if let .const declName _ := f then
+          let c := (← get).counters.find? declName |>.getD 0
+          modify fun s => { s with counters := s.counters.insert declName (c+1) }
+        visit f
+        args.forM visit
+      | .proj _ _ b      => visit b
+      | _                => return ()
+
+unsafe def main (e : Expr) : NameMap Nat :=
+  let (_, s) := NumApps.visit e |>.run {}
+  s.counters
+
+end NumApps
+
+/--
+Returns the number of applications for each declaration used in `e`.
+
+This operation is performed in `IO` because the result depends on the memory representation of the object.
+
+Note: Use this function primarily for diagnosing performance issues.
+-/
+def numApps (e : Expr) (threshold : Nat := 0) : IO (Array (Name × Nat)) := do
+  let counters := unsafe NumApps.main e
+  let mut result := #[]
+  for (declName, num) in counters do
+    if num > threshold then
+      result := result.push (declName, num)
+  return result.qsort fun a b => a.2 > b.2
+
+end Lean.Expr