feat: trace.profiler export to Firefox Profiler (#3801)

Reusing the best profiling UI out there Usage: ``` lean -Dtrace.profiler=true -Dtrace.profiler.output=profile.json foo.lean ... ``` then open `profile.json` in https://profiler.firefox.com/. See also `script/collideProfiles.lean` for minimizing and merging profiles.
2026-03-17 10:24:07 +00:00 · 2024-04-15 14:13:14 +02:00
parent 6712913bfe
commit 2dcd42f395
13 changed files with 450 additions and 42 deletions
--- a/.github/workflows/copyright-header.yml
+++ b/.github/workflows/copyright-header.yml
@@ -10,7 +10,7 @@ jobs:

    - name: Verify .lean files start with a copyright header.
      run: |
-        FILES=$(find . -type d \( -path "./tests" -o -path "./doc" -o -path "./src/lake/examples" -o -path "./src/lake/tests" -o -path "./build" -o -path "./nix" \) -prune -o -type f -name "*.lean" -exec perl -ne 'BEGIN { $/ = undef; } print "$ARGV\n" if !m{\A/-\nCopyright}; exit;' {} \;)
+        FILES=$(find ./src -type d \( -path "./src/lake/examples" -o -path "./src/lake/tests" \) -prune -o -type f -name "*.lean" -exec perl -ne 'BEGIN { $/ = undef; } print "$ARGV\n" if !m{\A/-\nCopyright}; exit;' {} \;)
        if [ -n "$FILES" ]; then
          echo "Found .lean files which do not have a copyright header:"
          echo "$FILES"
--- a/script/collideProfiles.lean
+++ b/script/collideProfiles.lean
@@ -0,0 +1,28 @@
+import Lean.Util.Profiler
+
+/-!
+
+Usage:
+```sh
+lean --run ./script/collideProfiles.lean **/*.lean.json ... > merged.json
+```
+
+Merges multiple `trace.profiler.output` profiles into a single one while deduplicating samples with
+the same stack. This is useful for building cumulative profiles of medium-to-large projects because
+Firefox Profiler cannot handle hundreds of tracks and the deduplication will also ensure that the
+profile is small enough for uploading.
+
+As ordering of samples is not meaningful after this transformation, only "Call Tree" and "Flame
+Graph" are useful for such profiles.
+-/
+
+open Lean
+
+def main (args : List String) : IO Unit := do
+  let profiles ← args.toArray.mapM fun path => do
+    let json ← IO.FS.readFile ⟨path⟩
+    let profile ← IO.ofExcept $ Json.parse json
+    IO.ofExcept <| fromJson? profile
+  -- NOTE: `collide` should not be interpreted
+  let profile := Firefox.Profile.collide profiles
+  IO.println <| Json.compress <| toJson profile
--- a/src/Lean/Compiler/Main.lean
+++ b/src/Lean/Compiler/Main.lean
@@ -12,7 +12,8 @@ Run the code generation pipeline for all declarations in `declNames`
 that fulfill the requirements of `shouldGenerateCode`.
 -/
 def compile (declNames : Array Name) : CoreM Unit := do profileitM Exception "compiler new" (← getOptions) do
-  discard <| LCNF.compile declNames
+  withTraceNode `Compiler (fun _ => return m!"compiling: {declNames}") do
+    discard <| LCNF.compile declNames

 builtin_initialize
  registerTraceClass `Compiler
--- a/src/Lean/CoreM.lean
+++ b/src/Lean/CoreM.lean
@@ -330,10 +330,13 @@ opaque compileDeclsNew (declNames : List Name) : CoreM Unit

 def compileDecl (decl : Declaration) : CoreM Unit := do
  let opts ← getOptions
+  let decls := Compiler.getDeclNamesForCodeGen decl
  if compiler.enableNew.get opts then
-    compileDeclsNew (Compiler.getDeclNamesForCodeGen decl)
-  match (← getEnv).compileDecl opts decl with
-  | Except.ok env   => setEnv env
+    compileDeclsNew decls
+  let res ← withTraceNode `compiler (fun _ => return m!"compiling old: {decls}") do
+    return (← getEnv).compileDecl opts decl
+  match res with
+  | Except.ok env => setEnv env
  | Except.error (KernelException.other msg) =>
    checkUnsupported decl -- Generate nicer error message for unsupported recursors and axioms
    throwError msg
--- a/src/Lean/Elab/Command.lean
+++ b/src/Lean/Elab/Command.lean
@@ -146,10 +146,13 @@ private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceStat

 private def addTraceAsMessages : CommandElabM Unit := do
  let ctx ← read
-  modify fun s => { s with
-    messages          := addTraceAsMessagesCore ctx s.messages s.traceState
-    traceState.traces := {}
-  }
+  -- do not add trace messages if `trace.profiler.output` is set as it would be redundant and
+  -- pretty printing the trace messages is expensive
+  if trace.profiler.output.get? (← getOptions) |>.isNone then
+    modify fun s => { s with
+      messages          := addTraceAsMessagesCore ctx s.messages s.traceState
+      traceState.traces := {}
+    }

 def liftCoreM (x : CoreM α) : CommandElabM α := do
  let s ← get
@@ -206,7 +209,8 @@ def runLinters (stx : Syntax) : CommandElabM Unit := do
      let linters ← lintersRef.get
      unless linters.isEmpty do
        for linter in linters do
-          withTraceNode `Elab.lint (fun _ => return m!"running linter: {linter.name}") do
+          withTraceNode `Elab.lint (fun _ => return m!"running linter: {linter.name}")
+              (tag := linter.name.toString) do
            let savedState ← get
            try
              linter.run stx
@@ -278,7 +282,9 @@ partial def elabCommand (stx : Syntax) : CommandElabM Unit := do
        -- list of commands => elaborate in order
        -- The parser will only ever return a single command at a time, but syntax quotations can return multiple ones
        args.forM elabCommand
-      else withTraceNode `Elab.command (fun _ => return stx) do
+      else withTraceNode `Elab.command (fun _ => return stx) (tag :=
+        -- special case: show actual declaration kind for `declaration` commands
+        (if stx.isOfKind ``Parser.Command.declaration then stx[1] else stx).getKind.toString) do
        let s ← get
        match (← liftMacroM <| expandMacroImpl? s.env stx) with
        | some (decl, stxNew?) =>
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -7,6 +7,7 @@ prelude
 import Lean.Language.Lean
 import Lean.Util.Profile
 import Lean.Server.References
+import Lean.Util.Profiler

 namespace Lean.Elab.Frontend

@@ -108,6 +109,7 @@ def runFrontend
    (trustLevel : UInt32 := 0)
    (ileanFileName? : Option String := none)
    : IO (Environment × Bool) := do
+  let startTime := (← IO.monoNanosNow).toFloat / 1000000000
  let inputCtx := Parser.mkInputContext input fileName
  -- TODO: replace with `#lang` processing
  if /- Lean #lang? -/ true then
@@ -119,6 +121,7 @@ def runFrontend
    let (env, messages) ← processHeader (leakEnv := true) header opts messages inputCtx trustLevel
    let env := env.setMainModule mainModuleName
    let mut commandState := Command.mkState env messages opts
+    let elabStartTime := (← IO.monoNanosNow).toFloat / 1000000000

    if ileanFileName?.isSome then
      -- Collect InfoTrees so we can later extract and export their info to the ilean file
@@ -135,6 +138,19 @@ def runFrontend
      let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
      IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean

+    if let some out := trace.profiler.output.get? opts then
+      let traceState := s.commandState.traceState
+      -- importing does not happen in an elaboration monad, add now
+      let traceState := { traceState with
+        traces := #[{
+          ref := .missing,
+          msg := .trace { cls := `Import, startTime, stopTime := elabStartTime }
+            (.ofFormat "importing") #[]
+        }].toPArray' ++ traceState.traces
+      }
+      let profile ← Firefox.Profile.export mainModuleName.toString startTime traceState opts
+      IO.FS.writeFile ⟨out⟩ <| Json.compress <| toJson profile
+
    return (s.commandState.env, !s.commandState.messages.hasErrors)

  let ctx := { inputCtx with mainModuleName, opts, trustLevel }
--- a/src/Lean/Elab/Tactic/Basic.lean
+++ b/src/Lean/Elab/Tactic/Basic.lean
@@ -147,7 +147,7 @@ partial def evalTactic (stx : Syntax) : TacticM Unit := do
      if k == nullKind then
        -- Macro writers create a sequence of tactics `t₁ ... tₙ` using `mkNullNode #[t₁, ..., tₙ]`
        stx.getArgs.forM evalTactic
-      else withTraceNode `Elab.step (fun _ => return stx) do
+      else withTraceNode `Elab.step (fun _ => return stx) (tag := stx.getKind.toString) do
        let evalFns := tacticElabAttribute.getEntries (← getEnv) stx.getKind
        let macros  := macroAttribute.getEntries (← getEnv) stx.getKind
        if evalFns.isEmpty && macros.isEmpty then
--- a/src/Lean/Elab/Term.lean
+++ b/src/Lean/Elab/Term.lean
@@ -1379,7 +1379,8 @@ where
 private partial def elabTermAux (expectedType? : Option Expr) (catchExPostpone : Bool) (implicitLambda : Bool) : Syntax → TermElabM Expr
  | .missing => mkSyntheticSorryFor expectedType?
  | stx => withFreshMacroScope <| withIncRecDepth do
-    withTraceNode `Elab.step (fun _ => return m!"expected type: {expectedType?}, term\n{stx}") do
+    withTraceNode `Elab.step (fun _ => return m!"expected type: {expectedType?}, term\n{stx}")
+      (tag := stx.getKind.toString) do
    checkSystem "elaborator"
    let env ← getEnv
    let result ← match (← liftMacroM (expandMacroImpl? env stx)) with
--- a/src/Lean/Message.lean
+++ b/src/Lean/Message.lean
@@ -46,6 +46,18 @@ structure PPFormat where
  /-- Searches for synthetic sorries in original input. Used to filter out certain messages. -/
  hasSyntheticSorry : MetavarContext → Bool := fun _ => false

+structure TraceData where
+  /-- Trace class, e.g. `Elab.step`. -/
+  cls       : Name
+  /-- Start time in seconds; 0 if unknown to avoid `Option` allocation. -/
+  startTime : Float := 0
+  /-- Stop time in seconds; 0 if unknown to avoid `Option` allocation. -/
+  stopTime  : Float := startTime
+  /-- Whether trace node defaults to collapsed in the infoview. -/
+  collapsed : Bool := true
+  /-- Optional tag shown in `trace.profiler.output` output after the trace class name. -/
+  tag       : String := ""
+
 /-- Structured message data. We use it for reporting errors, trace messages, etc. -/
 inductive MessageData where
  /-- Eagerly formatted text. We inspect this in various hacks, so it is not immediately subsumed by `ofPPFormat`. -/
@@ -65,7 +77,7 @@ 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
-  | trace (cls : Name) (msg : MessageData) (children : Array MessageData) (collapsed : Bool)
+  | trace (data : TraceData) (msg : MessageData) (children : Array MessageData)
  deriving Inhabited

 namespace MessageData
@@ -90,7 +102,7 @@ partial def hasTag : MessageData → Bool
  | group msg               => hasTag msg
  | compose msg₁ msg₂       => hasTag msg₁ || hasTag msg₂
  | tagged n msg            => p n || hasTag msg
-  | trace cls msg msgs _    => p cls || hasTag msg || msgs.any hasTag
+  | trace data msg msgs     => p data.cls || hasTag msg || msgs.any hasTag
  | _                       => false

 /-- An empty message. -/
@@ -133,7 +145,7 @@ where
  | group msg               => visit mctx? msg
  | compose msg₁ msg₂       => visit mctx? msg₁ || visit mctx? msg₂
  | tagged _ msg            => visit mctx? msg
-  | trace _ msg msgs _      => visit mctx? msg || msgs.any (visit mctx?)
+  | trace _ msg msgs        => visit mctx? msg || msgs.any (visit mctx?)
  | _                       => false

 partial def formatAux : NamingContext → Option MessageDataContext → MessageData → IO Format
@@ -147,8 +159,11 @@ partial def formatAux : NamingContext → Option MessageDataContext → MessageD
  | 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,       trace cls header children _ => do
-    let msg := f!"[{cls}] {(← formatAux nCtx ctx header).nest 2}"
+  | nCtx, ctx,       trace data header children => do
+    let mut msg := f!"[{data.cls}]"
+    if data.startTime != 0 then
+      msg := f!"{msg} [{data.stopTime - data.startTime}]"
+    msg := f!"{msg} {(← formatAux nCtx ctx header).nest 2}"
    let children ← children.mapM (formatAux nCtx ctx)
    return .nest 2 (.joinSep (msg::children.toList) "\n")

--- a/src/Lean/PrettyPrinter.lean
+++ b/src/Lean/PrettyPrinter.lean
@@ -75,8 +75,8 @@ private partial def noContext : MessageData → MessageData
  | MessageData.group msg  => MessageData.group (noContext msg)
  | MessageData.compose msg₁ msg₂ => MessageData.compose (noContext msg₁) (noContext msg₂)
  | MessageData.tagged tag msg => MessageData.tagged tag (noContext msg)
-  | MessageData.trace cls header children collapsed =>
-    MessageData.trace cls (noContext header) (children.map noContext) collapsed
+  | MessageData.trace data header children =>
+    MessageData.trace data (noContext header) (children.map noContext)
  | msg => msg

 -- strip context (including environments with registered pretty printers) to prevent infinite recursion when pretty printing pretty printer error
--- a/src/Lean/Util/Profiler.lean
+++ b/src/Lean/Util/Profiler.lean
@@ -0,0 +1,322 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Ullrich
+-/
+prelude
+import Lean.Util.Trace
+
+/-! `trace.profiler.output` Firefox Profiler integration -/
+
+namespace Lean.Firefox
+
+/-! Definitions from https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.js -/
+
+abbrev Milliseconds := Float
+
+structure Category where
+  name : String
+  color : String
+  subcategories : Array String := #[]
+deriving FromJson, ToJson
+
+structure ProfileMeta where
+  interval : Milliseconds := 0  -- should be irrelevant with `tracing-ms`
+  startTime : Milliseconds
+  categories : Array Category := #[]
+  processType : Nat := 0
+  product : String := "lean"
+  preprocessedProfileVersion : Nat := 48
+  markerSchema : Array Json := #[]
+deriving FromJson, ToJson
+
+structure StackTable where
+  frame : Array Nat
+  «prefix» : Array (Option Nat)
+  category : Array Nat
+  subcategory : Array Nat
+  length : Nat
+deriving FromJson, ToJson
+
+structure SamplesTable where
+  stack : Array Nat
+  time : Array Milliseconds
+  weight : Array Milliseconds
+  weightType : String := "tracing-ms"
+  length : Nat
+deriving FromJson, ToJson
+
+structure FuncTable where
+  name : Array Nat
+  isJS : Array Json := #[]
+  relevantForJS : Array Json := #[]
+  resource : Array Int
+  fileName : Array (Option Nat)
+  lineNumber : Array (Option Nat)
+  columnNumber : Array (Option Nat)
+  length : Nat
+deriving FromJson, ToJson
+
+structure FrameTable where
+  func : Array Nat
+  inlineDepth : Array Json := #[]
+  innerWindowID : Array Json := #[]
+  length : Nat
+deriving FromJson, ToJson
+
+structure RawMarkerTable where
+  data : Array Json := #[]
+  name : Array Json := #[]
+  length : Nat := 0
+deriving FromJson, ToJson
+
+structure ResourceTable where
+  type : Array Json := #[]
+  length : Nat := 0
+deriving FromJson, ToJson
+
+structure Thread where
+  name : String
+  processType : String := "default"
+  isMainThread : Bool := true
+  samples : SamplesTable
+  markers: RawMarkerTable := {}
+  stackTable : StackTable
+  frameTable : FrameTable
+  resourceTable : ResourceTable := {}
+  stringArray : Array String
+  funcTable : FuncTable
+deriving FromJson, ToJson
+
+structure Profile where
+  meta : ProfileMeta
+  libs : Array Json := #[]
+  threads : Array Thread
+deriving FromJson, ToJson
+
+/-- Thread with maps necessary for computing max sharing indices -/
+structure ThreadWithMaps extends Thread where
+  stringMap : HashMap String Nat := {}
+  funcMap : HashMap Nat Nat := {}
+  stackMap : HashMap (Nat × Option Nat) Nat := {}
+  /-- Last timestamp encountered: stop time of preceding sibling, or else start time of parent.  -/
+  lastTime : Float := 0
+
+-- TODO: add others, dynamically?
+def categories : Array Category := #[
+  { name := "Other", color := "gray" },
+  { name := "Elab", color := "red" },
+  { name := "Meta", color := "yellow" }
+]
+
+private partial def addTrace (pp : Bool) (thread : ThreadWithMaps) (trace : MessageData) :
+    IO ThreadWithMaps :=
+  (·.2) <$> StateT.run (go none trace) thread
+where
+  go parentStackIdx? : _ → StateT ThreadWithMaps IO Unit
+    | .trace data msg children => do
+      if data.startTime == 0 then
+        return  -- no time data, skip
+      let mut funcName := data.cls.toString
+      if !data.tag.isEmpty then
+        funcName := s!"{funcName}: {data.tag}"
+      if pp then
+        funcName := s!"{funcName}: {← msg.format}"
+      let strIdx ← modifyGet fun thread =>
+        if let some idx := thread.stringMap.find? funcName then
+          (idx, thread)
+        else
+          (thread.stringMap.size, { thread with
+            stringArray := thread.stringArray.push funcName
+            stringMap := thread.stringMap.insert funcName thread.stringMap.size })
+      let category := categories.findIdx? (·.name == data.cls.getRoot.toString) |>.getD 0
+      let funcIdx ← modifyGet fun thread =>
+        if let some idx := thread.funcMap.find? strIdx then
+          (idx, thread)
+        else
+          (thread.funcMap.size, { thread with
+            funcTable := {
+              name := thread.funcTable.name.push strIdx
+              resource := thread.funcTable.resource.push (-1)
+              -- the following fields could be inferred from `Syntax` in the message
+              fileName := thread.funcTable.fileName.push none
+              lineNumber := thread.funcTable.lineNumber.push none
+              columnNumber := thread.funcTable.columnNumber.push none
+              length := thread.funcTable.length + 1
+            }
+            frameTable := {
+              func := thread.frameTable.func.push thread.funcMap.size
+              length := thread.frameTable.length + 1
+            }
+            funcMap := thread.funcMap.insert strIdx thread.funcMap.size })
+      let frameIdx := funcIdx
+      let stackIdx ← modifyGet fun thread =>
+        if let some idx := thread.stackMap.find? (frameIdx, parentStackIdx?) then
+          (idx, thread)
+        else
+          (thread.stackMap.size, { thread with
+            stackTable := {
+              frame := thread.stackTable.frame.push frameIdx
+              «prefix» := thread.stackTable.prefix.push parentStackIdx?
+              category := thread.stackTable.category.push category
+              subcategory := thread.stackTable.subcategory.push 0
+              length := thread.stackTable.length + 1
+            }
+            stackMap := thread.stackMap.insert (frameIdx, parentStackIdx?) thread.stackMap.size })
+      modify fun thread => { thread with lastTime := data.startTime }
+      for c in children do
+        if let some nextStart := getNextStart? c then
+          -- add run time slice between children/before first child
+          modify fun thread => { thread with samples := {
+            stack := thread.samples.stack.push stackIdx
+            time := thread.samples.time.push (thread.lastTime * 1000)
+            weight := thread.samples.weight.push ((nextStart - thread.lastTime) * 1000)
+            length := thread.samples.length + 1
+          } }
+          go (some stackIdx) c
+      -- add remaining slice after last child
+      modify fun thread => { thread with
+        lastTime := data.stopTime
+        samples := {
+          stack := thread.samples.stack.push stackIdx
+          time := thread.samples.time.push (thread.lastTime * 1000)
+          weight := thread.samples.weight.push ((data.stopTime - thread.lastTime) * 1000)
+          length := thread.samples.length + 1
+        } }
+    | .withContext _ msg => go parentStackIdx? msg
+    | _ => return
+  /-- Returns first `startTime` in the trace tree, if any. -/
+  getNextStart?
+    | .trace data _ children => do
+      if data.startTime != 0 then
+        return data.startTime
+      if let some startTime := children.findSome? getNextStart? then
+        return startTime
+      none
+    | .withContext _ msg => getNextStart? msg
+    | _ => none
+
+def Thread.new (name : String) : Thread := {
+  name
+  samples := { stack := #[], time := #[], weight := #[], length := 0 }
+  stackTable := { frame := #[], «prefix» := #[], category := #[], subcategory := #[], length := 0 }
+  frameTable := { func := #[], length := 0 }
+  stringArray := #[]
+  funcTable := {
+    name := #[], resource := #[], fileName := #[], lineNumber := #[], columnNumber := #[],
+    length := 0 }
+}
+
+def Profile.export (name : String) (startTime : Milliseconds) (traceState : TraceState)
+    (opts : Options) : IO Profile := do
+  let thread := Thread.new name
+  -- wrap entire trace up to current time in `runFrontend` node
+  let trace := .trace {
+    cls := `runFrontend, startTime, stopTime := (← IO.monoNanosNow).toFloat / 1000000000,
+    collapsed := true } "" (traceState.traces.toArray.map (·.msg))
+  let thread ← addTrace (Lean.trace.profiler.output.pp.get opts) { thread with } trace
+  return {
+    meta := { startTime, categories }
+    threads := #[thread.toThread]
+  }
+
+structure ThreadWithCollideMaps extends ThreadWithMaps where
+  /-- Max sharing map for samples -/
+  sampleMap : HashMap Nat Nat := {}
+
+/--
+Adds samples from `add` to `thread`, increasing the weight of existing samples with identical stacks
+instead of pushing new samples.
+-/
+private partial def collideThreads (thread : ThreadWithCollideMaps) (add : Thread) :
+    ThreadWithCollideMaps :=
+  StateT.run collideSamples thread |>.2
+where
+  collideSamples : StateM ThreadWithCollideMaps Unit := do
+    for oldSampleIdx in [0:add.samples.length] do
+      let oldStackIdx := add.samples.stack[oldSampleIdx]!
+      let stackIdx ← collideStacks oldStackIdx
+      modify fun thread =>
+        if let some idx := thread.sampleMap.find? stackIdx then
+          -- imperative to preserve linear use of arrays here!
+          let ⟨⟨⟨t1, t2, t3, samples, t5, t6, t7, t8, t9, t10⟩, o2, o3, o4, o5⟩, o6⟩ := thread
+          let ⟨s1, s2, weight, s3, s4⟩ := samples
+          let weight := weight.set! idx <| weight[idx]! + add.samples.weight[oldSampleIdx]!
+          let samples := ⟨s1, s2, weight, s3, s4⟩
+          ⟨⟨⟨t1, t2, t3, samples, t5, t6, t7, t8, t9, t10⟩, o2, o3, o4, o5⟩, o6⟩
+        else
+          -- imperative to preserve linear use of arrays here!
+          let ⟨⟨⟨t1, t2, t3, samples, t5, t6, t7, t8, t9, t10⟩, o2, o3, o4, o5⟩, sampleMap⟩ :=
+            thread
+          let ⟨stack, time, weight, _, length⟩ := samples
+          let samples := {
+              stack := stack.push stackIdx
+              time := time.push time.size.toFloat
+              weight := weight.push add.samples.weight[oldSampleIdx]!
+              length := length + 1
+            }
+          let sampleMap := sampleMap.insert stackIdx sampleMap.size
+          ⟨⟨⟨t1, t2, t3, samples, t5, t6, t7, t8, t9, t10⟩, o2, o3, o4, o5⟩, sampleMap⟩
+  collideStacks oldStackIdx : StateM ThreadWithCollideMaps Nat := do
+    let oldParentStackIdx? := add.stackTable.prefix[oldStackIdx]!
+    let parentStackIdx? ← oldParentStackIdx?.mapM (collideStacks ·)
+    let oldFrameIdx := add.stackTable.frame[oldStackIdx]!
+    let oldFuncIdx := add.frameTable.func[oldFrameIdx]!
+    let oldStrIdx := add.funcTable.name[oldFuncIdx]!
+    let strIdx ← getStrIdx add.stringArray[oldStrIdx]!
+    let funcIdx ← modifyGet fun thread =>
+      if let some idx := thread.funcMap.find? strIdx then
+        (idx, thread)
+      else
+        (thread.funcMap.size, { thread with
+          funcTable := {
+            name := thread.funcTable.name.push strIdx
+            resource := thread.funcTable.resource.push (-1)
+            fileName := thread.funcTable.fileName.push none
+            lineNumber := thread.funcTable.lineNumber.push none
+            columnNumber := thread.funcTable.columnNumber.push none
+            length := thread.funcTable.length + 1
+          }
+          frameTable := {
+            func := thread.frameTable.func.push thread.funcMap.size
+            length := thread.frameTable.length + 1
+          }
+          funcMap := thread.funcMap.insert strIdx thread.funcMap.size })
+    let frameIdx := funcIdx
+    modifyGet fun thread =>
+      if let some idx := thread.stackMap.find? (frameIdx, parentStackIdx?) then
+        (idx, thread)
+      else
+        (thread.stackMap.size,
+          let ⟨⟨⟨t1,t2, t3, t4, t5, stackTable, t7, t8, t9, t10⟩, o2, o3, stackMap, o5⟩, o6⟩ :=
+            thread
+          let { frame, «prefix», category, subcategory, length } := stackTable
+          let stackTable := {
+            frame := frame.push frameIdx
+            «prefix» := prefix.push parentStackIdx?
+            category := category.push add.stackTable.category[oldStackIdx]!
+            subcategory := subcategory.push add.stackTable.subcategory[oldStackIdx]!
+            length := length + 1
+          }
+          let stackMap := stackMap.insert (frameIdx, parentStackIdx?) stackMap.size
+          ⟨⟨⟨t1,t2, t3, t4, t5, stackTable, t7, t8, t9, t10⟩, o2, o3, stackMap, o5⟩, o6⟩)
+  getStrIdx (s : String) :=
+    modifyGet fun thread =>
+      if let some idx := thread.stringMap.find? s then
+        (idx, thread)
+      else
+        (thread.stringMap.size, { thread with
+          stringArray := thread.stringArray.push s
+          stringMap := thread.stringMap.insert s thread.stringMap.size })
+
+/--
+Merges given profiles such that samples with identical stacks are deduplicated by adding up their
+weights. Minimizes profile size while preserving per-stack timing information.
+-/
+def Profile.collide (ps : Array Profile) : Option Profile := do
+  let base ← ps[0]?
+  let thread := Thread.new "collided"
+  let thread := ps.map (·.threads) |>.flatten.foldl collideThreads { thread with }
+  return { base with threads := #[thread.toThread] }
+
+end Lean.Firefox
--- a/src/Lean/Util/Trace.lean
+++ b/src/Lean/Util/Trace.lean
@@ -129,7 +129,7 @@ def addRawTrace (msg : MessageData) : m Unit := do
 def addTrace (cls : Name) (msg : MessageData) : m Unit := do
  let ref ← getRef
  let msg ← addMessageContext msg
-  modifyTraces (·.push { ref, msg := .trace (collapsed := false) cls msg #[] })
+  modifyTraces (·.push { ref, msg := .trace { collapsed := false, cls } msg #[] })

@[inline] def trace (cls : Name) (msg : Unit → MessageData) : m Unit := do
  if (← isTracingEnabledFor cls) then
@@ -141,18 +141,18 @@ def addTrace (cls : Name) (msg : MessageData) : m Unit := do
    addTrace cls msg

 private def addTraceNode (oldTraces : PersistentArray TraceElem)
-    (cls : Name) (ref : Syntax) (msg : MessageData) (collapsed : Bool) : m Unit :=
+    (data : TraceData) (ref : Syntax) (msg : MessageData) : m Unit :=
  withRef ref do
-  let msg := .trace cls msg ((← getTraces).toArray.map (·.msg)) collapsed
+  let msg := .trace data msg ((← getTraces).toArray.map (·.msg))
  let msg ← addMessageContext msg
  modifyTraces fun _ =>
    oldTraces.push { ref, msg }

-def withSeconds [Monad m] [MonadLiftT BaseIO m] (act : m α) : m (α × Float) := do
+def withStartStopSeconds [Monad m] [MonadLiftT BaseIO m] (act : m α) : m (α × Float × Float) := do
  let start ← IO.monoNanosNow
  let a ← act
  let stop ← IO.monoNanosNow
-  return (a, (stop - start).toFloat / 1000000000)
+  return (a, start.toFloat / 1000000000, stop.toFloat / 1000000000)

 register_builtin_option trace.profiler : Bool := {
  defValue := false
@@ -166,6 +166,20 @@ register_builtin_option trace.profiler.threshold : Nat := {
  descr    := "threshold in milliseconds, traces below threshold will not be activated"
 }

+register_builtin_option trace.profiler.output : String := {
+  defValue := ""
+  group    := "profiler"
+  descr    :=
+    "output `trace.profiler` data in Firefox Profiler-compatible format to given file path"
+}
+
+register_builtin_option trace.profiler.output.pp : Bool := {
+  defValue := false
+  group    := "profiler"
+  descr    :=
+    "if false, limit text in exported trace nodes to trace class name and `TraceData.tag`, if any"
+}
+
 def trace.profiler.threshold.getSecs (o : Options) : Float :=
  (trace.profiler.threshold.get o).toFloat / 1000

@@ -208,31 +222,32 @@ instance [always : MonadAlwaysExcept ε m] [STWorld ω m] [BEq α] [Hashable α]
  except := let _ := always.except; inferInstance

 def withTraceNode [always : MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] (cls : Name)
-    (msg : Except ε α → m MessageData) (k : m α) (collapsed := true) : m α := do
+    (msg : Except ε α → m MessageData) (k : m α) (collapsed := true) (tag := "") : m α := do
  let _ := always.except
  let opts ← getOptions
  let clsEnabled ← isTracingEnabledFor cls
  unless clsEnabled || trace.profiler.get opts do
    return (← k)
  let oldTraces ← getResetTraces
-  let (res, secs) ← withSeconds <| observing k
-  let aboveThresh := trace.profiler.get opts && secs > trace.profiler.threshold.getSecs opts
+  let (res, start, stop) ← withStartStopSeconds <| observing k
+  let aboveThresh := trace.profiler.get opts && stop - start > trace.profiler.threshold.getSecs opts
  unless clsEnabled || aboveThresh do
    modifyTraces (oldTraces ++ ·)
    return (← MonadExcept.ofExcept res)
  let ref ← getRef
  let mut m ← try msg res catch _ => pure m!"<exception thrown while producing trace node message>"
+  let mut data := { cls, collapsed, tag }
  if profiler.get opts || aboveThresh then
-    m := m!"[{secs}s] {m}"
-  addTraceNode oldTraces cls ref m collapsed
+    data := { data with startTime := start, stopTime := stop }
+  addTraceNode oldTraces data ref m
  MonadExcept.ofExcept res

 def withTraceNode' [MonadAlwaysExcept Exception m] [MonadLiftT BaseIO m] (cls : Name)
-    (k : m (α × MessageData)) (collapsed := true) : m α :=
+    (k : m (α × MessageData)) (collapsed := true) (tag := "") : m α :=
  let msg := fun
    | .ok (_, msg) => return msg
    | .error err => return err.toMessageData
-  Prod.fst <$> withTraceNode cls msg k collapsed
+  Prod.fst <$> withTraceNode cls msg k collapsed tag

 end

@@ -300,7 +315,7 @@ TODO: find better name for this function.
 -/
 def withTraceNodeBefore [MonadRef m] [AddMessageContext m] [MonadOptions m]
    [always : MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] [ExceptToEmoji ε α] (cls : Name)
-    (msg : m MessageData) (k : m α) (collapsed := true) : m α := do
+    (msg : m MessageData) (k : m α) (collapsed := true) (tag := "") : m α := do
  let _ := always.except
  let opts ← getOptions
  let clsEnabled ← isTracingEnabledFor cls
@@ -310,15 +325,16 @@ def withTraceNodeBefore [MonadRef m] [AddMessageContext m] [MonadOptions m]
  let ref ← getRef
  -- make sure to preserve context *before* running `k`
  let msg ← withRef ref do addMessageContext (← msg)
-  let (res, secs) ← withSeconds <| observing k
-  let aboveThresh := trace.profiler.get opts && secs > trace.profiler.threshold.getSecs opts
+  let (res, start, stop) ← withStartStopSeconds <| observing k
+  let aboveThresh := trace.profiler.get opts && stop - start > trace.profiler.threshold.getSecs opts
  unless clsEnabled || aboveThresh do
    modifyTraces (oldTraces ++ ·)
    return (← MonadExcept.ofExcept res)
  let mut msg := m!"{ExceptToEmoji.toEmoji res} {msg}"
+  let mut data := { cls, collapsed, tag }
  if profiler.get opts || aboveThresh then
-    msg := m!"[{secs}s] {msg}"
-  addTraceNode oldTraces cls ref msg collapsed
+    data := { data with startTime := start, stopTime := stop }
+  addTraceNode oldTraces data ref msg
  MonadExcept.ofExcept res

 end Lean
--- a/src/Lean/Widget/InteractiveDiagnostic.lean
+++ b/src/Lean/Widget/InteractiveDiagnostic.lean
@@ -143,10 +143,10 @@ where
  | ctx,       nest n d                 => Format.nest n <$> go nCtx ctx d
  | ctx,       compose d₁ d₂            => do let d₁ ← go nCtx ctx d₁; let d₂ ← go nCtx ctx d₂; pure $ d₁ ++ d₂
  | ctx,       group d                  => Format.group <$> go nCtx ctx d
-  | ctx,       .trace cls header children collapsed => do
+  | ctx,       .trace data header children => do
    let header := (← go nCtx ctx header).nest 4
    let nodes ←
-      if collapsed && !children.isEmpty then
+      if data.collapsed && !children.isEmpty then
        let children := children.map fun child =>
          MessageData.withNamingContext nCtx <|
            match ctx with
@@ -154,11 +154,11 @@ where
            | none     => child
        let blockSize := ctx.bind (infoview.maxTraceChildren.get? ·.opts)
          |>.getD infoview.maxTraceChildren.defValue
-        let children := chopUpChildren cls blockSize children.toSubarray
+        let children := chopUpChildren data.cls blockSize children.toSubarray
        pure (.lazy children)
      else
        pure (.strict (← children.mapM (go nCtx ctx)))
-    let e := .trace cls header collapsed nodes
+    let e := .trace data.cls header data.collapsed nodes
    return .tag (← pushEmbed e) ".\n"

  /-- Recursively moves child nodes after the first `blockSize` into a new "more" node. -/
@@ -167,7 +167,7 @@ where
    if children.size > blockSize + 1 then  -- + 1 to make idempotent
      let more := chopUpChildren cls blockSize children[blockSize:]
      children[:blockSize].toArray.push <|
-        .trace (collapsed := true) cls
+        .trace { collapsed := true, cls }
          f!"{children.size - blockSize} more entries..." more
    else children