Add an --insert-stack-labels modification to the profiler-edit tool

package.json

@@ -108,6 +108,7 @@
     "redux-logger": "^3.0.6",
     "redux-thunk": "^3.1.0",
     "reselect": "^4.1.8",
+    "smol-toml": "^1.6.1",
     "valibot": "^1.4.0",
     "workbox-window": "^7.4.1"
   },

src/node-tools/profiler-edit.ts

@@ -8,6 +8,7 @@ import { unserializeProfileOfArbitraryFormat } from 'firefox-profiler/profile-lo
 import { computeCompactedProfile } from 'firefox-profiler/profile-logic/profile-compacting';
 import { GOOGLE_STORAGE_BUCKET } from 'firefox-profiler/app-logic/constants';
 import { compress } from 'firefox-profiler/utils/gz';
+import { insertStackLabels } from 'firefox-profiler/profile-logic/insert-stack-labels';
 import { SymbolStore } from 'firefox-profiler/profile-logic/symbol-store';
 import {
   symbolicateProfile,
@@ -21,6 +22,10 @@ import {
 } from 'firefox-profiler/profile-logic/wasm-symbolication';
 import type { Profile } from 'firefox-profiler/types/profile';
 import { assertExhaustiveCheck } from 'firefox-profiler/utils/types';
+import {
+  parseLabelToml,
+  resolveAllLabels,
+} from 'firefox-profiler/utils/label-templates';
 
 /**
  * A CLI tool for editing profiles.
@@ -38,6 +43,9 @@ import { assertExhaustiveCheck } from 'firefox-profiler/utils/types';
  *   node node-tools-dist/profiler-edit.js -i input.json.gz -o out.json.gz \
  *     --symbolicate-wasm http://host/a.wasm=./a-unstripped.wasm \
  *     --symbolicate-wasm http://host/b.wasm=./b-unstripped.wasm
+ *
+ *   node node-tools-dist/profiler-edit.js --from-hash w1spyw917hg... -o out.json.gz \
+ *     --insert-label-frames known-functions.toml
  */
 
 type ProfileSource =
@@ -61,6 +69,7 @@ export interface CliOptions {
   output: string;
   symbolicateWithServer?: string;
   symbolicateWasm: WasmSymbolicationCliSpec[];
+  insertLabelFrames?: string;
 }
 
 function loadWasmSymbolicationSpecs(
@@ -77,6 +86,34 @@ function loadWasmSymbolicationSpecs(
   });
 }
 
+/**
+ * Reconstruct the func-name strings used by insertStackLabels' prefix matcher
+ * (mirrors getLabelIndexForFunc in insert-stack-labels.ts), so auto-discovery
+ * sees the same strings the labeler will compare against.
+ */
+function collectFuncNames(profile: Profile): string[] {
+  const { funcTable, sources, stringArray } = profile.shared;
+  const result: string[] = [];
+  for (let i = 0; i < funcTable.length; i++) {
+    let name = stringArray[funcTable.name[i]];
+    const sourceIndex = funcTable.source[i];
+    if (sourceIndex !== null) {
+      const filename = stringArray[sources.filename[sourceIndex]];
+      const line = funcTable.lineNumber[i];
+      const col = funcTable.columnNumber[i];
+      if (line !== null && col !== null) {
+        name += ` (${filename}:${line}:${col})`;
+      } else if (line !== null) {
+        name += ` (${filename}:${line})`;
+      } else {
+        name += ` (${filename})`;
+      }
+    }
+    result.push(name);
+  }
+  return result;
+}
+
 async function loadProfile(source: ProfileSource): Promise<Profile> {
   switch (source.type) {
     case 'FILE': {
@@ -129,7 +166,7 @@ async function loadProfile(source: ProfileSource): Promise<Profile> {
 }
 
 export async function run(options: CliOptions) {
-  const profile = await loadProfile(options.input);
+  let profile = await loadProfile(options.input);
 
   if (options.symbolicateWithServer !== undefined) {
     const server = options.symbolicateWithServer;
@@ -183,6 +220,15 @@ export async function run(options: CliOptions) {
     loadWasmSymbolicationSpecs(options.symbolicateWasm)
   );
 
+  if (options.insertLabelFrames !== undefined) {
+    console.log('Inserting label frames...');
+    const tomlText = fs.readFileSync(options.insertLabelFrames, 'utf8');
+    const parsed = parseLabelToml(tomlText);
+    const funcNames = collectFuncNames(profile);
+    const labels = resolveAllLabels(parsed, funcNames);
+    profile = insertStackLabels(profile, labels);
+  }
+
   const { profile: compactedProfile } = computeCompactedProfile(profile);
 
   console.log(`Saving profile to ${options.output}`);
@@ -243,7 +289,8 @@ export function makeOptionsFromArgv(processArgv: string[]): CliOptions {
       )
         .argParser(collectWasm)
         .default([] as WasmSymbolicationCliSpec[])
-    );
+    )
+    .option('--insert-label-frames <path>', 'TOML file with label definitions');
 
   program.parse(processArgv);
   const opts = program.opts();
@@ -290,6 +337,11 @@ export function makeOptionsFromArgv(processArgv: string[]): CliOptions {
         ? opts.symbolicateWithServer
         : undefined,
     symbolicateWasm: opts.symbolicateWasm,
+    insertLabelFrames:
+      typeof opts.insertLabelFrames === 'string' &&
+      opts.insertLabelFrames !== ''
+        ? opts.insertLabelFrames
+        : undefined,
   };
 }