feat: always record per-declaration heartbeat costs

.gitignore

@@ -1,6 +1,7 @@
 *~
 \#*
 .#*
+*.hb.json
 .lake
 lake-manifest.json
 /build

src/Lean/AddDecl.lean

@@ -182,6 +182,7 @@ private def addDeclCore (decl : Declaration) (forceExpose : Bool) : CoreM Unit :
 where
   doAdd := do
     profileitM Exception "type checking" (← getOptions) do
+      Core.withDeclHeartbeats (decl.getTopLevelNames.head?.getD .anonymous) `kernel do
       withTraceNode `Kernel (fun _ => return m!"typechecking declarations {decl.getTopLevelNames}") do
         warnIfUsesSorry decl
         try

src/Lean/CoreM.lean

@@ -9,6 +9,7 @@ prelude
 public import Lean.Util.RecDepth
 public import Lean.ResolveName
 public import Lean.Language.Basic
+public import Lean.Util.Profile
 import Init.While
 import Lean.Compiler.NoncomputableAttr
 
@@ -242,6 +243,8 @@ structure Context extends Context.Cold where
   openDecls      : List OpenDecl := []
   initHeartbeats : Nat := 0
   maxHeartbeats  : Nat := getMaxHeartbeats options
+  heartbeats?    : Option (IO.Ref (Array HeartbeatEntry)) := none
+  costOwner      : CostOwner := .unknown
   currMacroScope : MacroScope := firstFrontendMacroScope
   /--
   If `diag := true`, different parts of the system collect diagnostics.
@@ -557,6 +560,35 @@ instance : MonadLog CoreM where
     let msg := { msg with data := MessageData.withNamingContext { currNamespace := ctx.currNamespace, openDecls := ctx.openDecls } msg.data };
     modify fun s => { s with messages := s.messages.add msg }
 
+/-- Attributes `count` raw heartbeats to `declName`, owned by the surrounding `withCostOwner` scope. -/
+def recordDeclHeartbeats (declName : Name) (phase : Name) (count : Nat) : CoreM Unit := do
+  let ctx ← read
+  if let some sink := ctx.heartbeats? then
+    let entry := { owner := ctx.costOwner.name?.getD declName, declName, phase, heartbeats := count :
+      HeartbeatEntry }
+    sink.modify (·.push entry)
+
+/--
+Runs `act`, attributing the heartbeats it uses to `declName`. The counter is thread-local, so
+`act` must not fork off the work being measured.
+-/
+@[specialize] def withDeclHeartbeats {α : Type} (declName : Name) (phase : Name) (act : CoreM α) :
+    CoreM α := do
+  if (← read).heartbeats?.isNone then return (← act)
+  let startHeartbeats ← IO.getNumHeartbeats
+  let a ← act
+  let stopHeartbeats ← IO.getNumHeartbeats
+  recordDeclHeartbeats declName phase (stopHeartbeats - startHeartbeats)
+  return a
+
+/-- Attributes heartbeats recorded inside `x` to `declName`; an already-fixed owner is kept. -/
+def withCostOwner [Monad m] [MonadControlT CoreM m] (declName : Name) (x : m α) : m α :=
+  controlAt CoreM fun runInBase =>
+    withReader (fun ctx =>
+      match ctx.costOwner with
+      | .unknown | .pending _ => { ctx with costOwner := .fixed declName }
+      | .fixed _ => ctx) (runInBase x)
+
 /--
 Includes a given task (such as from `wrapAsyncAsSnapshot`) in the overall snapshot tree for this
 command's elaboration, making its result available to reporting and the language server. The

src/Lean/Data/Lsp/Capabilities.lean

@@ -76,6 +76,11 @@ structure LeanClientCapabilities where
   -/
   silentDiagnosticSupport? : Option Bool := none
   /--
+  Whether the client supports `DiagnosticWith.heartbeats?`.
+  If `none` or `false`, the machine-dependent field is cleared before diagnostics are served.
+  -/
+  heartbeatSupport? : Option Bool := none
+  /--
   The latest RPC wire format supported by the client.
   Defaults to `v0` when `none`.
   -/
@@ -104,6 +109,13 @@ def ClientCapabilities.silentDiagnosticSupport (c : ClientCapabilities) : Bool :
     | return false
   return silentDiagnosticSupport
 
