Reduce allocations for StackLineInfo and StackAddressInfo.

Before: https://share.firefox.dev/3NnthpO
After: https://share.firefox.dev/45dMNv1

getStackLineInfo now allocates 46MB instead of 220MB (4.8x reduction),
and takes 146ms instead of 338ms (2.3x reduction), when opening the
source view for mozjemalloc.cpp on https://share.firefox.dev/3NnzEcF

Author: mstange

src/profile-logic/address-timings.ts

@@ -76,7 +76,9 @@ import type {
   StackAddressInfo,
   AddressTimings,
   Address,
+  IndexIntoAddressSetTable,
 } from 'firefox-profiler/types';
+import { SetCollectionBuilder } from 'firefox-profiler/utils/set-collection';
 
 /**
  * For each stack in `stackTable`, and one specific native symbol, compute the
@@ -112,67 +114,38 @@ import type {
  * If there is recursion, and the same address is present in multiple frames in
  * the same stack, the address is only counted once - the addresses are stored
  * in a set.
- *
- * The returned StackAddressInfo is computed as follows:
- *   selfAddress[stack]:
- *     For stacks whose stack.frame.nativeSymbol is the given native symbol,
- *     this is stack.frame.address.
- *     For all other stacks this is null.
- *   stackAddresses[stack]:
- *     For stacks whose stack.frame.nativeSymbol is the given native symbol,
- *     this is the stackAddresses of its prefix stack, plus stack.frame.address
- *     added to the set.
- *     For all other stacks this is the same as the stackAddresses set of the
- *     stack's prefix.
  */
 export function getStackAddressInfo(
   stackTable: StackTable,
   frameTable: FrameTable,
   _funcTable: FuncTable,
   nativeSymbol: IndexIntoNativeSymbolTable
 ): StackAddressInfo {
-  // "self address" == "the address which a stack's self time is contributed to"
-  const selfAddressForAllStacks = [];
-  // "total addresses" == "the set of addresses whose total time this stack contributes to"
-  const totalAddressesForAllStacks: Array<Set<Address> | null> = [];
+  const builder = new SetCollectionBuilder<number>();
+  const stackIndexToAddressSetIndex = new Int32Array(stackTable.length);
 
-  // This loop takes advantage of the fact that the stack table is topologically ordered:
-  // Prefix stacks are always visited before their descendants.
-  // Each stack inherits the "total" addresses from its parent stack, and then adds its
-  // self address to that set. If the stack doesn't have a self address in the library, we just
-  // re-use the prefix's set object without copying it.
   for (let stackIndex = 0; stackIndex < stackTable.length; stackIndex++) {
-    const frame = stackTable.frame[stackIndex];
     const prefixStack = stackTable.prefix[stackIndex];
-    const nativeSymbolOfThisStack = frameTable.nativeSymbol[frame];
+    const prefixAddressSet: IndexIntoAddressSetTable | -1 =
+      prefixStack !== null ? stackIndexToAddressSetIndex[prefixStack] : -1;
 
-    let selfAddress: Address | null = null;
-    let totalAddresses: Set<Address> | null =
-      prefixStack !== null ? totalAddressesForAllStacks[prefixStack] : null;
+    const frame = stackTable.frame[stackIndex];
+    const nativeSymbolOfThisStack = frameTable.nativeSymbol[frame];
+    const matchesNativeSymbol = nativeSymbolOfThisStack === nativeSymbol;
+    if (prefixAddressSet === -1 && !matchesNativeSymbol) {
+      stackIndexToAddressSetIndex[stackIndex] = -1;
+    } else {
+      const selfAddress = matchesNativeSymbol ? frameTable.address[frame] : -1;
 
-    if (nativeSymbolOfThisStack === nativeSymbol) {
-      selfAddress = frameTable.address[frame];
-      if (selfAddress !== -1) {
-        // Add this stack's address to this stack's totalAddresses. The rest of this stack's
-        // totalAddresses is the same as for the parent stack.
-        // We avoid creating new Set objects unless the new set is actually
-        // different.
-        if (totalAddresses === null) {
-          // None of the ancestor stack nodes have hit a address in the given library.
-          totalAddresses = new Set([selfAddress]);
-        } else if (!totalAddresses.has(selfAddress)) {
-          totalAddresses = new Set(totalAddresses);
-          totalAddresses.add(selfAddress);
-        }
-      }
+      stackIndexToAddressSetIndex[stackIndex] = builder.extend(
+        prefixAddressSet !== -1 ? prefixAddressSet : null,
+        selfAddress
+      );
     }
-
-    selfAddressForAllStacks.push(selfAddress);
-    totalAddressesForAllStacks.push(totalAddresses);
   }
   return {
-    selfAddress: selfAddressForAllStacks,
-    stackAddresses: totalAddressesForAllStacks,
+    stackIndexToAddressSetIndex,
+    addressSetTable: builder.finish(),
   };
 }
 
@@ -192,30 +165,70 @@ export function getAddressTimings(
   if (stackAddressInfo === null) {
     return emptyAddressTimings;
   }
-  const { selfAddress, stackAddresses } = stackAddressInfo;
-  const totalAddressHits: Map<Address, number> = new Map();
-  const selfAddressHits: Map<Address, number> = new Map();
 
-  // Iterate over all the samples, and aggregate the sample's weight into the
-  // addresses which are hit by the sample's stack.
-  // TODO: Maybe aggregate sample count per stack first, and then visit each stack only once?
+  const { stackIndexToAddressSetIndex, addressSetTable } = stackAddressInfo;
+
+  // We do two passes to compute the timings:
+  // 1. One pass over the samples to accumulate the sample weight onto the
+  //    nodes in the addressSetTable.
+  // 2. One pass (from back to front) over the addressSetTable, propagating
+  //    values up the tree and, at the same time, accumulating per-address
+  //    totals.
+
+  // First, do the pass over the samples to compute the weight per address set.
+  const selfPerAddressSet = new Float64Array(addressSetTable.length);
   for (let sampleIndex = 0; sampleIndex < samples.length; sampleIndex++) {
     const stackIndex = samples.stack[sampleIndex];
     if (stackIndex === null) {
       continue;
     }
-    const weight = samples.weight ? samples.weight[sampleIndex] : 1;
-    const setOfHitAddresses = stackAddresses[stackIndex];
-    if (setOfHitAddresses !== null) {
-      for (const address of setOfHitAddresses) {
-        const oldHitCount = totalAddressHits.get(address) ?? 0;
-        totalAddressHits.set(address, oldHitCount + weight);
+    const addressSetIndex = stackIndexToAddressSetIndex[stackIndex];
+    if (addressSetIndex !== -1) {
+      const weight = samples.weight ? samples.weight[sampleIndex] : 1;
+      selfPerAddressSet[addressSetIndex] += weight;
+    }
+  }
+
+  // Now, do a pass over the addressSetTable, from back to front.
+  // This is a similar idea to what we do for the call tree or the function
+  // list. The upwards propagation of a sample's weight will not contribute
+  // to the same address multiple times thanks to the guarantees of the
+  // addressSetTable - there are no duplicate values on a node's path to the
+  // root.
+  const totalAddressHits: Map<Address, number> = new Map();
+  const selfAddressHits: Map<Address, number> = new Map();
+  const selfSumOfAddressSetDescendants = new Float64Array(
+    addressSetTable.length
+  );
+  for (
+    let addressSetIndex = addressSetTable.length - 1;
+    addressSetIndex >= 0;
+    addressSetIndex--
+  ) {
+    const selfWeight = selfPerAddressSet[addressSetIndex];
+    if (selfWeight !== 0) {
+      const selfAddress = addressSetTable.self[addressSetIndex];
+      if (selfAddress !== -1) {
+        const oldHitCount = selfAddressHits.get(selfAddress) ?? 0;
+        selfAddressHits.set(selfAddress, oldHitCount + selfWeight);
       }
     }
-    const address = selfAddress[stackIndex];
-    if (address !== null) {
-      const oldHitCount = selfAddressHits.get(address) ?? 0;
-      selfAddressHits.set(address, oldHitCount + weight);
+
+    const selfSumOfThisAddressSetDescendants =
+      selfSumOfAddressSetDescendants[addressSetIndex];
+    const thisAddressSetWeight =
+      selfWeight + selfSumOfThisAddressSetDescendants;
+    const addressSetParent = addressSetTable.parent[addressSetIndex];
+    if (addressSetParent !== null) {
+      selfSumOfAddressSetDescendants[addressSetParent] += thisAddressSetWeight;
+    }
+
+    if (thisAddressSetWeight !== 0) {
+      const address = addressSetTable.value[addressSetIndex];
+      if (address !== -1) {
+        const oldHitCount = totalAddressHits.get(address) ?? 0;
+        totalAddressHits.set(address, oldHitCount + thisAddressSetWeight);
+      }
     }
   }
   return { totalAddressHits, selfAddressHits };

src/profile-logic/line-timings.ts

@@ -11,7 +11,9 @@ import type {
   LineTimings,
   LineNumber,
   IndexIntoSourceTable,
+  IndexIntoLineSetTable,
 } from 'firefox-profiler/types';
+import { SetCollectionBuilder } from 'firefox-profiler/utils/set-collection';
 
 /**
  * For each stack in `stackTable`, and one specific source file, compute the
@@ -40,69 +42,41 @@ import type {
  * This last line is the stack's "self line".
  * If there is recursion, and the same line is present in multiple frames in the
  * same stack, the line is only counted once - the lines are stored in a set.
- *
- * The returned StackLineInfo is computed as follows:
- *   selfLine[stack]:
- *     For stacks whose stack.frame.func.file is the given file, this is stack.frame.line.
- *     For all other stacks this is null.
- *   stackLines[stack]:
- *     For stacks whose stack.frame.func.file is the given file, this is the stackLines
- *     of its prefix stack, plus stack.frame.line added to the set.
- *     For all other stacks this is the same as the stackLines set of the stack's prefix.
  */
 export function getStackLineInfo(
   stackTable: StackTable,
   frameTable: FrameTable,
   funcTable: FuncTable,
   sourceViewSourceIndex: IndexIntoSourceTable
 ): StackLineInfo {
-  // "self line" == "the line which a stack's self time is contributed to"
-  const selfLineForAllStacks = [];
-  // "total lines" == "the set of lines whose total time this stack contributes to"
-  const totalLinesForAllStacks: Array<Set<LineNumber> | null> = [];
-
-  // This loop takes advantage of the fact that the stack table is topologically ordered:
-  // Prefix stacks are always visited before their descendants.
-  // Each stack inherits the "total" lines from its parent stack, and then adds its
-  // self line to that set. If the stack doesn't have a self line in the file, we just
-  // re-use the prefix's set object without copying it.
+  const builder = new SetCollectionBuilder<number>();
+  const stackIndexToLineSetIndex = new Int32Array(stackTable.length);
+
   for (let stackIndex = 0; stackIndex < stackTable.length; stackIndex++) {
-    const frame = stackTable.frame[stackIndex];
     const prefixStack = stackTable.prefix[stackIndex];
+    const prefixLineSet: IndexIntoLineSetTable | -1 =
+      prefixStack !== null ? stackIndexToLineSetIndex[prefixStack] : -1;
+
+    const frame = stackTable.frame[stackIndex];
     const func = frameTable.func[frame];
     const sourceIndexOfThisStack = funcTable.source[func];
+    const matchesSource = sourceIndexOfThisStack === sourceViewSourceIndex;
+    if (prefixLineSet === -1 && !matchesSource) {
+      stackIndexToLineSetIndex[stackIndex] = -1;
+    } else {
+      const selfLineOrNull = matchesSource
+        ? (frameTable.line[frame] ?? funcTable.lineNumber[func])
+        : null;
 
-    let selfLine: LineNumber | null = null;
-    let totalLines: Set<LineNumber> | null =
-      prefixStack !== null ? totalLinesForAllStacks[prefixStack] : null;
-
-    if (sourceIndexOfThisStack === sourceViewSourceIndex) {
-      selfLine = frameTable.line[frame];
-      // Fallback to func line info if frame line info is not available
-      if (selfLine === null) {
-        selfLine = funcTable.lineNumber[func];
-      }
-      if (selfLine !== null) {
-        // Add this stack's line to this stack's totalLines. The rest of this stack's
-        // totalLines is the same as for the parent stack.
-        // We avoid creating new Set objects unless the new set is actually
-        // different.
-        if (totalLines === null) {
-          // None of the ancestor stack nodes have hit a line in the given file.
-          totalLines = new Set([selfLine]);
-        } else if (!totalLines.has(selfLine)) {
-          totalLines = new Set(totalLines);
-          totalLines.add(selfLine);
-        }
-      }
+      stackIndexToLineSetIndex[stackIndex] = builder.extend(
+        prefixLineSet !== -1 ? prefixLineSet : null,
+        selfLineOrNull !== null ? selfLineOrNull : -1
+      );
     }
-
-    selfLineForAllStacks.push(selfLine);
-    totalLinesForAllStacks.push(totalLines);
   }
   return {
-    selfLine: selfLineForAllStacks,
-    stackLines: totalLinesForAllStacks,
+    stackIndexToLineSetIndex,
+    lineSetTable: builder.finish(),
   };
 }
 
@@ -122,30 +96,66 @@ export function getLineTimings(
   if (stackLineInfo === null) {
     return emptyLineTimings;
   }
-  const { selfLine, stackLines } = stackLineInfo;
-  const totalLineHits: Map<LineNumber, number> = new Map();
-  const selfLineHits: Map<LineNumber, number> = new Map();
+  const { stackIndexToLineSetIndex, lineSetTable } = stackLineInfo;
+
+  // We do two passes to compute the timings:
+  // 1. One pass over the samples to accumulate the sample weight onto the
+  //    nodes in the lineSetTable.
+  // 2. One pass (from back to front) over the lineSetTable, propagating
+  //    values up the tree and, at the same time, accumulating per-line
+  //    totals.
 
-  // Iterate over all the samples, and aggregate the sample's weight into the
-  // lines which are hit by the sample's stack.
-  // TODO: Maybe aggregate sample count per stack first, and then visit each stack only once?
+  // First, do the pass over the samples to compute the weight per line set.
+  const selfPerLineSet = new Float64Array(lineSetTable.length);
   for (let sampleIndex = 0; sampleIndex < samples.length; sampleIndex++) {
     const stackIndex = samples.stack[sampleIndex];
     if (stackIndex === null) {
       continue;
     }
-    const weight = samples.weight ? samples.weight[sampleIndex] : 1;
-    const setOfHitLines = stackLines[stackIndex];
-    if (setOfHitLines !== null) {
-      for (const line of setOfHitLines) {
-        const oldHitCount = totalLineHits.get(line) ?? 0;
-        totalLineHits.set(line, oldHitCount + weight);
+    const lineSetIndex = stackIndexToLineSetIndex[stackIndex];
+    if (lineSetIndex !== -1) {
+      const weight = samples.weight ? samples.weight[sampleIndex] : 1;
+      selfPerLineSet[lineSetIndex] += weight;
+    }
+  }
+
+  // Now, do a pass over the lineSetTable, from back to front.
+  // This is a similar idea to what we do for the call tree or the function
+  // list. The upwards propagation of a sample's weight will not contribute
+  // to the same line multiple times thanks to the guarantees of the
+  // lineSetTable - there are no duplicate values on a node's path to the
+  // root.
+  const totalLineHits: Map<LineNumber, number> = new Map();
+  const selfLineHits: Map<LineNumber, number> = new Map();
+  const selfSumOfLineSetDescendants = new Float64Array(lineSetTable.length);
+  for (
+    let lineSetIndex = lineSetTable.length - 1;
+    lineSetIndex >= 0;
+    lineSetIndex--
+  ) {
+    const selfWeight = selfPerLineSet[lineSetIndex];
+    if (selfWeight !== 0) {
+      const selfLine = lineSetTable.self[lineSetIndex];
+      if (selfLine !== -1) {
+        const oldHitCount = selfLineHits.get(selfLine) ?? 0;
+        selfLineHits.set(selfLine, oldHitCount + selfWeight);
       }
     }
-    const line = selfLine[stackIndex];
-    if (line !== null) {
-      const oldHitCount = selfLineHits.get(line) ?? 0;
-      selfLineHits.set(line, oldHitCount + weight);
+
+    const selfSumOfThisLineSetDescendants =
+      selfSumOfLineSetDescendants[lineSetIndex];
+    const thisLineSetWeight = selfWeight + selfSumOfThisLineSetDescendants;
+    const lineSetParent = lineSetTable.parent[lineSetIndex];
+    if (lineSetParent !== null) {
+      selfSumOfLineSetDescendants[lineSetParent] += thisLineSetWeight;
+    }
+
+    if (thisLineSetWeight !== 0) {
+      const line = lineSetTable.value[lineSetIndex];
+      if (line !== -1) {
+        const oldHitCount = totalLineHits.get(line) ?? 0;
+        totalLineHits.set(line, oldHitCount + thisLineSetWeight);
+      }
     }
   }
   return { totalLineHits, selfLineHits };

src/test/unit/address-timings.test.ts

@@ -44,8 +44,8 @@ describe('getStackAddressInfo', function () {
 
     // Expect the returned arrays to have the same length as the stackTable.
     expect(stackTable.length).toBe(9);
-    expect(stackLineInfoOne.selfAddress.length).toBe(9);
-    expect(stackLineInfoOne.stackAddresses.length).toBe(9);
+    expect(stackLineInfoOne.stackIndexToAddressSetIndex.length).toBe(9);
+    expect(stackLineInfoOne.addressSetTable.length).toBe(4);
   });
 });
 

src/test/unit/line-timings.test.ts

@@ -44,8 +44,8 @@ describe('getStackLineInfo', function () {
 
     // Expect the returned arrays to have the same length as the stackTable.
     expect(stackTable.length).toBe(9);
-    expect(stackLineInfoOne.selfLine.length).toBe(9);
-    expect(stackLineInfoOne.stackLines.length).toBe(9);
+    expect(stackLineInfoOne.lineSetTable.length).toBe(7);
+    // expect(stackLineInfoOne.stackLines.length).toBe(9);
   });
 });