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

@@ -77,7 +77,9 @@ import type {
   StackAddressInfo,
   AddressTimings,
   Address,
+  IndexIntoAddressSetTable,
 } from 'firefox-profiler/types';
+import { IntSetTableBuilder } from 'firefox-profiler/utils/intset-table';
 
 import { getMatchingAncestorStackForInvertedCallNode } from './profile-data';
 import type { CallNodeInfo, CallNodeInfoInverted } from './call-node-info';
@@ -135,48 +137,32 @@ export function getStackAddressInfo(
   _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 IntSetTableBuilder();
+  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.indexForSetWithParentAndSelf(
+          prefixAddressSet !== -1 ? prefixAddressSet : null,
+          selfAddress
+        );
     }
-
-    selfAddressForAllStacks.push(selfAddress);
-    totalAddressesForAllStacks.push(totalAddresses);
   }
   return {
-    selfAddress: selfAddressForAllStacks,
-    stackAddresses: totalAddressesForAllStacks,
+    stackIndexToAddressSetIndex,
+    addressSetTable: builder.finish(),
   };
 }
 
@@ -502,30 +488,58 @@ 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;
+  // Iterate over all the samples, and aggregate the sample's weight into
+  // selfPerAddressSet.
+  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;
+    }
+  }
+
+  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 AddressSetPrefix = addressSetTable.prefix[addressSetIndex];
+    if (AddressSetPrefix !== null) {
+      selfSumOfAddressSetDescendants[AddressSetPrefix] += 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

@@ -12,7 +12,9 @@ import type {
   LineTimings,
   LineNumber,
   IndexIntoSourceTable,
+  IndexIntoLineSetTable,
 } from 'firefox-profiler/types';
+import { IntSetTableBuilder } from 'firefox-profiler/utils/intset-table';
 
 import { getMatchingAncestorStackForInvertedCallNode } from './profile-data';
 import type { CallNodeInfo, CallNodeInfoInverted } from './call-node-info';
@@ -60,53 +62,35 @@ export function getStackLineInfo(
   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> = [];
+  const builder = new IntSetTableBuilder();
+  const stackIndexToLineSetIndex = 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" 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.
   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.indexForSetWithParentAndSelf(
+          prefixLineSet !== -1 ? prefixLineSet : null,
+          selfLineOrNull !== null ? selfLineOrNull : -1
+        );
     }
-
-    selfLineForAllStacks.push(selfLine);
-    totalLinesForAllStacks.push(totalLines);
   }
   return {
-    selfLine: selfLineForAllStacks,
-    stackLines: totalLinesForAllStacks,
+    stackIndexToLineSetIndex,
+    lineSetTable: builder.finish(),
   };
 }
 
@@ -376,30 +360,57 @@ 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;
+
+  console.log('lineSetTable.length:', lineSetTable.length);
 
-  // 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?
+  // Iterate over all the samples, and aggregate the sample's weight into
+  // selfPerLineSet.
+  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;
+    }
+  }
+
+  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 lineSetPrefix = lineSetTable.prefix[lineSetIndex];
+    if (lineSetPrefix !== null) {
+      selfSumOfLineSetDescendants[lineSetPrefix] += 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

@@ -42,8 +42,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

@@ -42,8 +42,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);
   });
 });
 

src/types/profile-derived.ts

@@ -39,6 +39,7 @@ import type { IndexedArray } from './utils';
 import type { BitSet } from '../utils/bitset';
 import type { StackTiming } from '../profile-logic/stack-timing';
 import type { StringTable } from '../utils/string-table';
+import type { IndexIntoIntSetTable, IntSetTable } from 'firefox-profiler/utils/intset-table';
 export type IndexIntoCallNodeTable = number;
 
 /**
@@ -337,20 +338,12 @@ export type LineNumber = number;
 // stackLine may be null. This is fine because their values are not accessed
 // during the LineTimings computation.
 export type StackLineInfo = {
-  // An array that contains, for each "self" stack, the line number that this stack
-  // spends its self time in, in this file, or null if the self time of the
-  // stack is in a different file or if the line number is not known.
-  // For non-"self" stacks, i.e. stacks which are only used as prefix stacks and
-  // never referred to from a SamplesLikeTable, the value may be null.
-  selfLine: Array<LineNumber | null>;
-  // An array that contains, for each "self" stack, all the lines that the frames in
-  // this stack hit in this file, or null if this stack does not hit any line
-  // in the given file.
-  // For non-"self" stacks, i.e. stacks which are only used as prefix stacks and
-  // never referred to from a SamplesLikeTable, the value may be null.
-  stackLines: Array<Set<LineNumber> | null>;
+  stackIndexToLineSetIndex: Int32Array;
+  lineSetTable: IntSetTable;
 };
 
+export type IndexIntoLineSetTable = IndexIntoIntSetTable;
+
 // Stores, for all lines of one specific file, how many times each line is hit
 // by samples in a thread. The maps only contain non-zero values.
 // So map.get(line) === undefined should be treated as zero.
@@ -377,20 +370,12 @@ export type LineTimings = {
 // stackAddress may be null. This is fine because their values are not accessed
 // during the AddressTimings computation.
 export type StackAddressInfo = {
-  // An array that contains, for each "self" stack, the address that this stack
-  // spends its self time in, in this native symbol, or null if the self time of
-  // the stack is in a different native symbol or if the address is not known.
-  // For non-"self" stacks, i.e. stacks which are only used as prefix stacks and
-  // never referred to from a SamplesLikeTable, the value may be null.
-  selfAddress: Array<Address | null>;
-  // An array that contains, for each "self" stack, all the addresses that the
-  // frames in this stack hit in this native symbol, or null if this stack does
-  // not hit any address in the given native symbol.
-  // For non-"self" stacks, i.e. stacks which are only used as prefix stacks and
-  // never referred to from a SamplesLikeTable, the value may be null.
-  stackAddresses: Array<Set<Address> | null>;
+  stackIndexToAddressSetIndex: Int32Array;
+  addressSetTable: IntSetTable;
 };
 
+export type IndexIntoAddressSetTable = IndexIntoIntSetTable;
+
 // Stores, for all addresses of one specific library, how many times each
 // address is hit by samples in a thread. The maps only contain non-zero values.
 // So map.get(address) === undefined should be treated as zero.