+def ClientCapabilities.heartbeatSupport (c : ClientCapabilities) : Bool := Id.run do
+  let some lean := c.lean?
+    | return false
+  let some heartbeatSupport := lean.heartbeatSupport?
+    | return false
+  return heartbeatSupport
+
 def ClientCapabilities.rpcWireFormat (c : ClientCapabilities) : RpcWireFormat := Id.run do
   let some lean := c.lean?
     | return .v0

src/Lean/Data/Lsp/Diagnostics.lean

@@ -144,6 +144,11 @@ structure DiagnosticWith (α : Type) where
   tags? : Option (Array DiagnosticTag) := none
   /-- Additional Lean-specific metadata about the diagnostic. -/
   leanTags? : Option (Array LeanDiagnosticTag) := none
+  /--
+  Extension: heartbeats used to process the declaration this diagnostic reports on, in
+  `IO.getNumHeartbeats` units (divide by 1000 for the `maxHeartbeats` unit).
+  -/
+  heartbeats? : Option Nat := none
   /-- An array of related diagnostic information, e.g. when symbol-names within a scope collide all definitions can be marked via this property. -/
   relatedInformation? : Option (Array DiagnosticRelatedInformation) := none
   /-- A data entry field that is preserved between a `textDocument/publishDiagnostics` notification and `textDocument/codeAction` request. -/

src/Lean/Elab/Command.lean

@@ -36,6 +36,7 @@ structure State where
   auxDeclNGen    : DeclNameGenerator := .ofPrefix .anonymous
   infoState      : InfoState := {}
   traceState     : TraceState := {}
+  heartbeatsRef? : Option (IO.Ref (Array HeartbeatEntry)) := none
   snapshotTasks  : Array (Language.SnapshotTask Language.SnapshotTree) := #[]
   prevLinterStates : Option (Task (Array LinterState)) := none
   codeQualityEntryTasks : Array (Task (Array Linter.CodeQualityLogEntry)) := #[]
@@ -68,6 +69,7 @@ structure Context where
   errors; see also `logMessage` below.
   -/
   suppressElabErrors : Bool := false
+  costOwner      : CostOwner := .unknown
 
 abbrev CommandElabM := ReaderT Context $ StateRefT State $ EIO Exception
 abbrev CommandElab  := Syntax → CommandElabM Unit
@@ -258,6 +260,23 @@ instance : MonadQuotation CommandElabM where
   getContext          := do (← read).quotContext?.getDM getMainModule
   withFreshMacroScope := Command.withFreshMacroScope
 
+/-- Best-effort qualification of a declaration name as written; see `Core.withCostOwner`. -/
+def approxCostOwnerName (currNamespace shortName : Name) : Name :=
+  if (`_root_).isPrefixOf shortName then
+    shortName.replacePrefix `_root_ .anonymous
+  else currNamespace ++ shortName
+
+/--
+Marks the heartbeat cost owner of a declaration command; a command elaborated while an owner is
+pending or fixed is machine-generated, so that owner is fixed for its whole elaboration.
+-/
+def withCostOwner? (declName? : Option Name) (x : CommandElabM α) : CommandElabM α :=
+  withReader (fun ctx =>
+    match ctx.costOwner, declName? with
+    | .pending declName, _ | .fixed declName, _ => { ctx with costOwner := .fixed declName }
+    | .unknown, some declName => { ctx with costOwner := .pending declName }
+    | .unknown, none => ctx) x
+
 private def runCore (x : CoreM α) : CommandElabM α := do
   let s ← get
   let ctx ← read
@@ -277,7 +296,9 @@ private def runCore (x : CoreM α) : CommandElabM α := do
     currMacroScope     := ctx.currMacroScope
     options            := scope.opts
     cancelTk?          := ctx.cancelTk?
-    suppressElabErrors := ctx.suppressElabErrors }
+    suppressElabErrors := ctx.suppressElabErrors
+    heartbeats?        := s.heartbeatsRef?
+    costOwner          := ctx.costOwner }
   let x : EIO _ _ := x.run coreCtx {
     env
     ngen := s.ngen
@@ -1102,6 +1123,7 @@ private def liftCommandElabMCore (cmd : CommandElabM α) (throwOnError : Bool) :
       snap? := none
       cancelTk? := ctx.cancelTk?
       suppressElabErrors := ctx.suppressElabErrors
+      costOwner := ctx.costOwner
     } |>.run {
       env := s.env
       nextMacroScope := s.nextMacroScope
@@ -1110,6 +1132,7 @@ private def liftCommandElabMCore (cmd : CommandElabM α) (throwOnError : Bool) :
       auxDeclNGen := s.auxDeclNGen
       scopes := [{ header := "", opts := ctx.options }]
       infoState.enabled := s.infoState.enabled
+      heartbeatsRef? := ctx.heartbeats?
     }
   modify fun coreState => { coreState with
     env := commandState.env

src/Lean/Elab/Declaration.lean

@@ -159,6 +159,7 @@ def expandNamespacedDeclaration : Macro := fun stx => do
 
 @[builtin_command_elab declaration, builtin_incremental]
 def elabDeclaration : CommandElab := fun stx => do
+  withCostOwner? ((getDeclName? stx).map (approxCostOwnerName (← getScope).currNamespace)) do
   withExporting (isExporting := (← getScope).isPublic) do
   let modifiers : TSyntax ``Parser.Command.declModifiers := ⟨stx[0]⟩
   let decl     := stx[1]
@@ -287,6 +288,8 @@ def expandMutualPreamble : Macro := fun stx =>
 
 @[builtin_command_elab «mutual», builtin_incremental]
 def elabMutual : CommandElab := fun stx => do
+  withCostOwner? ((stx[1].getArgs[0]? |>.bind getDeclName?).map
+    (approxCostOwnerName (← getScope).currNamespace)) do
   withExporting (isExporting := (← getScope).isPublic) do
   if isMutualDefLike stx then
     -- only case implementing incrementality currently

src/Lean/Elab/Frontend.lean

@@ -401,6 +401,20 @@ def runFrontend
     }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
+  if let some oleanFileName := oleanFileName? then
+    if cmdState.heartbeatsRef?.isSome then
+      let entries ← Language.Lean.collectHeartbeatEntries snap
+      let path := oleanFileName.withExtension "hb.json"
+      IO.FS.writeFile path <| Json.compress <| Json.mkObj [
+        ("module", mainModuleName.toString),
+        ("unit", "raw"),
+        ("entries", Json.arr <| entries.map fun e => Json.mkObj [
+          ("owner", e.owner.toString),
+          ("decl", e.declName.toString),
+          ("phase", e.phase.toString),
+          ("heartbeats", toJson e.heartbeats)])
+      ]
+
   if let some out := trace.profiler.output.get? opts then
     let traceStates := snaps.getAll.map (·.traces)
     let profile ← Firefox.Profile.export mainModuleName.toString startTime traceStates opts

src/Lean/Elab/MutualDef.lean

@@ -514,6 +514,7 @@ private def useProofAsSorry (k : DefKind) : CoreM Bool := do
 
 private def elabFunValues (headers : Array DefViewElabHeader) (vars : Array Expr) (sc : Command.Scope) : TermElabM (Array Expr) :=
   headers.mapM fun header => do
+    Core.withCostOwner header.declName do
     let mut reusableResult? := none
     if let some snap := header.bodySnap? then
       if let some old := snap.old? then
@@ -582,6 +583,7 @@ private def elabFunValues (headers : Array DefViewElabHeader) (vars : Array Expr
                 in scope or explicitly omit them:\
               \n  omit {MessageData.joinSep unusedVars.toList " "} in theorem ..."
         return val
+    Core.recordDeclHeartbeats header.declName `elab state.«meta».core.passedHeartbeats
     if let some snap := header.bodySnap? then
       snap.new.resolve <| some {
         diagnostics :=
@@ -1317,6 +1319,8 @@ this warning can be disabled with `set_option warn.classDefReducibility false`."
       -- forked correctly.
       withDeprecationContextFromAttrs oldAttrs do
       withDeclName header.declName do
+      -- set the cost owner before going async so the forked context carries it
+      Core.withCostOwner header.declName do
       wrapAsyncAsSnapshot (desc := s!"elaborating proof of {declId.declName}")
         (cancelTk? := cancelTk) fun _ => do profileitM Exception "elaboration" (← getOptions) do
       setEnv async.asyncEnv
@@ -1470,7 +1474,7 @@ Logs a snapshot task that waits for the entire snapshot tree in `defsParsedSnap`
 is error-free and contains no syntactical `sorry`s.
 -/
 private def logGoalsAccomplishedSnapshotTask (views : Array DefView)
-    (defsParsedSnap : DefsParsedSnapshot) : TermElabM Unit := do
+    (defsParsedSnap : DefsParsedSnapshot) (sinkStart : Nat) : TermElabM Unit := do
   if ! Lean.Elab.inServer.get (← getOptions) then
     -- Skip 'goals accomplished' task if we are on the command line.
     -- These messages are only used in the language server.
@@ -1490,10 +1494,34 @@ private def logGoalsAccomplishedSnapshotTask (views : Array DefView)
         msg.severity matches .error || msg.data.hasTag (· == `hasSorry)
     if hasErrorOrSorry then
       return
+    -- `example`s share one `declName` per namespace, making their costs mutually
+    -- indistinguishable, so they get no count.
+    let blockOwners : Array Name := if views.any (·.2 matches .theorem) then
+      defsParsedSnap.defs.filterMap (·.headerProcessedSnap.get |>.map (·.view.declName))
+    else
+      #[]
+    let ownerCosts : Std.HashMap Name Nat ← do
+      match (← readThe Core.Context).heartbeats? with
+      | some sink =>
+        if blockOwners.isEmpty then pure ∅ else
+          let entries ← sink.get
+          pure <| entries.foldl (init := ∅) (start := sinkStart) fun m e =>
+            if blockOwners.contains e.owner then
+              m.alter e.owner (fun c => some (c.getD 0 + e.heartbeats))
+            else m
+      | none => pure ∅
     for d in defsParsedSnap.defs, (ref, kind) in views do
-      let logGoalsAccomplished :=
-        let msgData := .tagged `goalsAccomplished m!"Goals accomplished!"
-        logAt ref msgData (severity := .information) (isSilent := true)
+      let heartbeats? kind := do
+        guard (kind matches DefKind.theorem)
+        let s ← d.headerProcessedSnap.get
+        ownerCosts[s.view.declName]?
+      let logGoalsAccomplished := do
+        let msg := m!"Goals accomplished!"
+        let msg := match heartbeats? kind with
+          -- the outer tag must stay outermost so `Message.kind` remains `goalsAccomplished`
+          | some hb => MessageData.tagged (.num `heartbeats hb) msg
+          | none => msg
+        logAt ref (.tagged `goalsAccomplished msg) (severity := .information) (isSilent := true)
       match kind with
       | .theorem =>
         logGoalsAccomplished
@@ -1566,9 +1594,12 @@ def elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
   let sc ← getScope
   -- use hash of all names as stable quot context
   withInitQuotContext (some (hash (views.map (·.declId[0].getId)))) do
+  let mut sinkStart := 0
+  if let some sink := (← get).heartbeatsRef? then
+    sinkStart := (← sink.get).size
   runTermElabM fun vars => do
     Term.elabMutualDef vars sc views
-    Term.logGoalsAccomplishedSnapshotTask views defsParsedSnap
+    Term.logGoalsAccomplishedSnapshotTask views defsParsedSnap sinkStart
 
 builtin_initialize
   registerTraceClass `Elab.definition.mkClosure

src/Lean/Elab/PreDefinition/Main.lean

@@ -295,6 +295,7 @@ def addPreDefinitions (docCtx : LocalContext × LocalInstances) (preDefs : Array
       let preDefs ← betaReduceLetRecApps preDefs
       let cliques := partitionPreDefs preDefs
       for preDefs in cliques do
+        Core.withCostOwner preDefs[0]!.declName do
         trace[Elab.definition.scc] "{preDefs.map (·.declName)}"
         if preDefs.size == 1 && isNonRecursive preDefs[0]! then
           /-

src/Lean/Language/Lean.lean

@@ -520,6 +520,8 @@ where
       -- now that imports have been loaded, check options again
       opts ← reparseOptions opts
       let cmdState := Elab.Command.mkState headerEnv msgLog opts
+      -- enables heartbeat recording; `doElab` installs a fresh sink per command
+      let cmdState := { cmdState with heartbeatsRef? := some (← IO.mkRef #[]) }
       let cmdState := { cmdState with
         infoState := {
           enabled := true
@@ -685,8 +687,13 @@ where
       let mut reportedCmdState := cmdState
       let cmdline := internal.cmdlineSnapshots.get scope.opts && !Parser.isTerminalCommand stx
       if cmdline then
-        -- discard all metadata apart from the environment; see `internal.cmdlineSnapshots`
-        reportedCmdState := { env := reportedCmdState.env, maxRecDepth := 0 }
+        -- discard all metadata apart from the environment and heartbeat sink; see
+        -- `internal.cmdlineSnapshots`
+        reportedCmdState := {
+          env := reportedCmdState.env
+          maxRecDepth := 0
+          heartbeatsRef? := reportedCmdState.heartbeatsRef?
+        }
       resultPromise.resolve {
         diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
         traces := cmdState.traceState
@@ -755,9 +762,12 @@ where
       LeanProcessingM Command.State := do
     let ctx ← read
     let scope := cmdState.scopes.head!
+    -- fresh heartbeat sink per command; `runFrontend` aggregates across commands for the sidecar
+    let heartbeatsRef? : Option (IO.Ref (Array HeartbeatEntry)) ←
+      cmdState.heartbeatsRef?.mapM fun _ => IO.mkRef #[]
     -- reset per-command state
     let cmdStateRef ← IO.mkRef { cmdState with
-      messages := .empty, traceState := {}, snapshotTasks := #[] }
+      messages := .empty, traceState := {}, snapshotTasks := #[], heartbeatsRef? }
     let cmdCtx : Elab.Command.Context := { ctx with
       cmdPos       := beginPos
       snap?        := if internal.cmdlineSnapshots.get scope.opts then none else snap
@@ -811,6 +821,22 @@ where goCmd snap :=
   else
     snap.elabSnap.resultSnap.get.cmdState
 
+/-- Waits for and returns the heartbeat entries recorded by all commands. -/
+partial def collectHeartbeatEntries (snap : InitialSnapshot) : BaseIO (Array HeartbeatEntry) := do
+  let some parsed := snap.result? | return #[]
+  let some processed := parsed.processedSnap.get.result? | return #[]
+  goCmd processed.firstCmdSnap.get #[]
+where
+  goCmd (snap : CommandParsedSnapshot) (acc : Array HeartbeatEntry) :
+      BaseIO (Array HeartbeatEntry) := do
+    let mut acc := acc
+    if let some ref := snap.elabSnap.resultSnap.get.cmdState.heartbeatsRef? then
+      acc := acc ++ (← ref.get)
+    if let some next := snap.nextCmdSnap? then
+      goCmd next.get acc
+    else
+      return acc
+
 /--
 Returns `snap` with all elaborated command data discarded, retaining only the imported
 environment from the header.

src/Lean/Message.lean

@@ -193,6 +193,18 @@ partial def hasTag : MessageData → Bool
   | ofOriginatingSyntax _ msg   => hasTag msg
   | _                           => false
 
+/-- Returns the first result of `p` on a tag, traversing like `hasTag`. -/
+partial def findTag? (p : Name → Option α) : MessageData → Option α
+  | withContext _ msg           => findTag? p msg
+  | withNamingContext _ msg     => findTag? p msg
+  | nest _ msg                  => findTag? p msg
+  | group msg                   => findTag? p msg
+  | compose msg₁ msg₂           => findTag? p msg₁ <|> findTag? p msg₂
+  | tagged n msg                => p n <|> findTag? p msg
+  | trace data msg msgs         => p data.cls <|> findTag? p msg <|> msgs.findSome? (findTag? p)
+  | ofOriginatingSyntax _ msg    => findTag? p msg
+  | _                           => none
+
 /--
 Returns the top-level tag of the message.
 If none, returns `Name.anonymous`.

src/Lean/Server/FileWorker.lean

@@ -304,8 +304,10 @@ This option can only be set on the command line, not in the lakefile or via `set
             let mut msgs := node.element.diagnostics.msgLog.toArray
             if ! ctx.initParams.capabilities.silentDiagnosticSupport then
               msgs := msgs.filter (! ·.isSilent)
-            let diags ← msgs.mapM
+            let mut diags ← msgs.mapM
               (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
+            unless ctx.initParams.capabilities.heartbeatSupport do
+              diags := diags.map fun d => { d with heartbeats? := none }
             if let some cacheRef := node.element.diagnostics.interactiveDiagsRef? then
               cacheRef.set <| some <| .mk { diags : MemorizedInteractiveDiagnostics }
             pure diags